feat(treesitter): expand the API

4 files changed, 160 insertions, 19 deletions

diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt
index 23bb6d4343..c5261b739d 100644
--- a/runtime/doc/news.txt
+++ b/runtime/doc/news.txt
@@ -171,6 +171,14 @@ The following new APIs or features were added.
 
 • |vim.treesitter.foldexpr()| can be used for 'foldexpr' to use treesitter for folding.
 
+• Expanded the TSNode API with:
+  - |TSNode:tree()|
+  - |TSNode:has_changes()|
+  - |TSNode:extra()|
+  - |TSNode:equal()|
+
+  Additionally |TSNode:range()| now takes an optional {include_bytes} argument.
+
 ==============================================================================
 CHANGED FEATURES                                                 *news-changes*
 
diff --git a/runtime/doc/treesitter.txt b/runtime/doc/treesitter.txt
index 76cde91920..fdea84282f 100644
--- a/runtime/doc/treesitter.txt
+++ b/runtime/doc/treesitter.txt
@@ -136,9 +136,16 @@ TSNode:end_()                                           *TSNode:end_()*
     Get the node's end position. Return three values: the row, column and
     total byte count (all zero-based).
 
-TSNode:range()                                          *TSNode:range()*
-    Get the range of the node. Return four values: the row, column of the
-    start position, then the row, column of the end position.
+TSNode:range({include_bytes})                           *TSNode:range()*
+    Get the range of the node.
+
+    Return four or six values:
+        - start row
+        - start column
+        - start byte (if {include_bytes} is `true`)
+        - end row
+        - end column
+        - end byte (if {include_bytes} is `true`)
 
 TSNode:type()                                           *TSNode:type()*
     Get the node's type as a string.
@@ -155,6 +162,13 @@ TSNode:missing()                                        *TSNode:missing()*
     Check if the node is missing. Missing nodes are inserted by the parser in
     order to recover from certain kinds of syntax errors.
 
+TSNode:extra()                                          *TSNode:extra()*
+    Check if the node is extra. Extra nodes represent things like comments,
+    which are not required by the grammar but can appear anywhere.
+
+TSNode:has_changes()                                    *TSNode:has_changes()*
+    Check if a syntax node has been edited.
+
 TSNode:has_error()                                      *TSNode:has_error()*
     Check if the node is a syntax error or contains any syntax errors.
 
@@ -171,6 +185,8 @@ TSNode:id()                                             *TSNode:id()*
     Note: The `id` is not guaranteed to be unique for nodes from different
     trees.
 
