[mini.pick] Show framework with syntax highlighting

[drowning-cat]

Related

• [mini.pick] Syntax/treesitter highlight (~350 LOC) #2152

Similar

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

Description

I abstracted source.show as far as I could. I doubt someone in their right mind would use this code, but it's fun to try.

The framework focuses on highlighting and alignment, and is designed so you don't have to deal with extmarks offsets manually, much like how snacks.nvim handles it.

Note

Code highlighting is powered by Treesitter, so you'll need to configure it yourself.

Here is an example configuration:

function M.gen_show.buf_lines(local_opts)
  return make_show(function(item)
    local ret = {}
    if local_opts.scope == "all" then
      table.insert(ret, Slot.text(string.match(item.text, ".-%f[%z]")))
    end
    table.insert(ret, Slot.num(item.lnum))
    table.insert(ret, Slot.buf_lnum(item.bufnr, item.lnum))
    return ret
  end, { sep = "│", tab = true })
end

Format options:

• sep: string → used as a separator between slots.
• tab: boolean → to align slots by a separator using additional padding.

Slot options:

• align: "left"|"right" → text alignment inside the slot.
• min: integer → is the minimal slot length.

A Slot is the smallest unit of abstraction and is represented as a simple table. Chunks use offsets relative to the text (col and end_col default to the text boundaries).

{
  align = "left",
  min = 0,
  chunks = {
    { text = "some word", extms = { { col = 6, end_col = 9, hl_group = "ErrorMsg" } } },
    { ... },
  }
}

Predefined Slots 📦:

• Slot.text(text, hl_group, opts) Any text
• Slot.path(path) Displayed path
• Slot.num(lnum) Number (lnum/col)
• Slot.buf_lnum(buf, lnum) Highlighted line from the buffer
• Slot.code(code, path?, ft?) Highlighted code string

Code

The code is split across a few files, following the plugin/ structure while also using the lua/ directory.

There are three files:

• plugin/30_mini_pick.lua
• lua/share/mini_pick.lua
• lua/util/highlight.lua

Code is here (click me)

-- plugin/mini_pick.lua

vim.pack.add({
  { src = "https://github.com/nvim-mini/mini.pick" },
  { src = "https://github.com/nvim-mini/mini.extra" },
})

local pick_share = require("share.mini_pick")

local MiniPick = require("mini.pick")
local MiniExtra = require("mini.extra")

MiniPick.setup()

for name, gen_show in pairs(pick_share.gen_show) do
  local pick = MiniPick.builtin[name] or MiniExtra.pickers[name]
  MiniPick.registry[name] = function(local_opts, opts)
    local show = gen_show(local_opts)
    opts = vim.tbl_deep_extend("keep", opts or {}, { source = { show = show } })
    return pick(local_opts, opts)
  end
end

vim.keymap.set("n", "<Space>sf", function() MiniPick.registry.files() end, { desc = "Search files" })

-- stylua: ignore start
vim.keymap.set("n", "<Space>sg", function() MiniPick.registry.grep_live() end, { desc = "Search grep" })
vim.keymap.set("n", "<Space>sl", function() MiniPick.registry.buf_lines({ scope = "current" }) end, { desc = "Search lines (buf)" })
vim.keymap.set("n", "<Space>sL", function() MiniPick.registry.buf_lines({ scope = "all" }) end, { desc = "Search lines (all)" })
vim.keymap.set("n", "<Space>sm", function() MiniPick.registry.git_hunks() end, { desc = "Search modified hunks (all)" })
vim.keymap.set("n", "<Space>sR", function() MiniPick.registry.lsp({ scope = "references" }) end, { desc = "Search LSP refs" })
vim.keymap.set("n", "<Space>ss", function() MiniPick.registry.lsp({ scope = "document_symbol" }) end, { desc = "Search LSP symbols (doc)" })
vim.keymap.set("n", "<Space>sS", function() MiniPick.registry.lsp({ scope = "workspace_symbol_live" }) end, { desc = "Search LSP symbols (workspace)" })
vim.keymap.set("n", "<Space>sH", function() MiniPick.registry.hl_groups() end, { desc = "Search highlights" })
vim.keymap.set("n", "<Space>s:", function() MiniPick.registry.history({ scope = ":" }) end, { desc = "Search cmd history" })
vim.keymap.set("n", "<Space>s'", function() MiniPick.registry.marks() end, { desc = "Search marks" })

-- lua/share/mini_pick.lua

local M = {}

_G.MiniPick = _G.MiniPick

local util_hi = require("util.highlight")

