refactor(nvim): rewrite the highlight and extended marks plugins

1 files changed, 129 insertions, 96 deletions

diff --git a/.config/nvim/plugin/50-extmarks.lua b/.config/nvim/plugin/50-extmarks.lua
index 1324da0..50ba81d 100644
--- a/.config/nvim/plugin/50-extmarks.lua
+++ b/.config/nvim/plugin/50-extmarks.lua
@@ -1,129 +1,162 @@
 --
--- TODO highlighting (treesitter-aware extmarks)
+-- 50-extmarks.lua
+--
+-- * Creates a global highlight namespace for extended marks.
+-- * Keeps track of the current global highlight namespace in `vim.g.hl_ns`.
+-- * Creates an autocommand that initializes and updates extended marks.
+--   For now, only "TODO" strings are marked but the system can be easily extended by adding items
+--   in `marks`.
+--   Doesn't support multi-line patterns.
+--
+-- Global variables:
+--   `vim.g.hl_ns`: global highlight namespace
+--
+-- User commands:
+--   `ExtmarkToggle`: toggle extended marks highlight namespace
 --
--- The `dotfiles.Todo` highlight group and the storage namespace are owned by
--- 40-colors.lua, because the same per-window namespaces also drive the
--- focused/unfocused cursor line. That file is numbered 40- so it is sourced
--- before this one and `vim.g.dotfiles.todo_ns` / `.todo_hl` exist below. Here
--- we only place the extmarks.
 
--- Returns an iterator over the ascendants of `node`, starting at the root node
-local function node_ascendants(node)
-	local root = node:tree():root()
-	local current = root
-	return function()
-		if current:equal(node) then
-			return nil
-		end
-		current = assert(current:child_with_descendant(node))
-		return current
+vim.g.hl_ns = 0
+
+local extmarks_hl_ns = vim.api.nvim_create_namespace("dotfiles_extmarks")
+
+local hl_group_todo = "dotfiles.Todo"
+vim.api.nvim_set_hl(extmarks_hl_ns, hl_group_todo, { link = "Todo" })
+
+-- Wrapper around `nvim_set_hl_ns()`; sets the global highlight namespace to `hl_ns` and save it to
+-- `vim.g.hl_ns`. Useful because the global highlight namespace is not saved anywhere by default.
+local function set_global_hl_ns(hl_ns)
+	vim.g.hl_ns = hl_ns
+	vim.api.nvim_set_hl_ns(hl_ns)
+end
+
+-- Toggles extended marks highlight namespace
+local function toggle_marks()
+	local new_hl_ns = vim.g.hl_ns == 0 and extmarks_hl_ns or 0
+	set_global_hl_ns(new_hl_ns)
+end
+
+-- Parses the treesitter tree synchronously, including injections, in the 0-based, end-exclusive
+-- line range `first`:`last` of buffer `buf`.
+--- @return boolean parser_exists
+local function parse_tree(buf, first, last)
+	local parser = vim.treesitter.get_parser(buf)
+	if parser then
+		parser:parse({ first, last })
 	end
+	return parser ~= nil
 end
 
--- Returns true if `node` or any of its parent has the type `type`
-local function node_within_type(node, type)
-	local next_ascendant = node_ascendants(node)
-	local next = next_ascendant()
-	while next do
-		if next:type() == type then
+-- Returns a predicate telling whether a "TODO" string should be marked at a given position.
+-- In markdown buffers, positions inside code are invalid, otherwise positions outside comments are
+-- invalid. All positions are valid if there is no parser for the buffer filetype.
+-- Only the line range `first`:`last` is parsed; the predicate must not be called outside it.
+--- @param buf number
+--- @return fun(row: number, col: number): boolean predicate taking a 0-based position
+local function todo_filter(buf, first, last)
+	local has_parser = parse_tree(buf, first, last)
+	return function(row, col)
+		assert(first <= row and row < last, "Position outside the parsed line range")
+		if not has_parser then
 			return true
 		end
-		next = next_ascendant()
+		local block = vim.treesitter.get_node({
+			bufnr = buf,
+			pos = { row, col },
+		})
+		assert(block, "Expected node to be truthy, is the tree parsed?")
+		if vim.bo[buf].filetype == "markdown" then
+			if string.match(block:type(), "code") then
+				return false
+			end
+			local inline = vim.treesitter.get_node({
+				bufnr = buf,
+				ignore_injections = false,
+				pos = { row, col },
+			})
+			assert(inline, "Expected node to be truthy, is the tree parsed?")
+			return inline:type() ~= "code_span"
+		else
+			return string.match(block:type(), "comment") ~= nil
+		end
 	end
-	return false
 end
 
