feat: treesitter highlighter (#404)

* feat: treesitter highlighter

* perf(treesitter): cache treesitter highlights

* feat(menu): treesitter draw config option

Author: folke

lua/blink/cmp/completion/windows/render/treesitter.lua

@@ -0,0 +1,73 @@
+local treesitter = {}
+
+---@type table<string, blink.cmp.DrawHighlight[]>
+local cache = {}
+local cache_size = 0
+local MAX_CACHE_SIZE = 1000
+
+--- @param ctx blink.cmp.DrawItemContext
+--- @param opts? {offset?: number}
+function treesitter.highlight(ctx, opts)
+  -- early return if treesitter highlight is disabled
+  if not vim.b.ts_highlight then return {} end
+
+  local ret = cache[ctx.label]
+  if not ret then
+    -- cleanup cache if it's too big
+    cache_size = cache_size + 1
+    if cache_size > MAX_CACHE_SIZE then
+      cache = {}
+      cache_size = 0
+    end
+    ret = treesitter._highlight(ctx)
+    cache[ctx.label] = ret
+  end
+
+  -- offset highlights if needed
+  if opts and opts.offset then
+    ret = vim.deepcopy(ret)
+    for _, hl in ipairs(ret) do
+      hl[1] = hl[1] + opts.offset
+      hl[2] = hl[2] + opts.offset
+    end
+  end
+  return ret
+end
+
+--- @param ctx blink.cmp.DrawItemContext
+function treesitter._highlight(ctx)
+  local ret = {} ---@type blink.cmp.DrawHighlight[]
+
+  local source = ctx.label
+  local lang = vim.treesitter.language.get_lang(vim.bo.filetype)
+  if not lang then return ret end
+
+  local ok, parser = pcall(vim.treesitter.get_string_parser, source, lang)
+  if not ok then return ret end
+
+  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')
+    -- Some injected languages may not have highlight queries.
+    if not query then return end
+
+    for capture, node in query:iter_captures(tstree:root(), source) do
+      local _, start_col, _, end_col = node:range()
+
+      ---@type string
+      local name = query.captures[capture]
+      if name ~= 'spell' then
+        ret[#ret + 1] = {
+          start_col,
+          end_col,
+          group = '@' .. name .. '.' .. lang,
+        }
+      end
+    end
+  end)
+  return ret
+end
+
+return treesitter

lua/blink/cmp/completion/windows/render/types.lua

@@ -4,6 +4,7 @@
 --- @field gap? number Gap between columns
 --- @field columns? { [number]: string, gap?: number }[] Components to render, grouped by column
 --- @field components? table<string, blink.cmp.DrawComponent> Component definitions
+--- @field treesitter? boolean Use treesitter to highlight the label text
 ---
 --- @class blink.cmp.DrawHighlight
 --- @field [number] number Start and end index of the highlight

lua/blink/cmp/config/completion/menu.lua

@@ -44,6 +44,7 @@ local window = {
       padding = 1,
       -- Gap between columns
       gap = 1,
+      treesitter = false, -- Use treesitter to highlight the label text
       -- Components to render, grouped by column
       columns = { { 'kind_icon' }, { 'label', 'label_description', gap = 1 } },
       -- Definitions for possible components to render. Each component defines:
@@ -82,6 +83,11 @@ local window = {
               table.insert(highlights, { #label, #label + #ctx.label_detail, group = 'BlinkCmpLabelDetail' })
             end
 
+            if ctx.self.treesitter then
+              -- add treesitter highlights
+              vim.list_extend(highlights, require('blink.cmp.completion.windows.render.treesitter').highlight(ctx))
+            end
+
             -- characters matched on the label by the fuzzy matcher
             for _, idx in ipairs(ctx.label_matched_indices) do
               table.insert(highlights, { idx, idx + 1, group = 'BlinkCmpLabelMatch' })