tree-sitter: handle node equality

2 files changed, 26 insertions, 0 deletions

diff --git a/src/nvim/lua/treesitter.c b/src/nvim/lua/treesitter.c
index a71234d2c4..d2072402bb 100644
--- a/src/nvim/lua/treesitter.c
+++ b/src/nvim/lua/treesitter.c
@@ -48,6 +48,7 @@ static struct luaL_Reg tree_meta[] = {
 
 static struct luaL_Reg node_meta[] = {
   { "__tostring", node_tostring },
+  { "__eq", node_eq },
   { "__len", node_child_count },
   { "range", node_range },
   { "start", node_start },
@@ -431,6 +432,23 @@ static int node_tostring(lua_State *L)
   return 1;
 }
 
+static int node_eq(lua_State *L)
+{
+  TSNode node;
+  if (!node_check(L, &node)) {
+    return 0;
+  }
+  // This should only be called if both x and y in "x == y" has the
+  // treesitter_node metatable. So it is ok to error out otherwise.
+  TSNode *ud = luaL_checkudata(L, 2, "treesitter_node");
+  if (!ud) {
+    return 0;
+  }
+  TSNode node2 = *ud;
+  lua_pushboolean(L, ts_node_eq(node, node2));
+  return 1;
+}
+
 static int node_range(lua_State *L)
 {
   TSNode node;
diff --git a/test/functional/lua/treesitter_spec.lua b/test/functional/lua/treesitter_spec.lua
index d566f15649..8e21faca12 100644
--- a/test/functional/lua/treesitter_spec.lua
+++ b/test/functional/lua/treesitter_spec.lua
@@ -64,6 +64,13 @@ describe('tree-sitter API', function()
       eq({1,2,1,12}, exec_lua("return {descendant:range()}"))
       eq("(declaration (primitive_type) (init_declarator (identifier) (number_literal)))", exec_lua("return descendant:sexpr()"))
 
+      eq(true, exec_lua("return child == child"))
+      -- separate lua object, but represents same node
+      eq(true, exec_lua("return child == root:child(0)"))
+      eq(false, exec_lua("return child == descendant2"))
+      eq(false, exec_lua("return child == nil"))
+      eq(false, exec_lua("return child == tree"))
+
       feed("2G7|ay")
       exec_lua([[
         tree2 = parser:parse()
@@ -71,6 +78,7 @@ describe('tree-sitter API', function()
         descendant2 = root2:descendant_for_range(1,2,1,13)
       ]])
       eq(false, exec_lua("return tree2 == tree1"))
+      eq(false, exec_lua("return root2 == root"))
       eq("<node declaration>", exec_lua("return tostring(descendant2)"))
       eq({1,2,1,13}, exec_lua("return {descendant2:range()}"))