[mini.pick] Syntax/treesitter highlight

[drowning-cat]

Related:

• [mini.pick:grep] Treesitter highlight #2151
• https://github.com/folke/snacks.nvim/blob/main/lua/snacks/picker/util/highlight.lua

Similar:

• [mini.pick] pick "buffer lines" in current buffer and include treesitter syntax highlighting #988
• [mini.pick] alignment with grep_live? #1192

Version 1

-- See https://github.com/folke/snacks.nvim/blob/main/lua/snacks/picker/util/highlight.lua#L7-L110
local lang_bufs = {
  ["lua"] = vim.api.nvim_create_buf(false, true),
}
---@param code string[]
---@param ft string
local get_code_highlights = function(code, ft)
  local lang = vim.treesitter.language.get_lang(ft)
  if
    not lang
    or not vim.api.nvim_get_runtime_file("parser/" .. lang .. ".so", false)
    or not vim.api.nvim_get_runtime_file("queries/" .. lang .. "/highlights.scm", false)
  then
    return {}
  end
  -- Create scratch buffer
  local buf = lang_bufs[lang]
  if not (buf and vim.api.nvim_buf_is_valid(buf)) then
    ---@cast buf integer
    buf = vim.api.nvim_create_buf(false, true)
    lang_bufs[lang] = buf
  end
  -- vim.api.nvim_buf_set_lines(buf, 0, -1, false, vim.split(code, "\n", { plain = true }))
  vim.api.nvim_buf_set_lines(buf, 0, -1, false, code)
  -- Init parser
  local parser = vim.treesitter.get_parser(buf, lang, { error = false })
  if not parser then
    return {}
  end
  parser:parse(true)
  -- Get highlights (copy-paste from snacks.nvim)
  local ret = {}
  parser:for_each_tree(function(tstree, tree)
    if not tstree then
      return
    end
    local query = vim.treesitter.query.get(tree:lang(), "highlights")
    if not query then
      return
    end
    for capture, node, metadata in query:iter_captures(tstree:root(), buf) do
      ---@type string
      local name = query.captures[capture]
      if name ~= "spell" then
        local range = { node:range() } ---@type number[]
        local multi = range[1] ~= range[3]
        local text = multi
            and vim.split(vim.treesitter.get_node_text(node, buf, metadata[capture]), "\n", { plain = true })
          or {}
        for row = range[1] + 1, range[3] + 1 do
          local first, last = row == range[1] + 1, row == range[3] + 1
          local end_col = last and range[4] or #(text[row - range[1]] or "")
          end_col = multi and first and end_col + range[2] or end_col
          ret[row] = ret[row] or {}
          table.insert(ret[row], {
            col = first and range[2] or 0,
            end_col = end_col,
            priority = (tonumber(metadata.priority or metadata[capture] and metadata[capture].priority) or 100),
            conceal = metadata.conceal or metadata[capture] and metadata[capture].conceal,
            hl_group = "@" .. name .. "." .. lang,
          })
        end
      end
    end
  end)
  return ret
end

-- Show function

local grep_cache = {}
-- Clear cache
vim.api.nvim_create_autocmd("User", {
  pattern = "MiniPickStop",
  callback = function()
    grep_cache = {}
  end,
})

local grep_ns = vim.api.nvim_create_namespace("minipick_grep")
local minipick_grep_show = function(buf, data, query, opts)
  -- Clear previous highlights
  vim.api.nvim_buf_clear_namespace(buf, grep_ns, 0, -1)
  -- Default show
  MiniPick.default_show(buf, data, query, opts)
  -- Highlight code
  local lines = vim.api.nvim_buf_get_lines(buf, 0, -1, true)
  for ln in ipairs(data) do
    -- Parse
    local line = lines[ln]
    local file, file_row, file_col, code = unpack(vim.split(data[ln], "%z"))
    local ft = vim.filetype.match({ filename = file }) or ""
    -- Get highlights +cache
    local cache_id = string.format("%s|%s|%s", file, file_row, file_col)
    local line_hls = grep_cache[cache_id]
    if not line_hls then
      line_hls = get_code_highlights({ code }, ft)[1]
      grep_cache[cache_id] = line_hls
    end
    -- Apply highlights
    if line_hls then
      -- NOTE: Prevent cache mutation
      line_hls = vim.deepcopy(line_hls, true)
      -- Find code start
      local _, off = string.find(line, ".*│.*│.*│")
      -- Set extmarks
      local extm_row = ln - 1
      for _, hl in ipairs(line_hls) do
        local extm_col = off + hl.col
        hl.col = nil
        hl.end_col = off + hl.end_col
        local ok, err = pcall(vim.api.nvim_buf_set_extmark, buf, grep_ns, extm_row, extm_col, hl)
        if not ok then
          vim.notify("Set extmark: " .. err, vim.log.levels.WARN)
        end
      end
    end
  end
