Merge pull request #12590 from nvim-treesitter/ts-fix-highlight

[RDY] Treesitter: fix highlight, attempt 2

runtime/lua/vim/treesitter.lua

@@ -15,19 +15,27 @@ function Parser:parse()
   local changes
   self.tree, changes = self._parser:parse_buf(self.bufnr)
   self.valid = true
-  for _, cb in ipairs(self.change_cbs) do
-    cb(changes)
+
+  if not vim.tbl_isempty(changes) then
+    for _, cb in ipairs(self.changedtree_cbs) do
+      cb(changes)
+    end
   end
+
   return self.tree, changes
 end
 
-function Parser:_on_lines(bufnr, _, start_row, old_stop_row, stop_row, old_byte_size)
+function Parser:_on_lines(bufnr, changed_tick, start_row, old_stop_row, stop_row, old_byte_size)
   local start_byte = a.nvim_buf_get_offset(bufnr,start_row)
   local stop_byte = a.nvim_buf_get_offset(bufnr,stop_row)
   local old_stop_byte = start_byte + old_byte_size
   self._parser:edit(start_byte,old_stop_byte,stop_byte,
                     start_row,0,old_stop_row,0,stop_row,0)
   self.valid = false
+
+  for _, cb in ipairs(self.lines_cbs) do
+    cb(bufnr, changed_tick, start_row, old_stop_row, stop_row, old_byte_size)
+  end
 end
 
 function Parser:set_included_ranges(ranges)
@@ -80,7 +88,8 @@ function M.create_parser(bufnr, lang, id)
 
   local self = setmetatable({bufnr=bufnr, lang=lang, valid=false}, Parser)
   self._parser = vim._create_ts_parser(lang)
-  self.change_cbs = {}
+  self.changedtree_cbs = {}
+  self.lines_cbs = {}
   self:parse()
     -- TODO(bfredl): use weakref to self, so that the parser is free'd is no plugin is
     -- using it.
@@ -99,7 +108,7 @@ function M.create_parser(bufnr, lang, id)
   return self
 end
 
-function M.get_parser(bufnr, ft, cb)
+function M.get_parser(bufnr, ft, buf_attach_cbs)
   if bufnr == nil or bufnr == 0 then
     bufnr = a.nvim_get_current_buf()
   end
@@ -111,9 +120,15 @@ function M.get_parser(bufnr, ft, cb)
   if parsers[id] == nil then
     parsers[id] = M.create_parser(bufnr, ft, id)
   end
-  if cb ~= nil then
-    table.insert(parsers[id].change_cbs, cb)
+
+  if buf_attach_cbs and buf_attach_cbs.on_changedtree then
+    table.insert(parsers[id].changedtree_cbs, buf_attach_cbs.on_changedtree)
   end
+
+  if buf_attach_cbs and buf_attach_cbs.on_lines then
+    table.insert(parsers[id].lines_cbs, buf_attach_cbs.on_lines)
+  end
+
   return parsers[id]
 end
 

runtime/lua/vim/tshighlighter.lua

@@ -3,6 +3,7 @@ local a = vim.api
 -- support reload for quick experimentation
 local TSHighlighter = rawget(vim.treesitter, 'TSHighlighter') or {}
 TSHighlighter.__index = TSHighlighter
+local ts_hs_ns = a.nvim_create_namespace("treesitter_hl")
 
 -- These are conventions defined by tree-sitter, though it
 -- needs to be user extensible also.
