fix(treesitter): allow ranges in set_included_ranges

3 files changed, 97 insertions, 17 deletions

diff --git a/runtime/lua/vim/treesitter/languagetree.lua b/runtime/lua/vim/treesitter/languagetree.lua
index ed07e73a55..70e2ac4c62 100644
--- a/runtime/lua/vim/treesitter/languagetree.lua
+++ b/runtime/lua/vim/treesitter/languagetree.lua
@@ -1,3 +1,4 @@
+local a = vim.api
 local query = require'vim.treesitter.query'
 local language = require'vim.treesitter.language'
 
@@ -234,6 +235,24 @@ end
 --
 -- @param regions A list of regions this tree should manange and parse.
 function LanguageTree:set_included_regions(regions)
+  -- Transform the tables from 4 element long to 6 element long (with byte offset)
+  for _, region in ipairs(regions) do
+    for i, range in ipairs(region) do
+      if type(range) == "table" and #range == 4 then
+        -- TODO(vigoux): I don't think string parsers are useful for now
+        if type(self._source) == "number" then
+          local start_row, start_col, end_row, end_col = unpack(range)
+          -- Easy case, this is a buffer parser
+          -- TODO(vigoux): proper byte computation here, and account for EOL ?
+          local start_byte = a.nvim_buf_get_offset(self.bufnr, start_row) + start_col
+          local end_byte = a.nvim_buf_get_offset(self.bufnr, end_row) + end_col
+
+          region[i] = { start_row, start_col, start_byte, end_row, end_col, end_byte }
+        end
+      end
+    end
+  end
+
   self._regions = regions
   -- Trees are no longer valid now that we have changed regions.
   -- TODO(vigoux,steelsojka): Look into doing this smarter so we can use some of the
diff --git a/src/nvim/lua/treesitter.c b/src/nvim/lua/treesitter.c
index 19f9647fd7..a640b97d3b 100644
--- a/src/nvim/lua/treesitter.c
+++ b/src/nvim/lua/treesitter.c
@@ -180,7 +180,8 @@ int tslua_add_language(lua_State *L)
     return luaL_error(
         L,
         "ABI version mismatch : supported between %d and %d, found %d",
-        TREE_SITTER_MIN_COMPATIBLE_LANGUAGE_VERSION, TREE_SITTER_LANGUAGE_VERSION, lang_version);
+        TREE_SITTER_MIN_COMPATIBLE_LANGUAGE_VERSION,
+        TREE_SITTER_LANGUAGE_VERSION, lang_version);
   }
 
   pmap_put(cstr_t)(langs, xstrdup(lang_name), lang);
@@ -433,6 +434,72 @@ static int tree_edit(lua_State *L)
   return 0;
 }
 
+// Use the top of the stack (without popping it) to create a TSRange, it can be
+// either a lua table or a TSNode
+static void range_from_lua(lua_State *L, TSRange *range)
+{
+  TSNode node;
+
+  if (lua_istable(L, -1)) {
+    // should be a table of 6 elements
+    if (lua_objlen(L, -1) != 6) {
+      goto error;
+    }
+
+    uint32_t start_row, start_col, start_byte, end_row, end_col, end_byte;
+    lua_rawgeti(L, -1, 1);  // [ range, start_row]
+    start_row = luaL_checkinteger(L, -1);
+    lua_pop(L, 1);
+
+    lua_rawgeti(L, -1, 2);  // [ range, start_col]
+    start_col = luaL_checkinteger(L, -1);
+    lua_pop(L, 1);
+
+    lua_rawgeti(L, -1, 3);  // [ range, start_byte]
+    start_byte = luaL_checkinteger(L, -1);
+    lua_pop(L, 1);
+
+    lua_rawgeti(L, -1, 4);  // [ range, end_row]
+    end_row = luaL_checkinteger(L, -1);
+    lua_pop(L, 1);
+
+    lua_rawgeti(L, -1, 5);  // [ range, end_col]
+    end_col = luaL_checkinteger(L, -1);
+    lua_pop(L, 1);
+
+    lua_rawgeti(L, -1, 6);  // [ range, end_byte]
+    end_byte = luaL_checkinteger(L, -1);
+    lua_pop(L, 1);  // [ range ]
+
+    *range = (TSRange) {
+      .start_point = (TSPoint) {
+        .row = start_row,
+        .column = start_col
+      },
+      .end_point = (TSPoint) {
+        .row = end_row,
+        .column = end_col
+      },
+      .start_byte = start_byte,
+      .end_byte = end_byte,
+    };
+  } else if (node_check(L, -1, &node)) {
+    *range = (TSRange) {
+      .start_point = ts_node_start_point(node),
+      .end_point = ts_node_end_point(node),
+      .start_byte = ts_node_start_byte(node),
+      .end_byte = ts_node_end_byte(node)
+    };
+  } else {
+    goto error;
+  }
+  return;
+error:
+  luaL_error(
+      L,
+      "Ranges can only be made from 6 element long tables or nodes.");
+}
+
 static int parser_set_ranges(lua_State *L)
 {
   if (lua_gettop(L) < 2) {
@@ -459,22 +526,8 @@ static int parser_set_ranges(lua_State *L)
   // [ parser, ranges ]
   for (size_t index = 0; index < tbl_len; index++) {
     lua_rawgeti(L, 2, index + 1);  // [ parser, ranges, range ]
-
-    TSNode node;
-    if (!node_check(L, -1, &node)) {
-      xfree(ranges);
-      return luaL_error(
-          L,
-          "ranges should be tables of nodes.");
-    }
-    lua_pop(L, 1);  // [ parser, ranges ]
-
-    ranges[index] = (TSRange) {
-      .start_point = ts_node_start_point(node),
-      .end_point = ts_node_end_point(node),
-      .start_byte = ts_node_start_byte(node),
-      .end_byte = ts_node_end_byte(node)
-    };
+    range_from_lua(L, ranges + index);
+    lua_pop(L, 1);
   }
 
   // This memcpies ranges, thus we can free it afterwards
diff --git a/test/functional/lua/treesitter_spec.lua b/test/functional/lua/treesitter_spec.lua
index 273a5119cb..65dc1b3e03 100644
--- a/test/functional/lua/treesitter_spec.lua
+++ b/test/functional/lua/treesitter_spec.lua
@@ -781,6 +781,14 @@ local hl_query = [[
     ]]
 
     eq(range, { { 0, 0, 18, 1 } })
+
+    local range_tbl = exec_lua [[
+      parser:set_included_regions { { { 0, 0, 17, 1 } } }
+      parser:parse()
+      return parser:included_regions()
+    ]]
+
+    eq(range_tbl, { { { 0, 0, 0, 17, 1, 508 } } })
   end)
   it("allows to set complex ranges", function()
     if not check_parser() then return end