fix(treesitter): ensure TSLuaTree is always immutable

Problem:

The previous fix in #34314 relies on copying the tree in `tree_root` to
ensure the `TSNode`'s tree cannot be mutated. But that causes the
problem where two calls to `tree_root` return nodes from different
copies of a tree, which do not compare as equal. This has broken at
least one plugin.

Solution:

Make all `TSTree`s on the Lua side always immutable, avoiding the need
to copy the tree in `tree_root`, and make the only mutation point,
`tree_edit`, copy the tree instead.

(cherry picked from commit 168bf0024e)

runtime/lua/vim/treesitter/_meta/tstree.lua

@@ -26,6 +26,7 @@ function TSTree:root() end
 ---@param end_col_old integer
 ---@param end_row_new integer
 ---@param end_col_new integer
+---@return TSTree
 ---@nodoc
 function TSTree:edit(start_byte, end_byte_old, end_byte_new, start_row, start_col, end_row_old, end_col_old, end_row_new, end_col_new) end
 

runtime/lua/vim/treesitter/languagetree.lua

@@ -1082,8 +1082,8 @@ function LanguageTree:_edit(
   end_row_new,
   end_col_new
 )
-  for _, tree in pairs(self._trees) do
-    tree:edit(
+  for i, tree in pairs(self._trees) do
+    self._trees[i] = tree:edit(
       start_byte,
       end_byte_old,
       end_byte_new,

src/nvim/lua/treesitter.c

@@ -49,7 +49,8 @@ typedef struct {
 } TSLuaLoggerOpts;
 
 typedef struct {
-  TSTree *tree;
+  // We derive TSNode's, TSQueryCursor's, etc., from the TSTree, so it must not be mutated.
+  const TSTree *tree;
 } TSLuaTree;
 
 #ifdef INCLUDE_GENERATED_DECLARATIONS
@@ -490,7 +491,7 @@ static void push_ranges(lua_State *L, const TSRange *ranges, const size_t length
 static int parser_parse(lua_State *L)
 {
   TSParser *p = parser_check(L, 1);
-  TSTree *old_tree = NULL;
+  const TSTree *old_tree = NULL;
   if (!lua_isnil(L, 2)) {
     TSLuaTree *ud = luaL_checkudata(L, 2, TS_META_TREE);
     old_tree = ud ? ud->tree : NULL;
@@ -765,8 +766,8 @@ static struct luaL_Reg tree_meta[] = {
 /// Push tree interface on to the lua stack.
 ///
 /// The tree is not copied. Ownership of the tree is transferred from C to
-/// Lua. If needed use ts_tree_copy() in the caller
-static void push_tree(lua_State *L, TSTree *tree)
+/// Lua. If needed use ts_tree_copy() in the caller.
+static void push_tree(lua_State *L, const TSTree *tree)
 {
   if (tree == NULL) {
     lua_pushnil(L);
@@ -807,9 +808,12 @@ static int tree_edit(lua_State *L)
   TSInputEdit edit = { start_byte, old_end_byte, new_end_byte,
                        start_point, old_end_point, new_end_point };
 
-  ts_tree_edit(ud->tree, &edit);
+  TSTree *new_tree = ts_tree_copy(ud->tree);
+  ts_tree_edit(new_tree, &edit);
 
-  return 0;
+  push_tree(L, new_tree);  // [tree]
+
+  return 1;
 }
 
 static int tree_get_ranges(lua_State *L)
@@ -830,7 +834,12 @@ static int tree_get_ranges(lua_State *L)
 static int tree_gc(lua_State *L)
 {
   TSLuaTree *ud = luaL_checkudata(L, 1, TS_META_TREE);
-  ts_tree_delete(ud->tree);
+
+  // SAFETY: we can cast the const away because the tree is only garbage collected after all of its
+  // TSNode's, TSQuerCurors, etc., are unreachable (each contains a reference to the TSLuaTree)
+  TSTree *tree = (TSTree *)ud->tree;
+
+  ts_tree_delete(tree);
   return 0;
 }
 
@@ -844,11 +853,7 @@ static int tree_root(lua_State *L)
 {
   TSLuaTree *ud = luaL_checkudata(L, 1, TS_META_TREE);
 
-  // The tree may be mutate by tree_edit, but node needs that its tree is not
-  // mutated while the node is alive. So we make a copy of tree here.
-  TSTree *tree_copy = ts_tree_copy(ud->tree);
-
-  TSNode root = ts_tree_root_node(tree_copy);
+  TSNode root = ts_tree_root_node(ud->tree);
 
   TSNode *node_ud = lua_newuserdata(L, sizeof(TSNode));  // [node]
   *node_ud = root;
@@ -860,7 +865,7 @@ static int tree_root(lua_State *L)
   // Note: environments (fenvs) associated with userdata have no meaning in Lua
   // and are only used to associate a table.
   lua_createtable(L, 1, 0);  // [node, reftable]
-  push_tree(L, tree_copy);  // [node, reftable, tree]
+  lua_pushvalue(L, 1);  // [node, reftable, tree]
   lua_rawseti(L, -2, 1);  // [node, reftable]
   lua_setfenv(L, -2);  // [node]
 
@@ -1312,13 +1317,13 @@ static int node_root(lua_State *L)
 
 static int node_tree(lua_State *L)
 {
-  TSNode node = node_check(L, 1);
+  node_check(L, 1);
 
-  // The node's tree must not be mutated, but `tree_edit` may violate that. So
-  // we make a copy of the tree before pushing it to the LUA stack.
-  TSTree *tree = ts_tree_copy(node.tree);
+  // Get the tree from the node fenv. We cannot use `push_tree(node.tree)` here because that would
+  // cause a double free.
+  lua_getfenv(L, 1);  // [node, reftable]
+  lua_rawgeti(L, 2, 1);  // [node, reftable, tree]
 
-  push_tree(L, tree);
   return 1;
 }