Suggestion: use tree-sitter to highlight results

[fdschmidt93]

Have you RTFM'd?

• [X] ~~I have done proper research~~

Feature Request

Hi there!

A contributor and I've recently added TS Highlighting to the results buffer in my telescope-egrepify extension (see PR).

On a whim, I thought it'd be interesting to try this out for fzf-lua. Maybe I've overlooked this but I don't think fzf-lua has that yet. Anyways, here is a 15 minutes POC. It seems to work well (bold is merged with TS highlighting for queries) in 5 minutes testing and I thought you might want to integrate this properly (and more picker agnostic) :)

E: The approach should be super fast, but it won't get all highlighting right (cf. function in vim.keymap once correct and once incorrectly highlighted)

local api = vim.api

local function starts_with_pattern(line) return line:match "^.-:%d+:%d+:" ~= nil end

local function parse_line(line)
  local file_path, line_nr, col_nr, text = line:match "^(.-):(%d+):(%d+):(.*)$"
  return file_path, tonumber(line_nr), tonumber(col_nr), text
end

---@alias trouble.LangRegions table<string, number[][][]>

local TSInjector = {}

TSInjector.cache = {} ---@type table<number, table<string,{parser: vim.treesitter.LanguageTree, highlighter:vim.treesitter.highlighter, enabled:boolean}>>
local ns = vim.api.nvim_create_namespace "egrepify.highlighter"

local TSHighlighter = vim.treesitter.highlighter

local function wrap(name)
  return function(_, win, buf, ...)
    if not TSInjector.cache[buf] then
      return false
    end
    for _, hl in pairs(TSInjector.cache[buf] or {}) do
      if hl.enabled then
        TSHighlighter.active[buf] = hl.highlighter
        TSHighlighter[name](_, win, buf, ...)
      end
    end
    TSHighlighter.active[buf] = nil
  end
end

TSInjector.did_setup = false
function TSInjector.setup()
  if TSInjector.did_setup then
    return
  end
  TSInjector.did_setup = true

  vim.api.nvim_set_decoration_provider(ns, {
    on_win = wrap "_on_win",
    on_line = wrap "_on_line",
  })

  vim.api.nvim_create_autocmd("BufWipeout", {
    group = vim.api.nvim_create_augroup("egrepify.treesitter.hl", { clear = true }),
    callback = function(ev) TSInjector.cache[ev.buf] = nil end,
  })
end

---@param buf number
---@param regions trouble.LangRegions
function TSInjector.attach(buf, regions)
  TSInjector.setup()
  TSInjector.cache[buf] = TSInjector.cache[buf] or {}
  for lang in pairs(TSInjector.cache[buf]) do
    TSInjector.cache[buf][lang].enabled = regions[lang] ~= nil
  end

  for lang in pairs(regions) do
    TSInjector._attach_lang(buf, lang, regions[lang])
  end
end

---@param buf number
---@param lang? string
function TSInjector._attach_lang(buf, lang, regions)
  lang = lang or "markdown"
  lang = lang == "markdown" and "markdown_inline" or lang

  TSInjector.cache[buf] = TSInjector.cache[buf] or {}

  if not TSInjector.cache[buf][lang] then
    local ok, parser = pcall(vim.treesitter.languagetree.new, buf, lang)
    if not ok then
      return
    end
    parser:set_included_regions(regions)
    TSInjector.cache[buf][lang] = {
      parser = parser,
      highlighter = TSHighlighter.new(parser),
    }
  end
  TSInjector.cache[buf][lang].enabled = true
  local parser = TSInjector.cache[buf][lang].parser

  parser:set_included_regions(regions)
end

api.nvim_create_autocmd("FileType", {
  pattern = "fzf",
  callback = function(args)
    api.nvim_buf_attach(args.buf, false, {
      on_lines = function()
        local lines = api.nvim_buf_get_lines(args.buf, 0, -1, false)
        local regions = {}
        for i = 1, #lines do
          local line_idx = i - 1
          local line = lines[i]
          if starts_with_pattern(line) then
            local path, _, _, text = parse_line(line)
            local ft = vim.filetype.match { filename = path }

            if ft and regions[ft] == nil then
              regions[ft] = {}
            end

            if text ~= "" then
              local first_pos = string.find(line, text, 1, true)
              if first_pos == nil then
                return
              end
              first_pos = first_pos - 1
              table.insert(regions[ft], { { line_idx, first_pos, line_idx, line:len() } })
              TSInjector.attach(args.buf, regions)
            end
          end
        end
      end,
    })
  end,
})

return {
  "ibhagwan/fzf-lua",
  config = function()
    require("fzf-lua").setup {
      winopts = {
        height = 0.8, -- Full height
        width = 0.8, -- Full width
        preview = {
          layout = "vertical",
          vertical = "up:40%", -- Place preview window at the top, occupying 40% of the height
        },
        border = "single", -- Optional: sets a single border around the window
      },
      fzf_opts = {
        ["--layout"] = "reverse-list", -- Similar to `prompt_position = "top"`
      },
    }
  end,
}