feat(highlight): add treesitter context padding from disk

Problem: treesitter parses each diff hunk in isolation, so incomplete
syntax constructs at hunk boundaries (e.g., a function definition with
no body) produce ERROR nodes and drop captures.

Solution: read N lines from the on-disk file before/after each hunk and
prepend/append them as unmapped padding lines. The line_map guard in
highlight_treesitter skips extmarks for unmapped lines, so padding
provides syntax context without visual output. Controlled by
highlights.context (default 25, 0 to disable). Also applies to the vim
syntax fallback path via a leading_offset filter.
This commit is contained in:
Barrett Ruth 2026-02-07 13:05:53 -05:00
parent ba1f830629
commit 2e1ebdee03
7 changed files with 308 additions and 24 deletions

View file

@ -3,6 +3,33 @@ local M = {}
local dbg = require('diffs.log').dbg
local diff = require('diffs.diff')
---@param filepath string
---@param from_line integer
---@param count integer
---@return string[]
local function read_line_range(filepath, from_line, count)
if count <= 0 then
return {}
end
local f = io.open(filepath, 'r')
if not f then
return {}
end
local result = {}
local line_num = 0
for line in f:lines() do
line_num = line_num + 1
if line_num >= from_line then
table.insert(result, line)
if #result >= count then
break
end
end
end
f:close()
return result
end
local PRIORITY_CLEAR = 198
local PRIORITY_SYNTAX = 199
local PRIORITY_LINE_BG = 200
@ -177,8 +204,9 @@ end
---@param hunk diffs.Hunk
---@param code_lines string[]
---@param covered_lines? table<integer, true>
---@param leading_offset? integer
---@return integer
local function highlight_vim_syntax(bufnr, ns, hunk, code_lines, covered_lines)
local function highlight_vim_syntax(bufnr, ns, hunk, code_lines, covered_lines, leading_offset)
local ft = hunk.ft
if not ft then
return 0
@ -188,6 +216,8 @@ local function highlight_vim_syntax(bufnr, ns, hunk, code_lines, covered_lines)
return 0
end
leading_offset = leading_offset or 0
local scratch = vim.api.nvim_create_buf(false, true)
vim.api.nvim_buf_set_lines(scratch, 0, -1, false, code_lines)
vim.api.nvim_set_option_value('bufhidden', 'wipe', { buf = scratch })
@ -214,17 +244,21 @@ local function highlight_vim_syntax(bufnr, ns, hunk, code_lines, covered_lines)
vim.api.nvim_buf_delete(scratch, { force = true })
local hunk_line_count = #hunk.lines
local extmark_count = 0
for _, span in ipairs(spans) do
local buf_line = hunk.start_line + span.line - 1
pcall(vim.api.nvim_buf_set_extmark, bufnr, ns, buf_line, span.col_start, {
end_col = span.col_end,
hl_group = span.hl_name,
priority = PRIORITY_SYNTAX,
})
extmark_count = extmark_count + 1
if covered_lines then
covered_lines[buf_line] = true
local adj = span.line - leading_offset
if adj >= 1 and adj <= hunk_line_count then
local buf_line = hunk.start_line + adj - 1
pcall(vim.api.nvim_buf_set_extmark, bufnr, ns, buf_line, span.col_start, {
end_col = span.col_end,
hl_group = span.hl_name,
priority = PRIORITY_SYNTAX,
})
extmark_count = extmark_count + 1
if covered_lines then
covered_lines[buf_line] = true
end
end
end
@ -255,6 +289,20 @@ function M.highlight_hunk(bufnr, ns, hunk, opts)
---@type table<integer, true>
local covered_lines = {}
local context = opts.highlights.context or 0
local leading = {}
local trailing = {}
if (use_ts or use_vim) and context > 0 and hunk.file_new_start and hunk.repo_root then
local filepath = vim.fs.joinpath(hunk.repo_root, hunk.filename)
local lead_from = math.max(1, hunk.file_new_start - context)
local lead_count = hunk.file_new_start - lead_from
if lead_count > 0 then
leading = read_line_range(filepath, lead_from, lead_count)
end
local trail_from = hunk.file_new_start + (hunk.file_new_count or 0)
trailing = read_line_range(filepath, trail_from, context)
end
local extmark_count = 0
if use_ts then
---@type string[]
@ -266,6 +314,11 @@ function M.highlight_hunk(bufnr, ns, hunk, opts)
---@type table<integer, integer>
local old_map = {}
for _, pad_line in ipairs(leading) do
table.insert(new_code, pad_line)
table.insert(old_code, pad_line)
end
for i, line in ipairs(hunk.lines) do
local prefix = line:sub(1, 1)
local stripped = line:sub(2)
@ -284,6 +337,11 @@ function M.highlight_hunk(bufnr, ns, hunk, opts)
end
end
for _, pad_line in ipairs(trailing) do
table.insert(new_code, pad_line)
table.insert(old_code, pad_line)
end
extmark_count = highlight_treesitter(bufnr, ns, new_code, hunk.lang, new_map, 1, covered_lines)
extmark_count = extmark_count
+ highlight_treesitter(bufnr, ns, old_code, hunk.lang, old_map, 1, covered_lines)
@ -305,10 +363,16 @@ function M.highlight_hunk(bufnr, ns, hunk, opts)
elseif use_vim then
---@type string[]
local code_lines = {}
for _, pad_line in ipairs(leading) do
table.insert(code_lines, pad_line)
end
for _, line in ipairs(hunk.lines) do
table.insert(code_lines, line:sub(2))
end
extmark_count = highlight_vim_syntax(bufnr, ns, hunk, code_lines, covered_lines)
for _, pad_line in ipairs(trailing) do
table.insert(code_lines, pad_line)
end
extmark_count = highlight_vim_syntax(bufnr, ns, hunk, code_lines, covered_lines, #leading)
end
if

View file

@ -14,6 +14,7 @@
---@class diffs.Highlights
---@field background boolean
---@field gutter boolean
---@field context integer
---@field treesitter diffs.TreesitterConfig
---@field vim diffs.VimConfig
---@field intra diffs.IntraConfig
@ -80,6 +81,7 @@ local default_config = {
highlights = {
background = true,
gutter = true,
context = 25,
treesitter = {
enabled = true,
max_lines = 500,
@ -231,6 +233,7 @@ local function init()
vim.validate({
['highlights.background'] = { opts.highlights.background, 'boolean', true },
['highlights.gutter'] = { opts.highlights.gutter, 'boolean', true },
['highlights.context'] = { opts.highlights.context, 'number', true },
['highlights.treesitter'] = { opts.highlights.treesitter, 'table', true },
['highlights.vim'] = { opts.highlights.vim, 'table', true },
['highlights.intra'] = { opts.highlights.intra, 'table', true },
@ -291,6 +294,9 @@ local function init()
if opts.debounce_ms and opts.debounce_ms < 0 then
error('diffs: debounce_ms must be >= 0')
end
if opts.highlights and opts.highlights.context and opts.highlights.context < 0 then
error('diffs: highlights.context must be >= 0')
end
if
opts.highlights
and opts.highlights.treesitter

View file

@ -8,6 +8,11 @@
---@field lines string[]
---@field header_start_line integer?
---@field header_lines string[]?
---@field file_old_start integer?
---@field file_old_count integer?
---@field file_new_start integer?
---@field file_new_count integer?
---@field repo_root string?
local M = {}
@ -132,6 +137,14 @@ function M.parse_buffer(bufnr)
local header_start = nil
---@type string[]
local header_lines = {}
---@type integer?
local file_old_start = nil
---@type integer?
local file_old_count = nil
---@type integer?
local file_new_start = nil
---@type integer?
local file_new_count = nil
local function flush_hunk()
if hunk_start and #hunk_lines > 0 then
@ -143,6 +156,11 @@ function M.parse_buffer(bufnr)
header_context = hunk_header_context,
header_context_col = hunk_header_context_col,
lines = hunk_lines,
file_old_start = file_old_start,
file_old_count = file_old_count,
file_new_start = file_new_start,
file_new_count = file_new_count,
repo_root = repo_root,
}
if hunk_count == 1 and header_start and #header_lines > 0 then
hunk.header_start_line = header_start
@ -154,6 +172,10 @@ function M.parse_buffer(bufnr)
hunk_header_context = nil
hunk_header_context_col = nil
hunk_lines = {}
file_old_start = nil
file_old_count = nil
file_new_start = nil
file_new_count = nil
end
for i, line in ipairs(lines) do
@ -174,6 +196,13 @@ function M.parse_buffer(bufnr)
elseif line:match('^@@.-@@') then
flush_hunk()
hunk_start = i
local hs, hc, hs2, hc2 = line:match('^@@ %-(%d+),?(%d*) %+(%d+),?(%d*) @@')
if hs then
file_old_start = tonumber(hs)
file_old_count = tonumber(hc) or 1
file_new_start = tonumber(hs2)
file_new_count = tonumber(hc2) or 1
end
local prefix, context = line:match('^(@@.-@@%s*)(.*)')
if context and context ~= '' then
hunk_header_context = context