nvim/lua/core/todo.lua

---@class data.buf.node
---@field size number size of the node
---@field extmark_bg number? extmark id for the bg
---@field extmark_fg number? extmark id for the fg
---@field extmark_sn number? extmark id for the sign

---@class data.buf
---@field [integer] data.buf.node a list of nodes where the index is the row in file where the node resides

---@class data
---@field [integer] data.buf buffer number
local data = {}

local M = {}

local todo_comments_conf = {
  TODO = {
    -- TODO:
    -- NOTE:
    -- INFO:
    "TODO", "NOTE", "INFO",
    hlgroup = "TodoTODO"
  },
  BUG = {
    -- BUG:
    -- FIXME:
    "BUG", "FIXME",
    hlgroup = "TodoBUG"
  },
  TEST = {
    -- TEST:
    -- PERF:
    "TEST", "PERF",
    hlgroup = "TodoTEST"
  },
  WARN = {
    -- WARN:
    -- HACK:
    "WARN", "HACK",
    hlgroup = "TodoWARN"
  }
}

local todo_hl_ns = vim.api.nvim_create_namespace("todo_highlights")

--- create a new highlight or override an existing one
---@param bufnr number
---@param group string
---@param s_col number
---@param e_col number
---@param s_row number
---@param e_row number
---@param id number?
---@return number?
local function set_hl_at(bufnr, group, s_col, e_col, s_row, e_row, id)
  local ok, res = pcall(
    vim.api.nvim_buf_set_extmark,
    bufnr,
    todo_hl_ns,
    s_row,
    s_col,
    {
      hl_mode = "replace",
      hl_group = group,
      end_col = e_col,
      end_row = e_row,
      id = id,
    })

  if ok then
    return res
  end

  return nil
end

--- I'm wayyyyyy too lazy to make this function take a real list of arguments
---@param tab table everything in the whole universe in the most backwards incompatable way
local function create_extmarks(tab)
  local bufnr, s, e, node, only_continuation, t = vim.F.unpack_len(tab)

  local s_row, s_col = node:start()
  local e_row, e_col = node:end_()

  -- ensure that our string indicies are relative to the line
  s = s + s_col
  e = e + s_col

  if not data[bufnr][s_row] then
    local _, _, bytes = node:start()
    table.insert(data[bufnr], s_row, { size = bytes })
  end

  do
    local id = data[bufnr][s_row].extmark_fg
    local ext_id = set_hl_at(bufnr, t.hlgroup, only_continuation and s - 2 or e, e_col, s_row, e_row, id)
    if not id then
      data[bufnr][s_row].extmark_fg = ext_id
    end
  end

  if not only_continuation then
    do
      local id = data[bufnr][s_row].extmark_bg
      local ext_id = set_hl_at(bufnr, t.hlgroup.."BG", s - 2, e - 1, s_row, e_row, id)
      if not id then
        data[bufnr][s_row].extmark_bg = ext_id
      end
    end

    do
      local id = data[bufnr][s_row].extmark_sn
      local ext_id = set_hl_at(bufnr, t.hlgroup.."SIGN", e - 1, e, s_row, e_row, id)
      if not id then
        data[bufnr][s_row].extmark_sn = ext_id
      end
    end
  end
end

--- highlight todo comments in the current buffer.
--- No I won't use folke's super bloated plugin.
--- TODO: check if there's already an extmark at the location we're trying to
--- highlight, and if so don't add it again unless the node has changed.
---
--- ensure that we only change the extmark when the node has changed
function M.todo_comments()
  local bufnr = vim.api.nvim_win_get_buf(0)
  local ok, parser = pcall(vim.treesitter.get_parser, bufnr)
  if not ok or not parser then
    return
  end

  -- create an entry for the buffer in the data
  if not data[bufnr] then
    data[bufnr] = {}
  end

  -- Construct the query for comments.
  -- We're using treesitter so that I don't have to use external tooling.
  local ok, comment_query = pcall(vim.treesitter.query.parse,
    parser:lang(),
    "(comment) @comment"
  )
  if not ok then
    return
  end

  parser:parse(false, function(err, trees)
    if err then
      return
    end

    -- store the previous todo comment's location information
    local start_loc = nil
    local ls, le, lt = nil, nil, nil

    local root = trees[1]:root()
    for _, match in comment_query:iter_matches(root, bufnr, 0, -1) do
      for _, nodes in pairs(match) do
        for _, node in ipairs(nodes) do
          if not node or node:type() ~= "comment" then
            goto continue
          end
          local text = vim.treesitter.get_node_text(node, bufnr)

          -- is the node previous to this node a todo comment?
          local continuation = (start_loc and start_loc + 1 == node:start())

          -- the current todo comment's information
          local s, e = nil, nil
          local t = nil
          for _, type in pairs(todo_comments_conf) do
            for _, v in ipairs(type) do
              s, e = string.find(text, v..":")
              if not s or not e then
                s, e = string.find(text, v.."%b():")
              end
              if s and e then
                t = type
                goto work;
              end
            end
          end
          ::work::

          -- is the node previous to this node a todo comment and is this node
          -- not a todo comment?
          local only_continuation = (not s or not e) and continuation

          -- let's render this bad boy
          if (s and e) or continuation then
            if only_continuation then
              s, e, t = ls, le, lt
            end
            start_loc = node:start()

            if s and e then
              ls, le, lt = s, e, t
            end

            create_extmarks(vim.F.pack_len(
              bufnr,
              s,
              e,
              node,
              only_continuation,
              t
            ))
          end
          ::continue::
        end
      end
    end
  end)
end

return M