end

vim.api.nvim_set_hl(0, "MiniPickMatchRanges", {
  bold = true,
  underdotted = true,
})

-- Set `source.show` for grep pickers
MiniPick.registry.grep = function(local_opts, opts)
  return MiniPick.builtin.grep(
    local_opts,
    vim.tbl_deep_extend("keep", opts or {}, { source = { show = minipick_grep_show } })
  )
end
MiniPick.registry.grep_live = function(local_opts, opts)
  return MiniPick.builtin.grep_live(
    local_opts,
    vim.tbl_deep_extend("keep", opts or {}, { source = { show = minipick_grep_show } })
  )
end

Version 2

-- Get highlights (snacks.nvim)
-- NOTE: https://github.com/folke/snacks.nvim/blob/main/lua/snacks/picker/util/highlight.lua

local lang_bufs = {}
local get_code_highlights = function(code, ft)
  local lang = vim.treesitter.language.get_lang(ft or "")
  if
    not lang
    or not vim.api.nvim_get_runtime_file("parser/" .. lang .. ".so", false)
    or not vim.api.nvim_get_runtime_file("queries/" .. lang .. "/highlights.scm", false)
    or vim.list_contains({ "text" }, lang)
  then
    return {}
  end
  local buf = lang_bufs[lang]
  if not (buf and vim.api.nvim_buf_is_valid(buf)) then
    ---@cast buf integer
    buf = vim.api.nvim_create_buf(false, true)
    vim.api.nvim_buf_set_name(buf, "mini.pick://highlight/" .. lang)
    lang_bufs[lang] = buf
  end
  vim.api.nvim_buf_set_lines(buf, 0, -1, false, code)
  local parser = vim.treesitter.get_parser(buf, lang, { error = false })
  if not parser then
    return {}
  end
  -- NOTE: https://github.com/folke/snacks.nvim/blob/fe7cfe9800a182274d0f868a74b7263b8c0c020b/lua/snacks/picker/util/highlight.lua#L49-L84
  local ret = {}
  parser:parse(true)
  parser:for_each_tree(function(tstree, tree)
    if not tstree then
      return
    end
    local query = vim.treesitter.query.get(tree:lang(), "highlights")
    if not query then
      return
    end
    for capture, node, metadata in query:iter_captures(tstree:root(), buf) do
      ---@type string
      local name = query.captures[capture]
      if name ~= "spell" then
        local range = { node:range() } ---@type number[]
        local multi = range[1] ~= range[3]
        local text = multi
            and vim.split(vim.treesitter.get_node_text(node, buf, metadata[capture]), "\n", { plain = true })
          or {}
        for row = range[1] + 1, range[3] + 1 do
          local first, last = row == range[1] + 1, row == range[3] + 1
          local end_col = last and range[4] or #(text[row - range[1]] or "")
          end_col = multi and first and end_col + range[2] or end_col
          ret[row] = ret[row] or {}
          table.insert(ret[row], {
            col = first and range[2] or 0,
            end_col = end_col,
            priority = (tonumber(metadata.priority or metadata[capture] and metadata[capture].priority) or 100),
            conceal = metadata.conceal or metadata[capture] and metadata[capture].conceal,
            hl_group = "@" .. name .. "." .. lang,
          })
        end
      end
    end
  end)
  return ret
end

-- Set highlights