+TSNode:tree()                                           *TSNode:tree()*
+    Get the |TSTree| of the node.
                                                 *TSNode:descendant_for_range()*
 TSNode:descendant_for_range({start_row}, {start_col}, {end_row}, {end_col})
     Get the smallest node within this node that spans the given range of (row,
@@ -180,6 +196,9 @@ TSNode:descendant_for_range({start_row}, {start_col}, {end_row}, {end_col})
 TSNode:named_descendant_for_range({start_row}, {start_col}, {end_row}, {end_col})
     Get the smallest named node within this node that spans the given range of
     (row, column) positions
+                                                        *TSNode:equal()*
+TSNode:equal({node})
+    Check if {node} refers to the same node within the same tree.
 
 ==============================================================================
 TREESITTER QUERIES                                          *treesitter-query*
diff --git a/runtime/lua/vim/treesitter/_meta.lua b/runtime/lua/vim/treesitter/_meta.lua
index 87b4560798..731a5ebf9f 100644
--- a/runtime/lua/vim/treesitter/_meta.lua
+++ b/runtime/lua/vim/treesitter/_meta.lua
@@ -2,6 +2,7 @@
 
 ---@class TSNode
 ---@field id fun(self: TSNode): integer
+---@field tree fun(self: TSNode): TSTree
 ---@field range fun(self: TSNode): integer, integer, integer, integer
 ---@field start fun(self: TSNode): integer, integer, integer
 ---@field end_ fun(self: TSNode): integer, integer, integer
@@ -9,6 +10,7 @@
 ---@field symbol fun(self: TSNode): integer
 ---@field named fun(self: TSNode): boolean
 ---@field missing fun(self: TSNode): boolean
+---@field extra fun(self: TSNode): boolean
 ---@field child_count fun(self: TSNode): integer
 ---@field named_child_count fun(self: TSNode): integer
 ---@field child fun(self: TSNode, integer): TSNode
@@ -21,7 +23,8 @@
 ---@field next_named_sibling fun(self: TSNode): TSNode
 ---@field prev_named_sibling fun(self: TSNode): TSNode
 ---@field named_children fun(self: TSNode): TSNode[]
----@field has_error fun(self: TSNode): boolean
+---@field has_changes fun(self: TSNode): boolean
+---@field equal fun(self: TSNode, other: TSNode): boolean
 ---@field iter_children fun(self: TSNode): fun(): TSNode, string
 local TSNode = {}
 
@@ -41,8 +44,11 @@ function TSNode:_rawquery(query, captures, start, end_) end
 
 ---@class TSParser
 ---@field parse fun(self: TSParser, tree, source: integer|string): TSTree, integer[]
+---@field reset fun(self: TSParser)
 ---@field included_ranges fun(self: TSParser): integer[]
 ---@field set_included_ranges fun(self: TSParser, ranges: integer[][])
+---@field set_timeout fun(self: TSParser, timeout: integer)
+---@field timeout fun(self: TSParser): integer
 
 ---@class TSTree
 ---@field root fun(self: TSTree): TSNode
diff --git a/src/nvim/lua/treesitter.c b/src/nvim/lua/treesitter.c
index 5248ebed14..d72e7a1825 100644
--- a/src/nvim/lua/treesitter.c
+++ b/src/nvim/lua/treesitter.c
@@ -51,8 +51,11 @@ static struct luaL_Reg parser_meta[] = {
   { "__gc", parser_gc },
   { "__tostring", parser_tostring },
   { "parse", parser_parse },
+  { "reset", parser_reset },
   { "set_included_ranges", parser_set_ranges },
   { "included_ranges", parser_get_ranges },
+  { "set_timeout", parser_set_timeout },
+  { "timeout", parser_get_timeout },
   { NULL, NULL }
 };
 
@@ -78,6 +81,8 @@ static struct luaL_Reg node_meta[] = {
   { "field", node_field },
   { "named", node_named },
   { "missing", node_missing },
+  { "extra", node_extra },
+  { "has_changes", node_has_changes },
   { "has_error", node_has_error },
   { "sexpr", node_sexpr },
   { "child_count", node_child_count },
@@ -95,7 +100,9 @@ static struct luaL_Reg node_meta[] = {
   { "prev_named_sibling", node_prev_named_sibling },
   { "named_children", node_named_children },
   { "root", node_root },
+  { "tree", node_tree },
   { "byte_length", node_byte_length },
+  { "equal", node_equal },
 
   { NULL, NULL }
 };
@@ -446,6 +453,16 @@ static int parser_parse(lua_State *L)
   return 2;
 }
 
+static int parser_reset(lua_State *L)
+{
+  TSParser **p = parser_check(L, 1);
+  if (p && *p) {
+    ts_parser_reset(*p);
+  }
+
+  return 0;
+}
+
 static int tree_copy(lua_State *L)
 {
   TSTree **tree = tree_check(L, 1);
@@ -597,6 +614,33 @@ static int parser_get_ranges(lua_State *L)
   return 1;
 }
 
+static int parser_set_timeout(lua_State *L)
+{
+  TSParser **p = parser_check(L, 1);
+  if (!p) {
+    return 0;
+  }
+
+  if (lua_gettop(L) < 2) {
+    luaL_error(L, "integer expected");
+  }
+
+  uint32_t timeout = (uint32_t)luaL_checkinteger(L, 2);
+  ts_parser_set_timeout_micros(*p, timeout);
+  return 0;
+}
+
+static int parser_get_timeout(lua_State *L)
+{
+  TSParser **p = parser_check(L, 1);
+  if (!p) {
+    return 0;
+  }
+
+  lua_pushinteger(L, (long)ts_parser_timeout_micros(*p));
+  return 0;
+}
+
 // Tree methods
 
 /// push tree interface on lua stack.
@@ -740,12 +784,29 @@ static int node_range(lua_State *L)
   if (!node_check(L, 1, &node)) {
     return 0;
   }
+
+  bool include_bytes = false;
+  if (lua_gettop(L) >= 2) {
+    include_bytes = lua_toboolean(L, 2);
+  }
+
   TSPoint start = ts_node_start_point(node);
   TSPoint end = ts_node_end_point(node);
-  lua_pushnumber(L, start.row);
-  lua_pushnumber(L, start.column);
-  lua_pushnumber(L, end.row);
-  lua_pushnumber(L, end.column);
+
+  if (include_bytes) {
+    lua_pushinteger(L, start.row);
+    lua_pushinteger(L, start.column);
+    lua_pushinteger(L, ts_node_start_byte(node));
+    lua_pushinteger(L, end.row);
+    lua_pushinteger(L, end.column);
+    lua_pushinteger(L, ts_node_end_byte(node));
+    return 6;
+  }
+
+  lua_pushinteger(L, start.row);
+  lua_pushinteger(L, start.column);
+  lua_pushinteger(L, end.row);
+  lua_pushinteger(L, end.column);
   return 4;
 }
 