local pick_au = vim.api.nvim_create_augroup("minipick_show", { clear = true })
local pick_ns = vim.api.nvim_create_namespace("minipick_show")

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

vim.api.nvim_create_autocmd("User", {
  pattern = "MiniPickStart",
  desc = "Set `mini.pick` display options",
  group = pick_au,
  callback = function()
    local state = MiniPick.get_picker_state()
    local win = state.windows.main
    local buf = state.buffers.main
    vim.bo[buf].tabstop = vim.go.tabstop
    vim.wo[win].listchars = vim.go.listchars
  end,
})

local minipick_cache = {}

local pick_cache = function(id, func)
  if minipick_cache[id] == nil then
    minipick_cache[id] = { func() }
  end
  ---@diagnostic disable-next-line: redundant-return-value
  return unpack(minipick_cache[id])
end

vim.api.nvim_create_autocmd("User", {
  pattern = "MiniPickStop",
  desc = "Clear `mini.pick` cache",
  group = pick_au,
  callback = function()
    minipick_cache = {}
  end,
})

-- Format

---@class ChunkExtm : vim.api.keyset.set_extmark
---@field row? integer
---@field col? integer

---@class Chunk
---@field text string
---@field extms ChunkExtm[]

--- ```lua
--- {
---   align = "left",
---   min = 0,
---   chunks = {
---     { text = "some word", extms = { { col = 6, end_col = 9, hl_group = "ErrorMsg" }, ... } },
---     { ... },
---   }
--- }
--- ```
---@class Slot
---@field align? "left"|"right"
---@field min? integer
---@field chunks Chunk[]

---@class FormatOpts
---@field sep? string
---@field tbl? boolean

