treesitter: allow custom parser for highlighter

Also allow to get parser ranges.

This will be useful for language injection, allowing us to tweak the
parser's ranges on the fly.

Update runtime/lua/vim/treesitter.lua

Co-authored-by: Paul Burlumi <paul@burlumi.com>

runtime/lua/vim/treesitter.lua

@@ -59,6 +59,24 @@ function Parser:_on_bytes(bufnr, changed_tick,
   end
 end
 
+--- Registers callbacks for the parser
+-- @param cbs An `nvim_buf_attach`-like table argument with the following keys :
+--  `on_bytes` : see `nvim_buf_attach`, but this will be called _after_ the parsers callback.
+--  `on_changedtree` : a callback that will be called everytime the tree has syntactical changes.
+--      it will only be passed one argument, that is a table of the ranges (as node ranges) that
+--      changed.
+function Parser:register_cbs(cbs)
+  if not cbs then return end
+
+  if cbs.on_changedtree then
+    table.insert(self.changedtree_cbs, cbs.on_changedtree)
+  end
+
+  if cbs.on_bytes then
+    table.insert(self.bytes_cbs, cbs.on_bytes)
+  end
+end
+
 --- Sets the included ranges for the current parser
 --
 -- @param ranges A table of nodes that will be used as the ranges the parser should include.
@@ -68,6 +86,11 @@ function Parser:set_included_ranges(ranges)
   self.valid = false
 end
 
+--- Gets the included ranges for the parsers
+function Parser:included_ranges()
+  return self._parser:included_ranges()
+end
+
 local M = vim.tbl_extend("error", query, language)
 
 setmetatable(M, {
@@ -127,11 +150,7 @@ end
 --
 -- @param bufnr The buffer the parser should be tied to
 -- @param ft The filetype of this parser
--- @param buf_attach_cbs An `nvim_buf_attach`-like table argument with the following keys :
+-- @param buf_attach_cbs See Parser:register_cbs
---  `on_lines` : see `nvim_buf_attach`, but this will be called _after_ the parsers callback.
---  `on_changedtree` : a callback that will be called everytime the tree has syntactical changes.
---      it will only be passed one argument, that is a table of the ranges (as node ranges) that
---      changed.
 --
 -- @returns The parser
 function M.get_parser(bufnr, lang, buf_attach_cbs)
@@ -147,13 +166,7 @@ function M.get_parser(bufnr, lang, buf_attach_cbs)
     parsers[id] = M._create_parser(bufnr, lang, id)
   end
 
-  if buf_attach_cbs and buf_attach_cbs.on_changedtree then
+  parsers[id]:register_cbs(buf_attach_cbs)
-    table.insert(parsers[id].changedtree_cbs, buf_attach_cbs.on_changedtree)
-  end
-
-  if buf_attach_cbs and buf_attach_cbs.on_bytes then
-    table.insert(parsers[id].bytes_cbs, buf_attach_cbs.on_bytes)
-  end
 
   return parsers[id]
 end

runtime/lua/vim/treesitter/highlighter.lua

@@ -56,21 +56,14 @@ TSHighlighter.hl_map = {
     ["include"] = "Include",
 }
 
-function TSHighlighter.new(bufnr, ft, query)
+function TSHighlighter.new(parser, query)
-  if bufnr == nil or bufnr == 0 then
-    bufnr = a.nvim_get_current_buf()
-  end
-
   local self = setmetatable({}, TSHighlighter)
-  self.parser = vim.treesitter.get_parser(
-    bufnr,
-    ft,
-    {
-      on_changedtree = function(...) self:on_changedtree(...) end,
-    }
-  )
 
-  self.buf = self.parser.bufnr
+  self.parser = parser
+  parser:register_cbs {
+    on_changedtree = function(...) self:on_changedtree(...) end
+  }
+
   self:set_query(query)
   self.edit_count = 0
   self.redraw_count = 0
@@ -79,7 +72,7 @@ function TSHighlighter.new(bufnr, ft, query)
   a.nvim_buf_set_option(self.buf, "syntax", "")
 
   -- TODO(bfredl): can has multiple highlighters per buffer????
-  TSHighlighter.active[bufnr] = self
+  TSHighlighter.active[parser.bufnr] = self
 
   -- Tricky: if syntax hasn't been enabled, we need to reload color scheme
   -- but use synload.vim rather than syntax.vim to not enable
@@ -119,13 +112,6 @@ end
 function TSHighlighter:set_query(query)
   if type(query) == "string" then
     query = vim.treesitter.parse_query(self.parser.lang, query)
-  elseif query == nil then
-    query = vim.treesitter.get_query(self.parser.lang, 'highlights')
-
-    if query == nil then
-      a.nvim_err_writeln("No highlights.scm query found for " .. self.parser.lang)
-      query = vim.treesitter.parse_query(self.parser.lang, "")
-    end
   end
 
   self.query = query
@@ -139,7 +125,7 @@ function TSHighlighter:set_query(query)
     end
   })
 
-  a.nvim__buf_redraw_range(self.buf, 0, a.nvim_buf_line_count(self.buf))
+  a.nvim__buf_redraw_range(self.parser.bufnr, 0, a.nvim_buf_line_count(self.parser.bufnr))
 end
 
 function TSHighlighter._on_line(_, _win, buf, line)

src/nvim/lua/treesitter.c

@@ -43,6 +43,7 @@ static struct luaL_Reg parser_meta[] = {
   { "edit", parser_edit },
   { "tree", parser_tree },
   { "set_included_ranges", parser_set_ranges },
+  { "included_ranges", parser_get_ranges },
   { NULL, NULL }
 };
 
@@ -314,6 +315,26 @@ static const char *input_cb(void *payload, uint32_t byte_index,
 #undef BUFSIZE
 }
 
+static void push_ranges(lua_State *L,
+                        const TSRange *ranges,
+                        const unsigned int length)
+{
+  lua_createtable(L, length, 0);
+  for (size_t i = 0; i < length; i++) {
+    lua_createtable(L, 4, 0);
+    lua_pushinteger(L, ranges[i].start_point.row);
+    lua_rawseti(L, -2, 1);
+    lua_pushinteger(L, ranges[i].start_point.column);
+    lua_rawseti(L, -2, 2);
+    lua_pushinteger(L, ranges[i].end_point.row);
+    lua_rawseti(L, -2, 3);
+    lua_pushinteger(L, ranges[i].end_point.column);
+    lua_rawseti(L, -2, 4);
+
+    lua_rawseti(L, -2, i+1);
+  }
+}
+
 static int parser_parse(lua_State *L)
 {
   TSLua_parser *p = parser_check(L);
@@ -363,20 +384,8 @@ static int parser_parse(lua_State *L)
 
   tslua_push_tree(L, p->tree);
 
-  lua_createtable(L, n_ranges, 0);
+  push_ranges(L, changed, n_ranges);
-  for (size_t i = 0; i < n_ranges; i++) {
-    lua_createtable(L, 4, 0);
-    lua_pushinteger(L, changed[i].start_point.row);
-    lua_rawseti(L, -2, 1);
-    lua_pushinteger(L, changed[i].start_point.column);
-    lua_rawseti(L, -2, 2);
-    lua_pushinteger(L, changed[i].end_point.row);
-    lua_rawseti(L, -2, 3);
-    lua_pushinteger(L, changed[i].end_point.column);
-    lua_rawseti(L, -2, 4);
 
-    lua_rawseti(L, -2, i+1);
-  }
   xfree(changed);
   return 2;
 }
@@ -474,6 +483,21 @@ static int parser_set_ranges(lua_State *L)
   return 0;
 }
 