@@ -757,9 +818,9 @@ static int node_start(lua_State *L)
   }
   TSPoint start = ts_node_start_point(node);
   uint32_t start_byte = ts_node_start_byte(node);
-  lua_pushnumber(L, start.row);
-  lua_pushnumber(L, start.column);
-  lua_pushnumber(L, start_byte);
+  lua_pushinteger(L, start.row);
+  lua_pushinteger(L, start.column);
+  lua_pushinteger(L, start_byte);
   return 3;
 }
 
@@ -771,9 +832,9 @@ static int node_end(lua_State *L)
   }
   TSPoint end = ts_node_end_point(node);
   uint32_t end_byte = ts_node_end_byte(node);
-  lua_pushnumber(L, end.row);
-  lua_pushnumber(L, end.column);
-  lua_pushnumber(L, end_byte);
+  lua_pushinteger(L, end.row);
+  lua_pushinteger(L, end.column);
+  lua_pushinteger(L, end_byte);
   return 3;
 }
 
@@ -784,7 +845,7 @@ static int node_child_count(lua_State *L)
     return 0;
   }
   uint32_t count = ts_node_child_count(node);
-  lua_pushnumber(L, count);
+  lua_pushinteger(L, count);
   return 1;
 }
 
@@ -795,7 +856,7 @@ static int node_named_child_count(lua_State *L)
     return 0;
   }
   uint32_t count = ts_node_named_child_count(node);
-  lua_pushnumber(L, count);
+  lua_pushinteger(L, count);
   return 1;
 }
 
@@ -816,7 +877,7 @@ static int node_symbol(lua_State *L)
     return 0;
   }
   TSSymbol symbol = ts_node_symbol(node);
-  lua_pushnumber(L, symbol);
+  lua_pushinteger(L, symbol);
   return 1;
 }
 
@@ -882,6 +943,26 @@ static int node_missing(lua_State *L)
   return 1;
 }
 
+static int node_extra(lua_State *L)
+{
+  TSNode node;
+  if (!node_check(L, 1, &node)) {
+    return 0;
+  }
+  lua_pushboolean(L, ts_node_is_extra(node));
+  return 1;
+}
+
+static int node_has_changes(lua_State *L)
+{
+  TSNode node;
+  if (!node_check(L, 1, &node)) {
+    return 0;
+  }
+  lua_pushboolean(L, ts_node_has_changes(node));
+  return 1;
+}
+
 static int node_has_error(lua_State *L)
 {
   TSNode node;
@@ -1108,6 +1189,17 @@ static int node_root(lua_State *L)
   return 1;
 }
 
+static int node_tree(lua_State *L)
+{
+  TSNode node;
+  if (!node_check(L, 1, &node)) {
+    return 0;
+  }
+
+  push_tree(L, (TSTree *)node.tree, false);
+  return 1;
+}
+
 static int node_byte_length(lua_State *L)
 {
   TSNode node;
@@ -1118,7 +1210,23 @@ static int node_byte_length(lua_State *L)
   uint32_t start_byte = ts_node_start_byte(node);
   uint32_t end_byte = ts_node_end_byte(node);
 
-  lua_pushnumber(L, end_byte - start_byte);
+  lua_pushinteger(L, end_byte - start_byte);
+  return 1;
+}
+
+static int node_equal(lua_State *L)
+{
+  TSNode node1;
+  if (!node_check(L, 1, &node1)) {
+    return 0;
+  }
+
+  TSNode node2;
+  if (!node_check(L, 2, &node2)) {
+    return luaL_error(L, "TSNode expected");
+  }
+
+  lua_pushboolean(L, ts_node_eq(node1, node2));
   return 1;
 }
 
@@ -1367,7 +1475,7 @@ static int query_inspect(lua_State *L)
                                                        &strlen);
         lua_pushlstring(L, str, strlen);  // [retval, patterns, pat, pred, item]
       } else if (step[k].type == TSQueryPredicateStepTypeCapture) {
-        lua_pushnumber(L, step[k].value_id + 1);  // [..., pat, pred, item]
+        lua_pushinteger(L, step[k].value_id + 1);  // [..., pat, pred, item]
       } else {
         abort();
       }