---@param data Slot[][]
---@param opts? FormatOpts
local format = function(buf, data, opts)
  opts = vim.tbl_extend("keep", opts or {}, { sep = "│", tbl = true })
  -- Pre-process
  local slots = {} ---@type Slot[][]
  local slot_widths, slot_max = {}, 0
  for _, row_slots in ipairs(data) do
    local new_row = {}
    for _, slot in ipairs(row_slots) do
      if type(slot) == "table" then
        local slot_text, slot_chunk = "", {}
        for _, chunk in ipairs(slot.chunks or {}) do
          if type(chunk) == "table" and not vim.tbl_isempty(chunk) then
            local text = tostring(chunk.text or "")
            slot_text = slot_text .. text
            ---@type Chunk
            local new_chunk = { text = text, extms = vim.deepcopy(chunk.extms or {}) }
            table.insert(slot_chunk, new_chunk)
          end
        end
        local width = vim.fn.strdisplaywidth(slot_text)
        ---@type Slot
        local new_slot = vim.tbl_extend("force", { align = "left", min = 0 }, slot, {
          chunks = slot_chunk,
          _text = slot_text,
          _text_width = width,
        })
        table.insert(new_row, new_slot)
        local col = #new_row
        slot_widths[col] = math.max(slot_widths[col] or 0, new_slot.min, width)
      end
    end
    table.insert(slots, new_row)
    slot_max = math.max(slot_max, #new_row)
  end
  -- Pre-render
  local buf_lines, buf_extms = {}, {}
  for row, row_slots in ipairs(slots) do
    local line = ""
    for col, slot in ipairs(row_slots) do
      if col > 1 then
        line = line .. opts.sep
      end
      ---@diagnostic disable-next-line: undefined-field
      local text, text_width = slot._text, slot._text_width
      local width = math.max(slot.min --[[@as integer]], text_width)
      if opts.tbl and slot_max == #row_slots then
        width = slot_widths[col]
      end
      local pad = string.rep(" ", width - text_width)
      local left, right = "", ""
      if slot.align == "right" then
        left = pad
      end
      if slot.align == "left" and col < #row_slots then
        right = pad
      end
      local text_start = #line + #left
      local off = 0
      for _, chunk in ipairs(slot.chunks) do
        local shift = text_start + off
        for _, extm in ipairs(chunk.extms) do
          extm = vim.tbl_extend("force", extm, {
            row = row - 1,
            end_row = row - 1,
            col = (extm.col or 0) + shift,
            end_col = (extm.end_col and extm.end_col ~= -1) and (extm.end_col + shift) or (shift + #chunk.text),
          })
          table.insert(buf_extms, extm)
        end
        off = off + #chunk.text
      end
      line = line .. (left .. text .. right)
    end
    table.insert(buf_lines, line)
  end
  -- Render
  vim.api.nvim_buf_set_lines(buf, 0, -1, false, buf_lines)
  vim.api.nvim_buf_clear_namespace(buf, pick_ns, 0, -1)
  for _, extm in ipairs(buf_extms) do
    local row, col = extm.row or 0, extm.col or 0
    extm.row, extm.col = nil, nil
    vim.api.nvim_buf_set_extmark(buf, pick_ns, row, col, extm)
  end
end

local Slot = {}

function Slot.text(text, hl_group, opts)
  return vim.tbl_extend("force", opts or {}, {
    chunks = {
      { text = tostring(text or ""), extms = { { hl_group = hl_group } } },
    },
  } --[[@as Slot]])
end

function Slot.path(path)
  if #path > 30 then
    path = vim.fn.pathshorten(path)
  end
  return Slot.text(path)
end

function Slot.num(lnum)
  return Slot.text(lnum, nil, { align = "right" })
end

local hl2extm = function(hl)
  return vim.tbl_extend("force", hl.extm, { col = hl.col, row = nil })
end

function Slot.buf_lnum(buf, lnum)
  return pick_cache(vim.inspect({ "buf_lnum", buf, lnum }), function()
    local text = vim.api.nvim_buf_get_lines(buf, lnum - 1, lnum, true)[1]
    local extms = {}
    local row_hls = util_hi.get_buf_hls(buf, true, { lnum - 1, lnum })[1] or {}
    for _, hl in ipairs(row_hls) do
      if hl.extm.hl_group then
        table.insert(extms, hl2extm(hl))
      end
    end
    ---@type Slot
    return { chunks = { { text = text, extms = extms } } }
  end)
end

function Slot.code(code, path, ft)
  assert(not string.find(code, "\n"), "Code has new line")
  return pick_cache(vim.inspect({ "code", code, ft or path }), function()
    ft = ft or vim.filetype.match({ filename = path, contents = { code } })
    local extms = {}
    local row_hls = util_hi.get_text_patched_hls(code, ft)[1] or {}
    for _, hl in ipairs(row_hls) do
      table.insert(extms, hl2extm(hl))
    end
    ---@type Slot
    return { chunks = { { text = code, extms = extms } } }
  end)
end

M.gen_show = {}

local safe_show = function(show_fn)
  return function(buf, data, query, opts)
    xpcall(function()
      show_fn(buf, data, query, opts)
    end, function(err)
      vim.api.nvim_win_close(0, true)
      vim.api.nvim_echo({ { debug.traceback(err), "ErrorMsg" } }, true, {})
    end)
  end
end

---@param item_cb fun(item,i,buf,data,query): Slot[]
---@param opts? FormatOpts
local make_show = function(item_cb, opts)
  return safe_show(function(buf, data, query)
    local ret = {}
    for i, item in ipairs(data) do
      local row = item_cb(item, i, buf, data, query) or {}
      table.insert(ret, row)
    end
    format(buf, ret, opts)
  end)
end

function M.gen_show.buf_lines(local_opts)
  return make_show(function(item)
    local ret = {}
    if local_opts.scope == "all" then
      table.insert(ret, Slot.text(string.match(item.text, ".-%f[%z]")))
    end
    table.insert(ret, Slot.num(item.lnum))
    table.insert(ret, Slot.buf_lnum(item.bufnr, item.lnum))
    return ret
  end, { tbl = true })
end

function M.gen_show.git_hunks()
  local extract_code = function(item)
    for _, hunk_part in ipairs(item.hunk) do
      local sign, code = string.match(hunk_part, "^([+-])(.*)$")
      if sign and vim.trim(code) ~= "" then
        return sign, code
      end
    end
    return nil, item.hunk[2]
  end
  local code_slot = function(item)
    local sign, code = extract_code(item)
    local slot = vim.deepcopy(Slot.code(code, item.path))
    if sign then
      table.insert(slot, 1, { text = sign .. " " })
    end
    return slot
  end
  return make_show(function(item)
    local header = string.match(item.hunk[1], "@@ (.-) @@")
    return {
      Slot.path(item.path),
      Slot.text(" " .. header .. " "),
      code_slot(item),
    }
  end, { tbl = true })
end

local show_grep = function()
  return make_show(function(item)
    local path, lnum, col, code = unpack(vim.split(item, "%z"))
    return {
      Slot.path(path),
      Slot.num(lnum),
      Slot.num(col),
      Slot.code(code, path),
    }
  end)
end

M.gen_show.grep = show_grep
M.gen_show.grep_live = show_grep

function M.gen_show.hipatterns()
  return make_show(function(item)
    local path = vim.trim(string.match(item.text, ".-│(.-)│"))
    return {
      Slot.text(item.highlighter, item.hl_group),
      Slot.path(path),
      Slot.num(item.lnum),
      Slot.num(item.col),
      Slot.buf_lnum(item.bufnr, item.lnum),
    }
  end)
end

function M.gen_show.history()
  return make_show(function(item)
    if vim.startswith(item, ":") then
      return { Slot.code(item, nil, "vim") }
    else
      return { Slot.text(item) }
    end
  end)
end

function M.gen_show.lsp(local_opts)
  return make_show(function(item)
    if local_opts.scope == "document_symbol" then
      return { Slot.text(item.text, item.hl) }
    end
    local ret = {
      Slot.path(item.path),
      Slot.num(item.lnum),
      Slot.num(item.col),
    }
    local text = string.match(item.text, "^.-│.-│.-│(.*)")
    if local_opts.scope == "workspace_symbol_live" then
      local icon = string.match(item.text, "^%S+")
      table.insert(ret, Slot.text(" " .. icon .. text, item.hl))
    else
      table.insert(ret, Slot.code(text, item.path))
    end
    return ret
  end)
end

function M.gen_show.marks(local_opts)
  local_opts = vim.tbl_extend("keep", local_opts or {}, { scope = "buf" })
  local io_peek_line = function(path, lnum)
    if not vim.uv.fs_stat(path) then
      return
    end
    local ln = 0
    for line in io.lines(path) do
      ln = ln + 1
      if ln == lnum then
        return line
      end
    end
  end
  local code_slot = function(item)
    if item.bufnr then
      return Slot.buf_lnum(item.bufnr, item.lnum)
    end
    if item.path then
      local cache_id = vim.inspect({ "ln", item.path, item.lnum })
      local code = pick_cache(cache_id, function()
        return io_peek_line(item.path, item.lnum)
      end)
      return Slot.code(code or "", item.path)
    end
    return Slot.text()
  end
  local get_target_name = function()
    return pick_cache("buf_name", function()
      local target_win = MiniPick.get_picker_state().windows.target
      local target_buf = vim.api.nvim_win_get_buf(target_win)
      return vim.api.nvim_buf_get_name(target_buf)
    end)
  end
  return make_show(function(item)
    local ret = {}
    table.insert(ret, Slot.text(string.sub(item.text, 1, 1)))
    table.insert(ret, Slot.num(item.lnum))
    if local_opts.scope ~= "buf" then
      table.insert(ret, Slot.path(item.path or get_target_name()))
    end
    table.insert(ret, code_slot(item))
    return ret
  end, { tbl = true })
end

return M

-- lua/util/highlight.lua

local M = {}

local Range = require("vim.treesitter._range")

---@class HiExtmark
---@field row integer -- 0-based
---@field col integer -- 0-based
---@field extm vim.api.keyset.get_extmark

---@param query_range Range4? 0-based, end-exclusive
---@return HiExtmark[]
function M.get_ts_hls(source, lang, parse_range, query_range)
  parse_range = parse_range == nil and true or parse_range
  query_range = query_range
  local is_range4 = function(r)
    return r[1] and r[2] and r[3] and r[4]
  end
  if query_range and not is_range4(query_range) then
    error("`query_range` must be Range4|nil")
  end
  local parser, err ---@type vim.treesitter.LanguageTree?, string?
  if type(source) == "number" then
    parser, err = vim.treesitter.get_parser(source, lang)
  end
  if type(source) == "string" then
    parser, err = vim.treesitter.get_string_parser(source, lang)
  end
  if not parser then
    vim.notify(err or "Unable to get a parser", vim.log.levels.WARN)
    return {}
  end
  local ret = {}
  local prior_vals = {}
  parser:parse(parse_range)
  parser:for_each_tree(function(tstree, tree)
    if not tstree then
      return
    end
    local root_node = tstree:root()
    local root_range = { root_node:range() }
    local tsquery_range = query_range or root_range
    if not Range.intersection(root_range, tsquery_range) then
      return
    end
    local query = vim.treesitter.query.get(tree:lang(), "highlights")
    if not query then
      return
    end
    for id, node, metadata in
      query:iter_captures(
        root_node,
        source,
        tsquery_range[1],
        tsquery_range[3],
        { start_col = tsquery_range[2], end_col = tsquery_range[4] }
      )
    do
      local capture = query.captures[id]
      if capture ~= nil and capture ~= "spell" then
        local node_text = vim.treesitter.get_node_text(node, source, metadata[id])
        local text = vim.split(node_text, "\n")
        local range = { node:range() }
        local inter = Range.intersection(range, tsquery_range)
        if not inter then
          return
        end
        local col_start, col_end = inter[2], inter[4]
        local row_start, row_end = inter[1], col_end == 0 and inter[3] - 1 or inter[3]
        for row = row_start, row_end do
          local first, last = row == row_start, row == row_end
          local line = text[row - row_start + 1] or ""
          local col = first and col_start or 0
          local end_col = last and col_end or #line
          -- HACK: +1 `priority` for repeated columns
          local prior = tonumber(metadata.priority) or 100
          local prior_id = row .. ":" .. col
          prior = prior_vals[prior_id] and prior_vals[prior_id] + 1 or prior
          prior_vals[prior_id] = prior
          --
          ret[row] = ret[row] or {}
          table.insert(ret[row], {
            row = row,
            col = col,
            extm = {
              end_row = row,
              end_col = end_col,
              priority = prior,
              conceal = metadata.conceal,
              hl_group = "@" .. capture .. "." .. lang,
            },
          } --[[@as HiExtmark]])
        end
      end
    end
  end)
  local hls_flat = function(hls_tbl)
    local keys = vim.tbl_keys(hls_tbl)
    table.sort(keys)
    return vim.tbl_map(function(key)
      return hls_tbl[key]
    end, keys)
  end
  -- NOTE: Return highlights grouped by line
  ---@cast ret HiExtmark[][]
  ret = hls_flat(ret)
  return ret
end

function M.get_lang(ft)
  local lang = vim.treesitter.language.get_lang(ft or "")
  if not lang or not vim.treesitter.language.add(lang) then
    return nil
  end
  if not vim.api.nvim_get_runtime_file("queries/" .. lang .. "/highlights.scm", false) then
    return nil
  end
  return lang
end

function M.get_text_patched_hls(text, ft)
  text = type(text) == "table" and table.concat(text, "\n") or text
  local lang = M.get_lang(ft)
  if not lang then
    return {}
  end
  -- HACK: Complete keywords in order to get highlights
  if lang == "lua" then
    local prepend = { ["end"] = "do", ["until"] = "repeat" }
    local append = { ["do"] = "end", ["then"] = "end", ["repeat"] = "until" }
    local lines = vim.split(text, "\n")
    local first, last = lines[1], lines[#lines]
    local from, to = 0, #lines
    for _, word in ipairs(vim.split(first, "%s+")) do
      local add = prepend[word]
      if add then
        table.insert(lines, 1, add)
        from, to = from + 1, to + 1
      end
    end
    for _, word in ipairs(vim.split(last, "%s+")) do
      local add = append[word]
      if add then
        table.insert(lines, add)
      end
    end
    local patched_text = table.concat(lines, "\n")
    return M.get_ts_hls(patched_text, lang, true, { from, 0, to, 0 })
  end
  return M.get_ts_hls(text, lang)
end

function M.get_buf_hls(buf, with_extmarks, ln_range)
  if not vim.api.nvim_buf_is_valid(buf) then
    error(string.format("Invalid buffer: %s", buf))
  end
  ln_range[1] = ln_range[1] or 0
  ln_range[2] = ln_range[2] or vim.api.nvim_buf_line_count(buf)
  local lang = M.get_lang(vim.bo[buf].ft)
  if not lang then
    return {}
  end
  local range4 = { ln_range[1], 0, ln_range[2], 0 }
  local hls = M.get_ts_hls(buf, lang, range4, range4) or {}
  if with_extmarks then
    local start_pos, end_pos = { ln_range[1], 0 }, { ln_range[2] - 1, -1 }
    local extm_items = vim.api.nvim_buf_get_extmarks(buf, -1, start_pos, end_pos, { details = true })
    for _, extm_info in ipairs(extm_items) do
      local _, row, col, extm = unpack(extm_info)
      local ln = ln_range[1] - row + 1
      hls[ln] = hls[ln] or {}
      if extm then
        extm.sign_name = nil
        extm.sign_text = nil
        extm.ns_id = nil
        extm.end_row = nil
        if not vim.tbl_contains({ "eol", "overlay", "right_align", "inline" }, extm.virt_text_pos) then
          extm.virt_text = nil
          extm.virt_text_pos = nil
        end
        table.insert(hls[ln], { row = row, col = col, extm = extm })
      end
    end
  end
  return hls
end

return M

[echasnovski]

Thanks for sharing! That's a massive amount of work!

Yeah, this is indeed a bit too complex for 'mini.pick' itself. But it's good to have a hobby to tweak a Neovim config to this degree, I guess :)