@@ -24,9 +25,17 @@ TSHighlighter.hl_map = {
 
 function TSHighlighter.new(query, bufnr, ft)
   local self = setmetatable({}, TSHighlighter)
-  self.parser = vim.treesitter.get_parser(bufnr, ft, function(...) self:on_change(...) end)
+  self.parser = vim.treesitter.get_parser(
+    bufnr,
+    ft,
+    {
+      on_changedtree = function(...) self:on_changedtree(...) end,
+      on_lines = function() self.root = self.parser:parse():root() end
+    }
+  )
+
   self.buf = self.parser.bufnr
-  -- TODO(bfredl): perhaps on_start should be called uncondionally, instead for only on mod?
+
   local tree = self.parser:parse()
   self.root = tree:root()
   self:set_query(query)
@@ -34,11 +43,6 @@ function TSHighlighter.new(query, bufnr, ft)
   self.redraw_count = 0
   self.line_count = {}
   a.nvim_buf_set_option(self.buf, "syntax", "")
-  a.nvim__buf_set_luahl(self.buf, {
-    on_start=function(...) return self:on_start(...) end,
-    on_window=function(...) return self:on_window(...) end,
-    on_line=function(...) return self:on_line(...) end,
-  })
 
   -- Tricky: if syntax hasn't been enabled, we need to reload color scheme
   -- but use synload.vim rather than syntax.vim to not enable
@@ -50,73 +54,63 @@ function TSHighlighter.new(query, bufnr, ft)
   return self
 end
 
+local function is_highlight_name(capture_name)
+  local firstc = string.sub(capture_name, 1, 1)
+  return firstc ~= string.lower(firstc)
+end
+
+function TSHighlighter: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]
+  else
+    -- Default to false to avoid recomputing
+    return TSHighlighter.hl_map[name]
+  end
+end
+
 function TSHighlighter:set_query(query)
   if type(query) == "string" then
     query = vim.treesitter.parse_query(self.parser.lang, query)
   end
   self.query = query
 
-  self.id_map = {}
-  for i, capture in ipairs(self.query.captures) do
-    local hl = 0
-    local firstc = string.sub(capture, 1, 1)
-    local hl_group = self.hl_map[capture]
-    if firstc ~= string.lower(firstc) then
-      hl_group = vim.split(capture, '.', true)[1]
-    end
-    if hl_group then
-      hl = a.nvim_get_hl_id_by_name(hl_group)
-    end
-    self.id_map[i] = hl
-  end
+  self.hl_cache = setmetatable({}, {
+    __index = function(table, capture)
+      local hl = self:get_hl_from_capture(capture)
+      rawset(table, capture, hl)
 
-    a.nvim__buf_redraw_range(self.buf, 0, a.nvim_buf_line_count(self.buf))
+      return hl
+    end
+  })
+
+  self:on_changedtree({{self.root:range()}})
 end
 
-function TSHighlighter:on_change(changes)
+function TSHighlighter:on_changedtree(changes)
+  -- Get a fresh root
+  self.root = self.parser.tree:root()
+
   for _, ch in ipairs(changes or {}) do
-    a.nvim__buf_redraw_range(self.buf, ch[1], ch[3]+1)
-  end
-  self.edit_count = self.edit_count + 1
-end
+    -- Try to be as exact as possible
+    local changed_node = self.root:descendant_for_range(ch[1], ch[2], ch[3], ch[4])
 
-function TSHighlighter:on_start(_, _buf, _tick)
-  local tree = self.parser:parse()
-  self.root = tree:root()
-end
+    a.nvim_buf_clear_namespace(self.buf, ts_hs_ns, ch[1], ch[3])
 
-function TSHighlighter:on_window(_, _win, _buf, _topline, botline)
-  self.iter = nil
-  self.nextrow = 0
-  self.botline = botline
-  self.redraw_count = self.redraw_count + 1
-end
-
-function TSHighlighter:on_line(_, _win, buf, line)
-  if self.iter == nil then
-    self.iter = self.query:iter_captures(self.root,buf,line,self.botline)
-  end
-  while line >= self.nextrow do
-    local capture, node, match = self.iter()
-    local active = true
-    if capture == nil then
-      break
-    end
-    if match ~= nil then
-      active = self:run_pred(match)
-      match.active = active
-    end
-    local start_row, start_col, end_row, end_col = node:range()
-    local hl = self.id_map[capture]
-    if hl > 0 and active and end_row >= line then
-      a.nvim__put_attr(hl, start_row, start_col, end_row, end_col)
-    end
-    if start_row > line then
-      self.nextrow = start_row
+    for capture, node in self.query:iter_captures(changed_node, self.buf, ch[1], ch[3] + 1) do
+      local start_row, start_col, end_row, end_col = node:range()
+      local hl = self.hl_cache[capture]
+      if hl then
+        a.nvim__buf_add_decoration(self.buf, ts_hs_ns, hl,
+          start_row, start_col,
+          end_row, end_col,
+          {})
+      end
     end
   end
-  self.line_count[line] = (self.line_count[line] or 0) + 1
-  --return tostring(self.line_count[line])
 end
 
 return TSHighlighter

test/functional/lua/treesitter_spec.lua

@@ -237,9 +237,7 @@ static int nlua_schedule(lua_State *const lstate)
 (number_literal) @number
 (char_literal) @string
 