local buf_set_highlights = function(buf, ns, row_start, col_offset, highlights)
  row_start, col_offset = row_start or 0, col_offset or 0
  highlights = vim.deepcopy(highlights or {}, true)
  for i, line_hls in ipairs(highlights) do
    local row = row_start + i - 1
    for _, hl in ipairs(line_hls) do
      local col = col_offset + hl.col
      hl.col = nil
      hl.end_col = col_offset + hl.end_col
      local ok, err = pcall(vim.api.nvim_buf_set_extmark, buf, ns, row, col, hl)
      if not ok then
        vim.notify(string.format("Set extmark %s|%s: %s", row, col, err), vim.log.levels.WARN)
        -- NOTE: Debug
        local line = vim.api.nvim_buf_get_lines(buf, row, row + 1, true)[1]
        vim.print({ line, #line, hl })
      end
    end
  end
end

-- Cache

local minipick_cache = {}
local minipick_use_cache = function(id, func)
  local val = minipick_cache[id]
  if not val then
    val = func()
    minipick_cache[id] = val
  end
  return val
end

-- Cleanup

vim.api.nvim_create_autocmd("User", {
  pattern = "MiniPickStop",
  callback = function()
    minipick_cache = {}
    for _, buf in ipairs(lang_bufs) do
      vim.api.nvim_buf_delete(buf, { unload = true })
    end
  end,
})

-- Show functions

local minipick_hl_ns = vim.api.nvim_create_namespace("minipick_grep")

-- -- Grep

local minipick_grep_show = function(buf, data, query, opts)
  vim.api.nvim_buf_clear_namespace(buf, minipick_hl_ns, 0, -1)
  MiniPick.default_show(buf, data, query, opts)
  -- Syntax highlight
  local lines = vim.api.nvim_buf_get_lines(buf, 0, #data, true)
  for ln in ipairs(data) do
    local line = lines[ln]
    local info, code = string.match(line, "^(.-│.-│.-│)(.*)$")
    local line_hls = minipick_use_cache(info, function()
      local file, _, _, _ = unpack(vim.split(data[ln], "%z"))
      local ft = vim.filetype.match({ filename = file }) or ""
      return get_code_highlights({ code }, ft)
    end)
    local row_0 = ln - 1
    local col_0 = #info
    buf_set_highlights(buf, minipick_hl_ns, row_0, col_0, line_hls)
  end
end

-- -- Lsp (references)

local minipick_lsp_refs_show = function(buf, data, query, opts)
  vim.api.nvim_buf_clear_namespace(buf, minipick_hl_ns, 0, -1)
  MiniPick.default_show(buf, data, query, opts)
  local lines = vim.api.nvim_buf_get_lines(buf, 0, #data, true)
  for ln in ipairs(data) do
    local line = lines[ln]
    local item = data[ln]
    local info, code = string.match(line, "^(.-│.-│)(.*)$")
    local line_hls = minipick_use_cache(info, function()
      local file = item.path
      local ft = vim.filetype.match({ filename = file }) or ""
      return get_code_highlights({ code }, ft)
    end)
    local row_0 = ln - 1
    local col_0 = #info
    buf_set_highlights(buf, minipick_hl_ns, row_0, col_0, line_hls)
  end
end

-- -- Buffer lines

local minipick_buf_lines_show = function(buf, data, query, opts)
  vim.api.nvim_buf_clear_namespace(buf, minipick_hl_ns, 0, -1)
  MiniPick.default_show(buf, data, query, opts)
  -- Syntax highlight
  local bufs = {}
  for _, item in ipairs(data) do
    bufs[item.bufnr] = true
  end
  local buf_hls = {}
  for item_buf in pairs(bufs) do
    buf_hls[item_buf] = minipick_use_cache(item_buf, function()
      local lines = vim.api.nvim_buf_get_lines(item_buf, 0, -1, true)
      return get_code_highlights(lines, vim.bo[item_buf].ft)
    end)
  end
  local lines = vim.api.nvim_buf_get_lines(buf, 0, #data, true)
  for ln, item in ipairs(data) do
    local line = lines[ln]
    local item_buf, lnum = item.bufnr, item.lnum
    local line_hls = buf_hls[item_buf][lnum] or {}
    local row_0 = ln - 1
    local col_0 = string.match(line, "^.-│.-│()") - 1
    buf_set_highlights(buf, minipick_hl_ns, row_0, col_0, { line_hls })
  end
end

-- Main

vim.api.nvim_set_hl(0, "MiniPickMatchRanges", {
  bold = true,
  underdotted = true,
})

MiniPick.registry.grep = function(local_opts, opts)
  return MiniPick.builtin.grep(
    local_opts,
    vim.tbl_deep_extend("keep", opts or {}, { source = { show = minipick_grep_show } })
  )
end
MiniPick.registry.grep_live = function(local_opts, opts)
  return MiniPick.builtin.grep_live(
    local_opts,
    vim.tbl_deep_extend("keep", opts or {}, { source = { show = minipick_grep_show } })
  )
end
MiniPick.registry.lsp = function(local_opts, opts)
  local show = { ["references"] = minipick_lsp_refs_show }
  return MiniExtra.pickers.lsp(
    local_opts,
    vim.tbl_deep_extend("keep", opts or {}, { source = { show = show[local_opts.scope] } })
  )
end
MiniPick.registry.buf_lines = function(local_opts, opts)
  return MiniExtra.pickers.buf_lines(
    local_opts,
    vim.tbl_deep_extend("keep", opts or {}, { source = { show = minipick_buf_lines_show } })
  )
end

[echasnovski]

Thanks for sharing!

The code is... a bit too much to quickly review. But if it works - it works :)

One extra thing I realized about this is that syntax highlighting in items adds restrictions about how query matches need to be highlighted to be visible. For example, they probably need to be highlighted with visible background or only reverse = true.

[drowning-cat]

I tried reverse = true. It looks good from a distance but gets a bit noisy up close, so I switched to bold = true and underdotted = true. The only issue now is matching the highlight vibe between grep and the other pickers.

If you decide to review the code, feel free to be picky. I’m open to changing anything since others might be interested.
Thank you for your time :)

P.S. It seems like nobody needs this feature :(