Merge pull request #88 from barrettruth/fix/treesitter-injections

fix(highlight): include treesitter injections

README.md

@@ -74,4 +74,4 @@ luarocks install diffs.nvim
 - [`gitsigns.nvim`](https://github.com/lewis6991/gitsigns.nvim)
 - [`git-conflict.nvim`](https://github.com/akinsho/git-conflict.nvim)
 - [@phanen](https://github.com/phanen) - diff header highlighting, unknown
-  filetype fix, shebang/modeline detection
+  filetype fix, shebang/modeline detection, treesitter injection support

doc/diffs.nvim.txt

@@ -472,7 +472,8 @@ ACKNOWLEDGEMENTS                                     *diffs-acknowledgements*
 - vim-fugitive (https://github.com/tpope/vim-fugitive)
 - codediff.nvim (https://github.com/esmuellert/codediff.nvim)
 - diffview.nvim (https://github.com/sindrets/diffview.nvim)
-- @phanen (https://github.com/phanen) - diff header highlighting
+- @phanen (https://github.com/phanen) - diff header highlighting,
+  treesitter injection support
 
 ==============================================================================
  vim:tw=78:ts=8:ft=help:norl:

lua/diffs/highlight.lua

@@ -116,44 +116,50 @@ local function highlight_treesitter(
     return 0
   end
 
-  local trees = parser_obj:parse()
+  local trees = parser_obj:parse(true)
   if not trees or #trees == 0 then
     dbg('parse returned no trees for lang: %s', lang)
     return 0
   end
 
-  local query = vim.treesitter.query.get(lang, 'highlights')
-  if not query then
-    dbg('no highlights query for lang: %s', lang)
-    return 0
-  end
-
   local extmark_count = 0
-  for id, node, metadata in query:iter_captures(trees[1]:root(), code) do
-    local capture_name = '@' .. query.captures[id] .. '.' .. lang
-    local sr, sc, er, ec = node:range()
+  parser_obj:for_each_tree(function(tree, ltree)
+    local tree_lang = ltree:lang()
+    local query = vim.treesitter.query.get(tree_lang, 'highlights')
+    if not query then
+      return
+    end
 
-    local buf_sr = line_map[sr]
-    if buf_sr then
-      local buf_er = line_map[er] or buf_sr
+    for id, node, metadata in query:iter_captures(tree:root(), code) do
+      local capture = query.captures[id]
+      if capture ~= 'spell' and capture ~= 'nospell' then
+        local capture_name = '@' .. capture .. '.' .. tree_lang
+        local sr, sc, er, ec = node:range()
 
-      local buf_sc = sc + col_offset
-      local buf_ec = ec + col_offset
+        local buf_sr = line_map[sr]
+        if buf_sr then
+          local buf_er = line_map[er] or buf_sr
 
-      local priority = lang == 'diff' and (tonumber(metadata.priority) or 100) or PRIORITY_SYNTAX
+          local buf_sc = sc + col_offset
+          local buf_ec = ec + col_offset
 
-      pcall(vim.api.nvim_buf_set_extmark, bufnr, ns, buf_sr, buf_sc, {
-        end_row = buf_er,
-        end_col = buf_ec,
-        hl_group = capture_name,
-        priority = priority,
-      })
-      extmark_count = extmark_count + 1
-      if covered_lines then
-        covered_lines[buf_sr] = true
+          local priority = tree_lang == 'diff' and (tonumber(metadata.priority) or 100)
+            or PRIORITY_SYNTAX
+
+          pcall(vim.api.nvim_buf_set_extmark, bufnr, ns, buf_sr, buf_sc, {
+            end_row = buf_er,
+            end_col = buf_ec,
+            hl_group = capture_name,
+            priority = priority,
+          })
+          extmark_count = extmark_count + 1
+          if covered_lines then
+            covered_lines[buf_sr] = true
+          end
+        end
       end
     end
-  end
+  end)
 
   return extmark_count
 end

spec/helpers.lua

@@ -12,6 +12,7 @@ local function ensure_parser(lang)
 end
 
 ensure_parser('lua')
+ensure_parser('vim')
 
 local M = {}
 

spec/highlight_spec.lua

@@ -1164,6 +1164,94 @@ describe('highlight', function()
       assert.is_true(#extmarks > 0)
       delete_buffer(bufnr)
     end)
+
+    it('highlights treesitter injections', function()
+      local bufnr = create_buffer({
+        '@@ -1,1 +1,2 @@',
+        ' local x = 1',
+        '+vim.cmd([[ echo 1 ]])',
+      })
+
+      local hunk = {
+        filename = 'test.lua',
+        lang = 'lua',
+        start_line = 1,
+        lines = { ' local x = 1', '+vim.cmd([[ echo 1 ]])' },
+      }
+
+      highlight.highlight_hunk(bufnr, ns, hunk, default_opts())
+
+      local extmarks = get_extmarks(bufnr)
+      local has_vim_capture = false
+      for _, mark in ipairs(extmarks) do
+        if mark[4] and mark[4].hl_group and mark[4].hl_group:match('^@.*%.vim$') then
+          has_vim_capture = true
+          break
+        end
+      end
+      assert.is_true(has_vim_capture)
+      delete_buffer(bufnr)
+    end)
+
+    it('includes captures from both base and injected languages', function()
+      local bufnr = create_buffer({
+        '@@ -1,1 +1,2 @@',
+        ' local x = 1',
+        '+vim.cmd([[ echo 1 ]])',
+      })
+
+      local hunk = {
+        filename = 'test.lua',
+        lang = 'lua',
+        start_line = 1,
+        lines = { ' local x = 1', '+vim.cmd([[ echo 1 ]])' },
+      }
+
+      highlight.highlight_hunk(bufnr, ns, hunk, default_opts())
+
+      local extmarks = get_extmarks(bufnr)
+      local has_lua = false
+      local has_vim = false
+      for _, mark in ipairs(extmarks) do
+        if mark[4] and mark[4].hl_group then
+          if mark[4].hl_group:match('^@.*%.lua$') then
+            has_lua = true
+          end
+          if mark[4].hl_group:match('^@.*%.vim$') then
+            has_vim = true
+          end
+        end
+      end
+      assert.is_true(has_lua)
+      assert.is_true(has_vim)
+      delete_buffer(bufnr)
+    end)
+
+    it('filters @spell and @nospell captures from injections', function()
+      local bufnr = create_buffer({
+        '@@ -1,1 +1,2 @@',
+        ' local x = 1',
+        '+vim.cmd([[ echo 1 ]])',
+      })
+
+      local hunk = {
+        filename = 'test.lua',
+        lang = 'lua',
+        start_line = 1,
+        lines = { ' local x = 1', '+vim.cmd([[ echo 1 ]])' },
+      }
+
+      highlight.highlight_hunk(bufnr, ns, hunk, default_opts())
+
+      local extmarks = get_extmarks(bufnr)
+      for _, mark in ipairs(extmarks) do
+        if mark[4] and mark[4].hl_group then
+          assert.is_falsy(mark[4].hl_group:match('@spell'))
+          assert.is_falsy(mark[4].hl_group:match('@nospell'))
+        end
+      end
+      delete_buffer(bufnr)
+    end)
   end)
 
   describe('diff header highlighting', function()