-; TODO(bfredl): overlapping matches are unreliable,
-; we need a proper priority mechanism
-;(type_identifier) @type
+(type_identifier) @type
 ((type_identifier) @Special (#eq? @Special "LuaRef"))
 
 (primitive_type) @type
@@ -264,7 +262,7 @@ static int nlua_schedule(lua_State *const lstate)
       [4] = {bold = true, foreground = Screen.colors.Brown},
       [5] = {foreground = Screen.colors.Magenta},
       [6] = {foreground = Screen.colors.Red},
-      [7] = {foreground = Screen.colors.SlateBlue},
+      [7] = {bold = true, foreground = Screen.colors.SlateBlue},
       [8] = {foreground = Screen.colors.Grey100, background = Screen.colors.Red},
       [9] = {foreground = Screen.colors.Magenta, background = Screen.colors.Red},
       [10] = {foreground = Screen.colors.Red, background = Screen.colors.Red},
@@ -300,7 +298,7 @@ static int nlua_schedule(lua_State *const lstate)
     ]], hl_query)
     screen:expect{grid=[[
       {2:/// Schedule Lua callback on main loop's event queue}             |
-      {3:static} {3:int} nlua_schedule(lua_State *{3:const} lstate)                |
+      {3:static} {3:int} nlua_schedule({3:lua_State} *{3:const} lstate)                |
       {                                                                |
         {4:if} ({11:lua_type}(lstate, {5:1}) != {5:LUA_TFUNCTION}                       |
             || {6:lstate} != {6:lstate}) {                                     |
@@ -311,7 +309,7 @@ static int nlua_schedule(lua_State *const lstate)
         {7:LuaRef} cb = nlua_ref(lstate, {5:1});                               |
                                                                        |
         multiqueue_put(main_loop.events, nlua_schedule_event,          |
-                       {5:1}, ({3:void} *)(ptrdiff_t)cb);                      |
+                       {5:1}, ({3:void} *)({3:ptrdiff_t})cb);                      |
         {4:return} {5:0};                                                      |
       ^}                                                                |
       {1:~                                                                }|
@@ -322,7 +320,7 @@ static int nlua_schedule(lua_State *const lstate)
     feed('7Go*/<esc>')
     screen:expect{grid=[[
       {2:/// Schedule Lua callback on main loop's event queue}             |
-      {3:static} {3:int} nlua_schedule(lua_State *{3:const} lstate)                |
+      {3:static} {3:int} nlua_schedule({3:lua_State} *{3:const} lstate)                |
       {                                                                |
         {4:if} ({11:lua_type}(lstate, {5:1}) != {5:LUA_TFUNCTION}                       |
             || {6:lstate} != {6:lstate}) {                                     |
@@ -334,7 +332,7 @@ static int nlua_schedule(lua_State *const lstate)
         {7:LuaRef} cb = nlua_ref(lstate, {5:1});                               |
                                                                        |
         multiqueue_put(main_loop.events, nlua_schedule_event,          |
-                       {5:1}, ({3:void} *)(ptrdiff_t)cb);                      |
+                       {5:1}, ({3:void} *)({3:ptrdiff_t})cb);                      |
         {4:return} {5:0};                                                      |
       }                                                                |
       {1:~                                                                }|
@@ -344,7 +342,7 @@ static int nlua_schedule(lua_State *const lstate)
     feed('3Go/*<esc>')
     screen:expect{grid=[[
       {2:/// Schedule Lua callback on main loop's event queue}             |
-      {3:static} {3:int} nlua_schedule(lua_State *{3:const} lstate)                |
+      {3:static} {3:int} nlua_schedule({3:lua_State} *{3:const} lstate)                |
       {                                                                |
       {2:/^*}                                                               |
       {2:  if (lua_type(lstate, 1) != LUA_TFUNCTION}                       |
@@ -357,7 +355,7 @@ static int nlua_schedule(lua_State *const lstate)
         {7:LuaRef} cb = nlua_ref(lstate, {5:1});                               |
                                                                        |
         multiqueue_put(main_loop.events, nlua_schedule_event,          |
-                       {5:1}, ({3:void} *)(ptrdiff_t)cb);                      |
+                       {5:1}, ({3:void} *)({3:ptrdiff_t})cb);                      |
         {4:return} {5:0};                                                      |
       {8:}}                                                                |
                                                                        |