Merge pull request neovim#14245 from tjdevries/tjdevries/ts_override_hl

ts: Add per-language highlight links
mengtanhzc · Mar 31, 2021 · d55a691 · d55a691
2 parents 94c2ce2 + d50f99f
commit d55a691
Show file tree

Hide file tree

Showing 2 changed files with 77 additions and 7 deletions.
diff --git a/runtime/lua/vim/treesitter/highlighter.lua b/runtime/lua/vim/treesitter/highlighter.lua
@@ -12,6 +12,16 @@ TSHighlighterQuery.__index = TSHighlighterQuery
 
 local ns = a.nvim_create_namespace("treesitter/highlighter")
 
+local _default_highlights = {}
+local _link_default_highlight_once = function(from, to)
+  if not _default_highlights[from] then
+    _default_highlights[from] = true
+    vim.cmd(string.format("highlight default link %s %s", from, to))
+  end
+
+  return from
+end
+
 -- These are conventions defined by nvim-treesitter, though it
 -- needs to be user extensible also.
 TSHighlighter.hl_map = {
@@ -70,9 +80,12 @@ function TSHighlighterQuery.new(lang, query_string)
 
   self.hl_cache = setmetatable({}, {
     __index = function(table, capture)
-      local hl = self:get_hl_from_capture(capture)
-      rawset(table, capture, hl)
+      local hl, is_vim_highlight = self:_get_hl_from_capture(capture)
+      if not is_vim_highlight then
+        hl = _link_default_highlight_once(lang .. hl, hl)
+      end
 
+      rawset(table, capture, hl)
       return hl
     end
   })
@@ -90,16 +103,16 @@ function TSHighlighterQuery:query()
   return self._query
 end
 
-function TSHighlighterQuery:get_hl_from_capture(capture)
+--- Get the hl from capture.
+--- Returns a tuple { highlight_name: string, is_builtin: bool }
+function TSHighlighterQuery:_get_hl_from_capture(capture)
   local name = self._query.captures[capture]
 
   if is_highlight_name(name) then
     -- From "Normal.left" only keep "Normal"
-    return vim.split(name, '.', true)[1]
+    return vim.split(name, '.', true)[1], true
   else
-    -- Default to false to avoid recomputing
-    local hl = TSHighlighter.hl_map[name]
-    return hl and a.nvim_get_hl_id_by_name(hl) or 0
+    return TSHighlighter.hl_map[name] or name, false
   end
 end
 

diff --git a/test/functional/treesitter/highlight_spec.lua b/test/functional/treesitter/highlight_spec.lua
@@ -513,4 +513,61 @@ describe('treesitter highlighting', function()
                                                                        |
     ]]}
   end)
+
+  it("supports highlighting with custom highlight groups", function()
+    if pending_c_parser(pending) then return end
+
+    insert(hl_text)
+
+    exec_lua [[
+      local parser = vim.treesitter.get_parser(0, "c")
+      test_hl = vim.treesitter.highlighter.new(parser, {queries = {c = hl_query}})
+    ]]
+
+    screen:expect{grid=[[
+      {2:/// Schedule Lua callback on main loop's event queue}             |
+      {3:static} {3:int} {11:nlua_schedule}({3:lua_State} *{3:const} lstate)                |
+      {                                                                |
+        {4:if} ({11:lua_type}(lstate, {5:1}) != {5:LUA_TFUNCTION}                       |
+            || {6:lstate} != {6:lstate}) {                                     |
+          {11:lua_pushliteral}(lstate, {5:"vim.schedule: expected function"});  |
+          {4:return} {11:lua_error}(lstate);                                    |
+        }                                                              |
+                                                                       |
+        {7:LuaRef} cb = {11:nlua_ref}(lstate, {5:1});                               |
+                                                                       |
+        multiqueue_put(main_loop.events, {11:nlua_schedule_event},          |
+                       {5:1}, ({3:void} *)({3:ptrdiff_t})cb);                      |
+        {4:return} {5:0};                                                      |
+      ^}                                                                |
+      {1:~                                                                }|
+      {1:~                                                                }|
+                                                                       |
+    ]]}
+
+    -- This will change ONLY the literal strings to look like comments
+    -- The only literal string is the "vim.schedule: expected function" in this test.
+    exec_lua [[vim.cmd("highlight link cString comment")]]
+    screen:expect{grid=[[
+      {2:/// Schedule Lua callback on main loop's event queue}             |
+      {3:static} {3:int} {11:nlua_schedule}({3:lua_State} *{3:const} lstate)                |
+      {                                                                |
+        {4:if} ({11:lua_type}(lstate, {5:1}) != {5:LUA_TFUNCTION}                       |
+            || {6:lstate} != {6:lstate}) {                                     |
+          {11:lua_pushliteral}(lstate, {2:"vim.schedule: expected function"});  |
+          {4:return} {11:lua_error}(lstate);                                    |
+        }                                                              |
+                                                                       |
+        {7:LuaRef} cb = {11:nlua_ref}(lstate, {5:1});                               |
+                                                                       |
+        multiqueue_put(main_loop.events, {11:nlua_schedule_event},          |
+                       {5:1}, ({3:void} *)({3:ptrdiff_t})cb);                      |
+        {4:return} {5:0};                                                      |
+      ^}                                                                |
+      {1:~                                                                }|
+      {1:~                                                                }|
+                                                                       |
+    ]]}
+    screen:expect{ unchanged=true }
+  end)
 end)