nvim-config/lua/lgwtdark/go_extra.lua

-- Extra highlights Tree-sitter cannot express alone:
-- 1) fmt verbs (%q, %s, ...) inside interpreted strings -> warm (same as numbers/params)
-- 2) rune literals: single token — paint ' delimiters and inner body
-- 3) interpreted " and raw ` delimiters white (TS anonymous quotes are unreliable)
-- 4) comment // and /* */ edge markers white
--
-- Note: pointer `*` vertical alignment vs letters is controlled by the font/terminal,
-- not Neovim highlights. Try a font with a centered asterisk (e.g. Iosevka, JetBrains Mono).

local M = {}

local ns = vim.api.nvim_create_namespace('lgwtdark_go')

-- package fmt / fmt.Errorf verbs (trailing character; %% is handled separately)
local FMT_VERB = {
  ['%'] = true,
  b = true,
  c = true,
  d = true,
  e = true,
  E = true,
  f = true,
  F = true,
  g = true,
  G = true,
  h = true,
  o = true,
  O = true,
  p = true,
  q = true,
  s = true,
  t = true,
  T = true,
  U = true,
  v = true,
  V = true,
  w = true,
  x = true,
  X = true,
}

--- Return inclusive [start, end] byte indices in `s` for each fmt verb span (package fmt style).
local function fmt_verb_spans(s)
  local spans = {}
  local i = 1
  local n = #s
  while i <= n do
    if s:byte(i) ~= 37 then -- %
      i = i + 1
    elseif i < n and s:byte(i + 1) == 37 then
      i = i + 2 -- %%
    else
      local j = i + 1
      if j <= n and s:byte(j) == 91 then -- %[n]
        local close = s:find('%]', j + 1, true)
        if not close then
          break
        end
        j = close + 1
      end
      while j <= n do
        local c = s:byte(j)
        if c == 35 or c == 43 or c == 45 or c == 32 or c == 48 then -- # + - space 0
          j = j + 1
        else
          break
        end
      end
      if j <= n and s:byte(j) == 42 then
        j = j + 1
      else
        while j <= n do
          local c = s:byte(j)
          if c >= 48 and c <= 57 then
            j = j + 1
          else
            break
          end
        end
      end
      if j <= n and s:byte(j) == 46 then
        j = j + 1
        if j <= n and s:byte(j) == 42 then
          j = j + 1
        else
          while j <= n do
            local c = s:byte(j)
            if c >= 48 and c <= 57 then
              j = j + 1
            else
              break
            end
          end
        end
      end
      if j <= n then
        local ch = s:sub(j, j)
        if FMT_VERB[ch] then
          spans[#spans + 1] = { i, j }
          i = j + 1
        else
          i = i + 1
        end
      else
        i = i + 1
      end
    end
  end
  return spans
end

local function clear_buf(bufnr)
  vim.api.nvim_buf_clear_namespace(bufnr, ns, 0, -1)
end

local function extmark_verb(bufnr, row, col_base, text, a, b)
  local start_col = col_base + #text:sub(1, a - 1)
  local end_col = col_base + #text:sub(1, b)
  vim.api.nvim_buf_set_extmark(bufnr, ns, row, start_col, {
    end_col = end_col,
    hl_group = 'LgwtdFmtVerb',
    priority = 200,
  })
end

local function hl_interpreted_double_quotes(bufnr, node)
  local text = vim.treesitter.get_node_text(node, bufnr) or ''
  if #text < 2 or text:sub(1, 1) ~= '"' or text:sub(-1) ~= '"' then
    return
  end
  local sr, sc, er, ec = node:range()
  vim.api.nvim_buf_set_extmark(bufnr, ns, sr, sc, {
    end_col = sc + 1,
    hl_group = 'LgwtdQuote',
    priority = 210,
  })
  vim.api.nvim_buf_set_extmark(bufnr, ns, er, ec - 1, {
    end_col = ec,
    hl_group = 'LgwtdQuote',
    priority = 210,
  })
end

local function hl_raw_string_delims(bufnr, node)
  local text = vim.treesitter.get_node_text(node, bufnr) or ''
  if #text < 2 or text:sub(1, 1) ~= '`' or text:sub(-1) ~= '`' then
    return
  end
  local sr, sc, er, ec = node:range()
  vim.api.nvim_buf_set_extmark(bufnr, ns, sr, sc, {
    end_col = sc + 1,
    hl_group = 'LgwtdQuote',
    priority = 210,
  })
  vim.api.nvim_buf_set_extmark(bufnr, ns, er, ec - 1, {
    end_col = ec,
    hl_group = 'LgwtdQuote',
    priority = 210,
  })
end

local function hl_comment_prefix(bufnr, node)
  local sr, sc, er, ec = node:range()
  local text = vim.treesitter.get_node_text(node, bufnr) or ''
  if text:sub(1, 2) == '//' then
    vim.api.nvim_buf_set_extmark(bufnr, ns, sr, sc, {
      end_col = sc + 2,
      hl_group = 'LgwtdQuote',
      priority = 208,
    })
  elseif text:sub(1, 2) == '/*' then
    vim.api.nvim_buf_set_extmark(bufnr, ns, sr, sc, {
      end_col = sc + 2,
      hl_group = 'LgwtdQuote',
      priority = 208,
    })
    if text:sub(-2) == '*/' then
      vim.api.nvim_buf_set_extmark(bufnr, ns, er, ec - 2, {
        end_col = ec,
        hl_group = 'LgwtdQuote',
        priority = 208,
      })
    end
  end
end

local function hl_verbs_in_interpreted_string(bufnr, str_node)
  for idx = 0, str_node:child_count() - 1 do
    local ch = str_node:child(idx)
    if ch and ch:type() == 'interpreted_string_literal_content' then
      local text = vim.treesitter.get_node_text(ch, bufnr) or ''
      local sr, sc, er, _ = ch:range()
      if sr == er then
        for _, span in ipairs(fmt_verb_spans(text)) do
          extmark_verb(bufnr, sr, sc, text, span[1], span[2])
        end
      end
    end
  end
end

-- `nil` is bundled with iota as @constant.builtin; LSP can also tag it as keyword.
-- Extmarks beat both so it matches true/false (warm).
local function hl_nil_literal(bufnr, node)
  local sr, sc, er, ec = node:range()
  vim.api.nvim_buf_set_extmark(bufnr, ns, sr, sc, {
    end_row = er,
    end_col = ec,
    hl_group = 'LgwtdNil',
    priority = 10000,
  })
end

local function hl_rune_literal(bufnr, node)
  local text = vim.treesitter.get_node_text(node, bufnr) or ''
  if #text < 2 or text:sub(1, 1) ~= "'" or text:sub(-1) ~= "'" then
    return
  end
  local sr, sc, er, ec = node:range()
  vim.api.nvim_buf_set_extmark(bufnr, ns, sr, sc, {
    end_col = sc + 1,
    hl_group = 'LgwtdQuote',
    priority = 210,
  })
  local q0 = ec - 1
  vim.api.nvim_buf_set_extmark(bufnr, ns, er, q0, {
    end_col = ec,
    hl_group = 'LgwtdQuote',
    priority = 210,
  })
  if ec - sc <= 2 then
    return
  end
  local inner = text:sub(2, -2)
  local inner_hl = inner:match('^\\') and 'LgwtdEscape' or 'LgwtdStrInner'
  vim.api.nvim_buf_set_extmark(bufnr, ns, sr, sc + 1, {
    end_col = ec - 1,
    hl_group = inner_hl,
    priority = 190,
  })
end

function M.refresh(bufnr)
  if not vim.api.nvim_buf_is_valid(bufnr) or vim.bo[bufnr].filetype ~= 'go' then
    return
  end
  clear_buf(bufnr)
  local ok, parser = pcall(vim.treesitter.get_parser, bufnr, 'go')
  if not ok or not parser then
    return
  end
  parser:parse()
  local trees = parser:trees()
  local tree = trees and trees[1]
  if not tree then
    return
  end
  local root = tree:root()
  local q = vim.treesitter.query.parse('go', '(interpreted_string_literal) @s')
  local nlines = vim.api.nvim_buf_line_count(bufnr)
  for _, node in q:iter_captures(root, bufnr, 0, nlines) do
    if node:type() == 'interpreted_string_literal' then
      hl_verbs_in_interpreted_string(bufnr, node)
      hl_interpreted_double_quotes(bufnr, node)
    end
  end
  local qraw = vim.treesitter.query.parse('go', '(raw_string_literal) @raw')
  for _, node in qraw:iter_captures(root, bufnr, 0, nlines) do
    if node:type() == 'raw_string_literal' then
      hl_raw_string_delims(bufnr, node)
    end
  end
  local qr = vim.treesitter.query.parse('go', '(rune_literal) @r')
  for _, node in qr:iter_captures(root, bufnr, 0, nlines) do
    if node:type() == 'rune_literal' then
      hl_rune_literal(bufnr, node)
    end
  end
  local qc = vim.treesitter.query.parse('go', '(comment) @c')
  for _, node in qc:iter_captures(root, bufnr, 0, nlines) do
    if node:type() == 'comment' then
      hl_comment_prefix(bufnr, node)
    end
  end
  local qnil = vim.treesitter.query.parse('go', '(nil) @n')
  for _, node in qnil:iter_captures(root, bufnr, 0, nlines) do
    if node:type() == 'nil' then
      hl_nil_literal(bufnr, node)
    end
  end
end

local scheduled = {}

function M.schedule_refresh(bufnr)
  if scheduled[bufnr] then
    return
  end
  scheduled[bufnr] = true
  vim.schedule(function()
    scheduled[bufnr] = nil
    M.refresh(bufnr)
  end)
end

local function attach_parser_cbs(bufnr)
  if vim.b[bufnr].lgwtd_go_ts_cbs then
    return
  end
  local ok, parser = pcall(vim.treesitter.get_parser, bufnr, 'go')
  if not ok or not parser or not parser.register_cbs then
    return
  end
  vim.b[bufnr].lgwtd_go_ts_cbs = true
  parser:register_cbs({
    on_changed = function()
      M.schedule_refresh(bufnr)
    end,
  })
end

local setup_done = false

function M.setup()
  if setup_done then
    return
  end
  setup_done = true
  local aug = vim.api.nvim_create_augroup('LgwtdarkGoExtra', { clear = true })
  vim.api.nvim_create_autocmd('FileType', {
    group = aug,
    pattern = 'go',
    callback = function(ev)
      M.attach(ev.buf)
    end,
  })
  vim.api.nvim_create_autocmd({ 'TextChanged', 'TextChangedI' }, {
    group = aug,
    callback = function(ev)
      if vim.bo[ev.buf].filetype == 'go' then
        M.schedule_refresh(ev.buf)
      end
    end,
  })
  vim.api.nvim_create_autocmd('ColorScheme', {
    group = aug,
    pattern = '*',
    callback = function()
      for _, buf in ipairs(vim.api.nvim_list_bufs()) do
        if vim.bo[buf].filetype == 'go' then
          M.schedule_refresh(buf)
        end
      end
    end,
  })
end

function M.attach(bufnr)
  if not pcall(vim.treesitter.start, bufnr, 'go') then
    if vim.g.lgwtd_go_parser_warned == nil then
      vim.g.lgwtd_go_parser_warned = true
      vim.notify(
        'Go Tree-sitter parser is missing. Run :Lazy sync then :TSInstall go',
        vim.log.levels.WARN,
        { title = 'lgwtdark' }
      )
    end
    return
  end
  M.refresh(bufnr)
  attach_parser_cbs(bufnr)
end

return M