1 files changed, 222 insertions, 0 deletions
diff --git a/lua/quicker/highlight.lua b/lua/quicker/highlight.lua
new file mode 100644
index 0000000..a7323aa
--- /dev/null
+++ b/lua/quicker/highlight.lua
@@ -0,0 +1,222 @@
+local M = {}
+
+---@class quicker.TSHighlight
+---@field [1] integer start_col
+---@field [2] integer end_col
+---@field [3] string highlight group
+
+local _cached_queries = {}
+---@param lang string
+---@return vim.treesitter.Query?
+local function get_highlight_query(lang)
+  local query = _cached_queries[lang]
+  if query == nil then
+    query = vim.treesitter.query.get(lang, "highlights") or false
+    _cached_queries[lang] = query
+  end
+  if query then
+    return query
+  end
+end
+
+---@param bufnr integer
+---@param lnum integer
+---@return quicker.TSHighlight[]
+function M.buf_get_ts_highlights(bufnr, lnum)
+  local filetype = vim.bo[bufnr].filetype
+  if not filetype or filetype == "" then
+    filetype = vim.filetype.match({ buf = bufnr }) or ""
+  end
+  local lang = vim.treesitter.language.get_lang(filetype) or filetype
+  if lang == "" then
+    return {}
+  end
+  local ok, parser = pcall(vim.treesitter.get_parser, bufnr, lang)
+  if not ok or not parser then
+    return {}
+  end
+
+  local row = lnum - 1
+  if not parser:is_valid() then
+    parser:parse(true)
+  end
+
+  local highlights = {}
+  parser:for_each_tree(function(tstree, tree)
+    if not tstree then
+      return
+    end
+
+    local root_node = tstree:root()
+    local root_start_row, _, root_end_row, _ = root_node:range()
+
+    -- Only worry about trees within the line range
+    if root_start_row > row or root_end_row < row then
+      return
+    end
+
+    local query = get_highlight_query(tree:lang())
+
+    -- Some injected languages may not have highlight queries.
+    if not query then
+      return
+    end
+
+    for capture, node, metadata in query:iter_captures(root_node, bufnr, row, root_end_row + 1) do
+      if capture == nil then
+        break
+      end
+
+      local range = vim.treesitter.get_range(node, bufnr, metadata[capture])
+      local start_row, start_col, _, end_row, end_col, _ = unpack(range)
+      if start_row > row then
+        break
+      end
+      local capture_name = query.captures[capture]
+      local hl = string.format("@%s.%s", capture_name, tree:lang())
+      if end_row > start_row then
+        end_col = -1
+      end
+      table.insert(highlights, { start_col, end_col, hl })
+    end
+  end)
+
+  return highlights
+end
+
+---@class quicker.LSPHighlight
+---@field [1] integer start_col
+---@field [2] integer end_col
+---@field [3] string highlight group
+---@field [4] integer priority modifier
+
+-- We're accessing private APIs here. This could break in the future.
+local STHighlighter = vim.lsp.semantic_tokens.__STHighlighter
+
+--- Copied from Neovim semantic_tokens.lua
+--- Do a binary search of the tokens in the half-open range [lo, hi).
+---
+--- Return the index i in range such that tokens[j].line < line for all j < i, and
+--- tokens[j].line >= line for all j >= i, or return hi if no such index is found.
+---
+---@private
+local function lower_bound(tokens, line, lo, hi)
+  while lo < hi do
+    local mid = bit.rshift(lo + hi, 1) -- Equivalent to floor((lo + hi) / 2).
+    if tokens[mid].line < line then
+      lo = mid + 1
+    else
+      hi = mid
+    end
+  end
+  return lo
+end
+
+---@param bufnr integer
+---@param lnum integer
+---@return quicker.LSPHighlight[]
+function M.buf_get_lsp_highlights(bufnr, lnum)
+  local highlighter = STHighlighter.active[bufnr]
+  if not highlighter then
+    return {}
+  end
+  local ft = vim.bo[bufnr].filetype
+
+  local lsp_highlights = {}
+  for _, client in pairs(highlighter.client_state) do
+    local highlights = client.current_result.highlights
+    if highlights then
+      local idx = lower_bound(highlights, lnum - 1, 1, #highlights + 1)
+      for i = idx, #highlights do
+        local token = highlights[i]
+
+        if token.line >= lnum then
+          break
+        end
+
+        table.insert(
+          lsp_highlights,
+          { token.start_col, token.end_col, string.format("@lsp.type.%s.%s", token.type, ft), 0 }
+        )
+        for modifier, _ in pairs(token.modifiers) do
+          table.insert(
+            lsp_highlights,
+            { token.start_col, token.end_col, string.format("@lsp.mod.%s.%s", modifier, ft), 1 }
+          )
+          table.insert(lsp_highlights, {
+            token.start_col,
+            token.end_col,
+            string.format("@lsp.typemod.%s.%s.%s", token.type, modifier, ft),
+            2,
+          })
+        end
+      end
+    end
+  end
+
+  return lsp_highlights
+end
+
+---@param item QuickFixItem
+---@param line string
+---@return quicker.TSHighlight[]
+M.get_heuristic_ts_highlights = function(item, line)
+  local filetype = vim.filetype.match({ buf = item.bufnr })
+  if not filetype then
+    return {}
+  end
+
+  local lang = vim.treesitter.language.get_lang(filetype)
+  if not lang then
+    return {}
+  end
+
+  local has_parser, parser = pcall(vim.treesitter.get_string_parser, line, lang)
+  if not has_parser then
+    return {}
+  end
+
+  local root = parser:parse(true)[1]:root()
+  local query = vim.treesitter.query.get(lang, "highlights")
+  if not query then
+    return {}
+  end
+
+  local highlights = {}
+  for capture, node, metadata in query:iter_captures(root, line) do
+    if capture == nil then
+      break
+    end
+
+    local range = vim.treesitter.get_range(node, line, metadata[capture])
+    local start_row, start_col, _, end_row, end_col, _ = unpack(range)
+    local capture_name = query.captures[capture]
+    local hl = string.format("@%s.%s", capture_name, lang)
+    if end_row > start_row then
+      end_col = -1
+    end
+    table.insert(highlights, { start_col, end_col, hl })
+  end
+
+  return highlights
+end
+
+function M.set_highlight_groups()
+  if vim.tbl_isempty(vim.api.nvim_get_hl(0, { name = "QuickFixHeaderHard" })) then
+    vim.api.nvim_set_hl(0, "QuickFixHeaderHard", { link = "Delimiter", default = true })
+  end
+  if vim.tbl_isempty(vim.api.nvim_get_hl(0, { name = "QuickFixHeaderSoft" })) then
+    vim.api.nvim_set_hl(0, "QuickFixHeaderSoft", { link = "Comment", default = true })
+  end
+  if vim.tbl_isempty(vim.api.nvim_get_hl(0, { name = "QuickFixFilename" })) then
+    vim.api.nvim_set_hl(0, "QuickFixFilename", { link = "Directory", default = true })
+  end
+  if vim.tbl_isempty(vim.api.nvim_get_hl(0, { name = "QuickFixFilenameInvalid" })) then
+    vim.api.nvim_set_hl(0, "QuickFixFilenameInvalid", { link = "Comment", default = true })
+  end
+  if vim.tbl_isempty(vim.api.nvim_get_hl(0, { name = "QuickFixLineNr" })) then
+    vim.api.nvim_set_hl(0, "QuickFixLineNr", { link = "LineNr", default = true })
+  end
+end
+
+return M