--- Returns an iterator over the positions (`{ col1, col2, row }`) of matches of pattern in lines
-local function pmatches(lines, pattern)
-	local col1, col2, row = nil, nil, 1
+-- Returns an iterator over the positions of valid "TODO" strings in buffer `buf` in the line range
+-- `first`:`last`. Each item is a table with 0-based fields: row `y`, start column `x1` (inclusive)
+-- and end column `x2` (exclusive). See `todo_filter()` for what makes a position valid.
+local function find_todo(buf, first, last)
+	local lines = vim.api.nvim_buf_get_lines(buf, first, last, false)
+	local should_mark = todo_filter(buf, first, first + #lines)
+	local x1, x2, y = nil, nil, 1
 	return function()
-		while row <= #lines do
-			col1, col2 = string.find(lines[row], pattern, (col2 or 0) + 1)
-			if col1 then
-				return { col1 = col1, col2 = col2, row = row }
+		while y <= #lines do
+			x1, x2 = string.find(lines[y], "TODO", (x2 or 0) + 1)
+			if x1 and x2 then
+				local row = y + first - 1
+				if should_mark(row, x1 - 1) then
+					return { y = row, x1 = x1 - 1, x2 = x2 }
+				end
 			else
-				col1, col2, row = 0, 0, row + 1
+				x1, x2, y = 0, 0, y + 1
 			end
 		end
 		return nil
 	end
 end
 
--- Set extmarks are the positions (`{ col1, col2, row }`) given by `next_pos()`
--- `opts`: {
---     ns: highlight namespace
---     predicate: is passed row and col, if true then set the mark
---     hl_group: highlight group
---  }
-local function set_extmarks(buf, next_pos, opts)
-	for pos in next_pos do
-		local marks = vim.api.nvim_buf_get_extmarks(
-			buf,
-			opts.ns,
-			{ pos.row - 1, pos.col1 - 1 },
-			{ pos.row - 1, pos.col2 - 1 },
-			{}
-		)
-		if #marks == 0 and opts.predicate(buf, pos.row - 1, pos.col1 - 1) then
+-- List of extended marks to apply with `mark_all()`
+--
+-- Each item is a table with two fields: `hl_group`, the name of the highlight group to use
+-- for the marks, and `find`, which must return an iterator over the positions to mark (see
+-- `find_todo()` for the format).
+local marks = {
+	{ hl_group = hl_group_todo, find = find_todo },
+}
+
+-- Updates extended marks in buffer `buf` in range `first`:`last`.
+-- For each item in `marks`, sets extended marks on positions returned by `item.find` to the
+-- highlight group `item.hl_group`.
+local function mark_all(buf, first, last)
+	vim.api.nvim_buf_clear_namespace(buf, extmarks_hl_ns, first, last)
+	for _, marks_opts in ipairs(marks) do
+		for p in marks_opts.find(buf, first, last) do
 			vim.api.nvim_buf_set_extmark(
 				buf,
-				opts.ns,
-				pos.row - 1,
-				pos.col1 - 1,
-				{ end_col = pos.col2, hl_group = opts.hl_group }
+				extmarks_hl_ns,
+				p.y,
+				p.x1,
+				{ end_col = p.x2, hl_group = marks_opts.hl_group }
 			)
 		end
 	end
 end
 
--- TODO a better way to do this:
--- For code, only check in "comment" node ranges
--- For non-code, do it per type. For Markdown, only check "paragraph" nodes
-local function hl_todo_predicate(buf, row, col)
-	if vim.bo[buf].filetype == "markdown" or vim.treesitter.get_parser(buf) == nil then
-		return true
+-- Create extended marks in buffer `ev.buf` and updates them when the buffer changes.
+local function watch_extmarks(ev)
+	mark_all(ev.buf, 0, -1)
+	if vim.b[ev.buf].dotfiles_todo_attached then
+		return
 	end
-	local node = assert(vim.treesitter.get_node({ bufnr = buf, pos = { row, col } }))
-	return node_within_type(node, "comment")
+	vim.b[ev.buf].dotfiles_todo_attached = true
+	vim.api.nvim_buf_attach(ev.buf, false, {
+		on_lines = function(_, b, _, first, _, last_new)
+			vim.schedule(function()
+				if vim.api.nvim_buf_is_valid(b) then
+					mark_all(b, first, last_new)
+				end
+			end)
+		end,
+	})
 end
 
-local patterns = { todo = "TODO" }
-local extmarks_todo_opts = {
-	ns = vim.g.dotfiles.todo_ns,
-	hl_group = vim.g.dotfiles.todo_hl,
-	predicate = hl_todo_predicate,
-}
-
-local function scan(buf, first, last)
-	local lines = vim.api.nvim_buf_get_lines(buf, first, last, false)
-	local iter = pmatches(lines, patterns.todo)
-	set_extmarks(buf, function()
-		local pos = iter()
-		if pos then
-			pos.row = pos.row + first
-		end
-		return pos
-	end, extmarks_todo_opts)
-end
+----------------------------------------------------------------------------------------------------
 
 vim.api.nvim_create_autocmd("FileType", {
-	desc = "Initialize and watch TODO extmarks",
+	desc = "Set extended marks watcher",
 	group = vim.g.dotfiles.augroup,
-	callback = function(args)
-		local buf = args.buf
-		scan(buf, 0, -1)
-		if vim.b[buf].dotfiles_todo_attached then
-			return
-		end
-		vim.b[buf].dotfiles_todo_attached = true
-		vim.api.nvim_buf_attach(buf, false, {
-			on_lines = function(_, b, _, first, _, last_new)
-				vim.schedule(function()
-					if vim.api.nvim_buf_is_valid(b) then
-						scan(b, first, last_new)
-					end
-				end)
-			end,
-		})
-	end,
+	callback = watch_extmarks,
 })
+
+vim.api.nvim_create_user_command("ExtmarkToggle", toggle_marks, { desc = "Toggle extended marks" })
+
+set_global_hl_ns(extmarks_hl_ns)