+static int parser_get_ranges(lua_State *L)
+{
+  TSLua_parser *p = parser_check(L);
+  if (!p || !p->parser) {
+    return 0;
+  }
+
+  unsigned int len;
+  const TSRange *ranges = ts_parser_included_ranges(p->parser, &len);
+
+  push_ranges(L, ranges, len);
+
+  return 1;
+}
+
 
 // Tree methods
 

test/functional/lua/treesitter_spec.lua

@@ -429,9 +429,10 @@ static int nlua_schedule(lua_State *const lstate)
     ]]}
 
     exec_lua([[
+      local parser = vim.treesitter.get_parser(0, "c")
       local highlighter = vim.treesitter.highlighter
       local query = ...
-      test_hl = highlighter.new(0, "c", query)
+      test_hl = highlighter.new(parser, query)
     ]], hl_query)
     screen:expect{grid=[[
       {2:/// Schedule Lua callback on main loop's event queue}             |
@@ -568,6 +569,72 @@ static int nlua_schedule(lua_State *const lstate)
     ]]}
   end)
 
+  it("supports highlighting with custom parser", function()
+    if not check_parser() then return end
+
+    local screen = Screen.new(65, 18)
+    screen:attach()
+    screen:set_default_attr_ids({ {bold = true, foreground = Screen.colors.SeaGreen4} })
+
+    insert(test_text)
+
+    screen:expect{ grid= [[
+      int width = INT_MAX, height = INT_MAX;                         |
+      bool ext_widgets[kUIExtCount];                                 |
+      for (UIExtension i = 0; (int)i < kUIExtCount; i++) {           |
+        ext_widgets[i] = true;                                       |
+      }                                                              |
+                                                                     |
+      bool inclusive = ui_override();                                |
+      for (size_t i = 0; i < ui_count; i++) {                        |
+        UI *ui = uis[i];                                             |
+        width = MIN(ui->width, width);                               |
+        height = MIN(ui->height, height);                            |
+        foo = BAR(ui->bazaar, bazaar);                               |
+        for (UIExtension j = 0; (int)j < kUIExtCount; j++) {         |
+          ext_widgets[j] &= (ui->ui_ext[j] || inclusive);            |
+        }                                                            |
+      }                                                              |
+    ^}                                                                |
+                                                                     |
+    ]] }
+
+    exec_lua([[
+    parser = vim.treesitter.get_parser(0, "c")
+    query = vim.treesitter.parse_query("c", "(declaration) @decl")
+
+    local nodes = {}
+    for _, node in query:iter_captures(parser:parse():root(), 0, 0, 19) do
+      table.insert(nodes, node)
+    end
+
+    parser:set_included_ranges(nodes)
+
+    local hl = vim.treesitter.highlighter.new(parser, "(identifier) @type")
+    ]])
+
+    screen:expect{ grid = [[
+      int {1:width} = {1:INT_MAX}, {1:height} = {1:INT_MAX};                         |
+      bool {1:ext_widgets}[{1:kUIExtCount}];                                 |
+      for (UIExtension {1:i} = 0; (int)i < kUIExtCount; i++) {           |
+        ext_widgets[i] = true;                                       |
+      }                                                              |
+                                                                     |
+      bool {1:inclusive} = {1:ui_override}();                                |
+      for (size_t {1:i} = 0; i < ui_count; i++) {                        |
+        UI *{1:ui} = {1:uis}[{1:i}];                                             |
+        width = MIN(ui->width, width);                               |
+        height = MIN(ui->height, height);                            |
+        foo = BAR(ui->bazaar, bazaar);                               |
+        for (UIExtension {1:j} = 0; (int)j < kUIExtCount; j++) {         |
+          ext_widgets[j] &= (ui->ui_ext[j] || inclusive);            |
+        }                                                            |
+      }                                                              |
+    ^}                                                                |
+                                                                     |
+    ]] }
+  end)
+
   it('inspects language', function()
     if not check_parser() then return end
 
@@ -613,23 +680,29 @@ static int nlua_schedule(lua_State *const lstate)
 
     insert(test_text)
 
-    local res = exec_lua([[
+    local res = exec_lua [[
     parser = vim.treesitter.get_parser(0, "c")
     return { parser:parse():root():range() }
-    ]])
+    ]]
 
     eq({0, 0, 19, 0}, res)
 
     -- The following sets the included ranges for the current parser
     -- As stated here, this only includes the function (thus the whole buffer, without the last line)
-    local res2 = exec_lua([[
+    local res2 = exec_lua [[
     local root = parser:parse():root()
     parser:set_included_ranges({root:child(0)})
     parser.valid = false
     return { parser:parse():root():range() }
-    ]])
+    ]]
 
     eq({0, 0, 18, 1}, res2)
+
+    local range = exec_lua [[
+      return parser:included_ranges()
+    ]]
+
+    eq(range, { { 0, 0, 18, 1 } })
   end)
   it("allows to set complex ranges", function()
     if not check_parser() then return end
@@ -637,7 +710,7 @@ static int nlua_schedule(lua_State *const lstate)
     insert(test_text)
 
 
-    local res = exec_lua([[
+    local res = exec_lua [[
     parser = vim.treesitter.get_parser(0, "c")
     query = vim.treesitter.parse_query("c", "(declaration) @decl")
 
@@ -655,7 +728,7 @@ static int nlua_schedule(lua_State *const lstate)
       table.insert(res, { root:named_child(i):range() })
     end
     return res
-    ]])
+    ]]
 
     eq({
       { 2, 2, 2, 40 },