Merge pull request #85 from barrettruth/feat/context-padding

feat(highlight): add treesitter context padding from disk

README.md

@@ -41,12 +41,11 @@ luarocks install diffs.nvim
 
 ## Known Limitations
 
-- **Incomplete syntax context**: Treesitter parses each diff hunk in isolation
+- **Incomplete syntax context**: Treesitter parses each diff hunk in isolation.
-  without surrounding code context. When a hunk shows lines added to an existing
+  To improve accuracy, `diffs.nvim` reads lines from disk before and after each
-  block (e.g., adding a plugin inside `return { ... }`), the parser doesn't see
+  hunk for parsing context (`highlights.context`, enabled by default with 25
-  the `return` statement and may produce incorrect highlighting. This is
+  lines). This resolves most boundary issues. Set
-  inherent to parsing code fragments—no diff tooling solves this without
+  `highlights.context.enabled = false` to disable.
-  significant complexity.
 
 - **Syntax flashing**: `diffs.nvim` hooks into the `FileType fugitive` event
   triggered by `vim-fugitive`, at which point the buffer is preliminarily

doc/diffs.nvim.txt

@@ -56,6 +56,10 @@ Configuration is done via `vim.g.diffs`. Set this before the plugin loads:
       highlights = {
         background = true,
         gutter = true,
+        context = {
+          enabled = true,
+          lines = 25,
+        },
         treesitter = {
           enabled = true,
           max_lines = 500,
@@ -113,6 +117,10 @@ Configuration is done via `vim.g.diffs`. Set this before the plugin loads:
                              Highlight line numbers with matching colors.
                              Only visible if line numbers are enabled.
 
+        {context}            (table, default: see below)
+                             Syntax parsing context options.
+                             See |diffs.ContextConfig| for fields.
+
         {treesitter}         (table, default: see below)
                              Treesitter highlighting options.
                              See |diffs.TreesitterConfig| for fields.
@@ -125,6 +133,20 @@ Configuration is done via `vim.g.diffs`. Set this before the plugin loads:
                              Character-level (intra-line) diff highlighting.
                              See |diffs.IntraConfig| for fields.
 
+                                                      *diffs.ContextConfig*
+    Context config fields: ~
+        {enabled}            (boolean, default: true)
+                             Read lines from disk before and after each hunk
+                             to provide surrounding syntax context. Improves
+                             accuracy at hunk boundaries where incomplete
+                             constructs (e.g., a function definition with no
+                             body) would otherwise confuse the parser.
+
+        {lines}              (integer, default: 25)
+                             Number of context lines to read in each
+                             direction. Lines are read with early exit —
+                             cost scales with this value, not file size.
+
                                                     *diffs.TreesitterConfig*
     Treesitter config fields: ~
         {enabled}            (boolean, default: true)
@@ -305,14 +327,15 @@ KNOWN LIMITATIONS                                       *diffs-limitations*
 
 Incomplete Syntax Context ~
                                                       *diffs-syntax-context*
-Treesitter parses each diff hunk in isolation without surrounding code
+Treesitter parses each diff hunk in isolation. To provide surrounding code
-context. When a hunk shows lines added to an existing block (e.g., adding a
+context, diffs.nvim reads lines from disk before and after each hunk
-plugin inside `return { ... }`), the parser doesn't see the `return`
+(see |diffs.ContextConfig|, enabled by default). This resolves most boundary
-statement and may produce incorrect or unusual highlighting.
+issues where incomplete constructs (e.g., a function definition at the edge
+of a hunk with no body) would confuse the parser.
 
-This is inherent to parsing code fragments. No diff tooling solves this
+Set `highlights.context.enabled = false` to disable context padding. In rare
-problem without significant complexity—the parser simply doesn't have enough
+cases, context padding may not help if the relevant surrounding code is very
-information to understand the full syntactic structure.
+far from the hunk boundaries.
 
 Syntax Highlighting Flash ~
                                                              *diffs-flash*

lua/diffs/highlight.lua

@@ -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
+    local adj = span.line - leading_offset
-    pcall(vim.api.nvim_buf_set_extmark, bufnr, ns, buf_line, span.col_start, {
+    if adj >= 1 and adj <= hunk_line_count then
-      end_col = span.col_end,
+      local buf_line = hunk.start_line + adj - 1
-      hl_group = span.hl_name,
+      pcall(vim.api.nvim_buf_set_extmark, bufnr, ns, buf_line, span.col_start, {
-      priority = PRIORITY_SYNTAX,
+        end_col = span.col_end,
-    })
+        hl_group = span.hl_name,
-    extmark_count = extmark_count + 1
+        priority = PRIORITY_SYNTAX,
-    if covered_lines then
+      })
-      covered_lines[buf_line] = true
+      extmark_count = extmark_count + 1
+      if covered_lines then
+        covered_lines[buf_line] = true
+      end
     end
   end
 
@@ -255,6 +289,21 @@ function M.highlight_hunk(bufnr, ns, hunk, opts)
   ---@type table<integer, true>
   local covered_lines = {}
 
+  local ctx_cfg = opts.highlights.context
+  local context = (ctx_cfg and ctx_cfg.enabled) and ctx_cfg.lines 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 +315,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 +338,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 +364,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

lua/diffs/init.lua

@@ -11,9 +11,14 @@
 ---@field algorithm string
 ---@field max_lines integer
 
+---@class diffs.ContextConfig
+---@field enabled boolean
+---@field lines integer
+
 ---@class diffs.Highlights
 ---@field background boolean
 ---@field gutter boolean
+---@field context diffs.ContextConfig
 ---@field treesitter diffs.TreesitterConfig
 ---@field vim diffs.VimConfig
 ---@field intra diffs.IntraConfig
@@ -80,6 +85,10 @@ local default_config = {
   highlights = {
     background = true,
     gutter = true,
+    context = {
+      enabled = true,
+      lines = 25,
+    },
     treesitter = {
       enabled = true,
       max_lines = 500,
@@ -231,11 +240,19 @@ local function init()
     vim.validate({
       ['highlights.background'] = { opts.highlights.background, 'boolean', true },
       ['highlights.gutter'] = { opts.highlights.gutter, 'boolean', true },
+      ['highlights.context'] = { opts.highlights.context, 'table', true },
       ['highlights.treesitter'] = { opts.highlights.treesitter, 'table', true },
       ['highlights.vim'] = { opts.highlights.vim, 'table', true },
       ['highlights.intra'] = { opts.highlights.intra, 'table', true },
     })
 
+    if opts.highlights.context then
+      vim.validate({
+        ['highlights.context.enabled'] = { opts.highlights.context.enabled, 'boolean', true },
+        ['highlights.context.lines'] = { opts.highlights.context.lines, 'number', true },
+      })
+    end
+
     if opts.highlights.treesitter then
       vim.validate({
         ['highlights.treesitter.enabled'] = { opts.highlights.treesitter.enabled, 'boolean', true },
@@ -291,6 +308,14 @@ 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.lines
+    and opts.highlights.context.lines < 0
+  then
+    error('diffs: highlights.context.lines must be >= 0')
+  end
   if
     opts.highlights
     and opts.highlights.treesitter

lua/diffs/parser.lua

@@ -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

spec/highlight_spec.lua

@@ -37,6 +37,7 @@ describe('highlight', function()
         highlights = {
           background = false,
           gutter = false,
+          context = { enabled = false, lines = 0 },
           treesitter = {
             enabled = true,
             max_lines = 500,
@@ -1055,6 +1056,114 @@ describe('highlight', function()
       assert.is_true(min_line_bg < min_char_bg)
       delete_buffer(bufnr)
     end)
+
+    it('context padding produces no extmarks on padding lines', function()
+      local repo_root = '/tmp/diffs-test-context'
+      vim.fn.mkdir(repo_root, 'p')
+
+      local f = io.open(repo_root .. '/test.lua', 'w')
+      f:write('local M = {}\n')
+      f:write('function M.hello()\n')
+      f:write('  return "hi"\n')
+      f:write('end\n')
+      f:write('return M\n')
+      f:close()
+
+      local bufnr = create_buffer({
+        '@@ -3,1 +3,2 @@',
+        ' return "hi"',
+        '+"bye"',
+      })
+
+      local hunk = {
+        filename = 'test.lua',
+        lang = 'lua',
+        start_line = 1,
+        lines = { '  return "hi"', '+"bye"' },
+        file_old_start = 3,
+        file_old_count = 1,
+        file_new_start = 3,
+        file_new_count = 2,
+        repo_root = repo_root,
+      }
+
+      highlight.highlight_hunk(
+        bufnr,
+        ns,
+        hunk,
+        default_opts({ highlights = { context = { enabled = true, lines = 25 } } })
+      )
+
+      local extmarks = get_extmarks(bufnr)
+      for _, mark in ipairs(extmarks) do
+        local row = mark[2]
+        assert.is_true(row >= 1 and row <= 2)
+      end
+
+      delete_buffer(bufnr)
+      os.remove(repo_root .. '/test.lua')
+      vim.fn.delete(repo_root, 'rf')
+    end)
+
+    it('context disabled matches behavior without padding', function()
+      local bufnr = create_buffer({
+        '@@ -1,1 +1,2 @@',
+        ' local x = 1',
+        '+local y = 2',
+      })
+
+      local hunk = {
+        filename = 'test.lua',
+        lang = 'lua',
+        start_line = 1,
+        lines = { ' local x = 1', '+local y = 2' },
+        file_new_start = 1,
+        file_new_count = 2,
+        repo_root = '/nonexistent',
+      }
+
+      highlight.highlight_hunk(
+        bufnr,
+        ns,
+        hunk,
+        default_opts({ highlights = { context = { enabled = false, lines = 0 } } })
+      )
+
+      local extmarks = get_extmarks(bufnr)
+      assert.is_true(#extmarks > 0)
+      delete_buffer(bufnr)
+    end)
+
+    it('gracefully handles missing file for context padding', function()
+      local bufnr = create_buffer({
+        '@@ -1,1 +1,2 @@',
+        ' local x = 1',
+        '+local y = 2',
+      })
+
+      local hunk = {
+        filename = 'test.lua',
+        lang = 'lua',
+        start_line = 1,
+        lines = { ' local x = 1', '+local y = 2' },
+        file_new_start = 1,
+        file_new_count = 2,
+        repo_root = '/nonexistent/path',
+      }
+
+      assert.has_no.errors(function()
+        highlight.highlight_hunk(
+          bufnr,
+          ns,
+          hunk,
+          default_opts({ highlights = { context = { enabled = true, lines = 25 } } })
+        )
+      end)
+
+      local extmarks = get_extmarks(bufnr)
+      assert.is_true(#extmarks > 0)
+      delete_buffer(bufnr)
+    end)
   end)
 
   describe('diff header highlighting', function()
@@ -1086,6 +1195,7 @@ describe('highlight', function()
         highlights = {
           background = false,
           gutter = false,
+          context = { enabled = false, lines = 0 },
           treesitter = { enabled = true, max_lines = 500 },
           vim = { enabled = false, max_lines = 200 },
         },
@@ -1242,6 +1352,7 @@ describe('highlight', function()
         highlights = {
           background = false,
           gutter = false,
+          context = { enabled = false, lines = 0 },
           treesitter = { enabled = true, max_lines = 500 },
           vim = { enabled = false, max_lines = 200 },
         },

spec/parser_spec.lua

@@ -421,5 +421,84 @@ describe('parser', function()
       os.remove(file_path)
       vim.fn.delete(repo_root, 'rf')
     end)
+
+    it('extracts file line numbers from @@ header', function()
+      local bufnr = create_buffer({
+        'M lua/test.lua',
+        '@@ -1,3 +1,4 @@',
+        ' local M = {}',
+        '+local new = true',
+        ' return M',
+      })
+      local hunks = parser.parse_buffer(bufnr)
+
+      assert.are.equal(1, #hunks)
+      assert.are.equal(1, hunks[1].file_old_start)
+      assert.are.equal(3, hunks[1].file_old_count)
+      assert.are.equal(1, hunks[1].file_new_start)
+      assert.are.equal(4, hunks[1].file_new_count)
+      delete_buffer(bufnr)
+    end)
+
+    it('extracts large line numbers from @@ header', function()
+      local bufnr = create_buffer({
+        'M lua/test.lua',
+        '@@ -100,20 +200,30 @@',
+        ' local M = {}',
+      })
+      local hunks = parser.parse_buffer(bufnr)
+
+      assert.are.equal(1, #hunks)
+      assert.are.equal(100, hunks[1].file_old_start)
+      assert.are.equal(20, hunks[1].file_old_count)
+      assert.are.equal(200, hunks[1].file_new_start)
+      assert.are.equal(30, hunks[1].file_new_count)
+      delete_buffer(bufnr)
+    end)
+
+    it('defaults count to 1 when omitted in @@ header', function()
+      local bufnr = create_buffer({
+        'M lua/test.lua',
+        '@@ -1 +1 @@',
+        ' local M = {}',
+      })
+      local hunks = parser.parse_buffer(bufnr)
+
+      assert.are.equal(1, #hunks)
+      assert.are.equal(1, hunks[1].file_old_start)
+      assert.are.equal(1, hunks[1].file_old_count)
+      assert.are.equal(1, hunks[1].file_new_start)
+      assert.are.equal(1, hunks[1].file_new_count)
+      delete_buffer(bufnr)
+    end)
+
+    it('stores repo_root on hunk when available', function()
+      local bufnr = create_buffer({
+        'M lua/test.lua',
+        '@@ -1,3 +1,4 @@',
+        ' local M = {}',
+        '+local new = true',
+        ' return M',
+      })
+      vim.api.nvim_buf_set_var(bufnr, 'diffs_repo_root', '/tmp/test-repo')
+      local hunks = parser.parse_buffer(bufnr)
+
+      assert.are.equal(1, #hunks)
+      assert.are.equal('/tmp/test-repo', hunks[1].repo_root)
+      delete_buffer(bufnr)
+    end)
+
+    it('repo_root is nil when not available', function()
+      local bufnr = create_buffer({
+        'M lua/test.lua',
+        '@@ -1,3 +1,4 @@',
+        ' local M = {}',
+      })
+      local hunks = parser.parse_buffer(bufnr)
+
+      assert.are.equal(1, #hunks)
+      assert.is_nil(hunks[1].repo_root)
+      delete_buffer(bufnr)
+    end)
   end)
 end)