fix(treesitter): find buffer in multiple windows #28922

Problem:
1. When interacting with multiple :InspectTree and the source buffer
   windows there is a high chance of errors due to the window ids not
   being updated and validated.
2. Not all InspectTree windows were closed when the source buffer was
   closed.

Solution:
1. Update InspectTree window id on `CursorMoved` event and validate
   source buffer window id before trying to navigate to it.
2. Close all InspectTree windows

2 files changed, 74 insertions, 1 deletions

diff --git a/runtime/lua/vim/treesitter/dev.lua b/runtime/lua/vim/treesitter/dev.lua
index 5c91f101c0..ca8cf85eda 100644
--- a/runtime/lua/vim/treesitter/dev.lua
+++ b/runtime/lua/vim/treesitter/dev.lua
@@ -325,7 +325,10 @@ function M.inspect_tree(opts)
 
   opts = opts or {}
 
+  -- source buffer
   local buf = api.nvim_get_current_buf()
+
+  -- window id for source buffer
   local win = api.nvim_get_current_win()
   local treeview = assert(TSTreeView:new(buf, opts.lang))
 
@@ -334,12 +337,14 @@ function M.inspect_tree(opts)
     close_win(vim.b[buf].dev_inspect)
   end
 
+  -- window id for tree buffer
   local w = opts.winid
   if not w then
     vim.cmd(opts.command or '60vnew')
     w = api.nvim_get_current_win()
   end
 
+  -- tree buffer
   local b = opts.bufnr
   if b then
     api.nvim_win_set_buf(w, b)
@@ -375,6 +380,12 @@ function M.inspect_tree(opts)
     callback = function()
       local row = api.nvim_win_get_cursor(w)[1]
       local lnum, col = treeview:get(row).node:start()
+
+      -- update source window if original was closed
+      if not api.nvim_win_is_valid(win) then
+        win = vim.fn.win_findbuf(buf)[1]
+      end
+
       api.nvim_set_current_win(win)
       api.nvim_win_set_cursor(win, { lnum + 1, col })
     end,
@@ -432,6 +443,7 @@ function M.inspect_tree(opts)
         return true
       end
 
+      w = api.nvim_get_current_win()
       api.nvim_buf_clear_namespace(buf, treeview.ns, 0, -1)
       local row = api.nvim_win_get_cursor(w)[1]
       local lnum, col, end_lnum, end_col = treeview:get(row).node:range()
@@ -441,6 +453,11 @@ function M.inspect_tree(opts)
         hl_group = 'Visual',
       })
 
+      -- update source window if original was closed
+      if not api.nvim_win_is_valid(win) then
+        win = vim.fn.win_findbuf(buf)[1]
+      end
+
       local topline, botline = vim.fn.line('w0', win), vim.fn.line('w$', win)
 
       -- Move the cursor if highlighted range is completely out of view
@@ -506,7 +523,10 @@ function M.inspect_tree(opts)
     buffer = buf,
     once = true,
     callback = function()
-      close_win(w)
+      -- close all tree windows
+      for _, window in pairs(vim.fn.win_findbuf(b)) do
+        close_win(window)
+      end
     end,
   })
 end
diff --git a/test/functional/treesitter/inspect_tree_spec.lua b/test/functional/treesitter/inspect_tree_spec.lua
index f5acfe7c4a..b403cca735 100644
--- a/test/functional/treesitter/inspect_tree_spec.lua
+++ b/test/functional/treesitter/inspect_tree_spec.lua
@@ -114,4 +114,57 @@ describe('vim.treesitter.inspect_tree', function()
             (fenced_code_block_delimiter)))) ; [2, 0] - [2, 3] markdown
       ]]
   end)
+
+  it('updates source and tree buffer windows and closes them correctly', function()
+    insert([[
+      print()
+      ]])
+
+    -- setup two windows for the source buffer
+    exec_lua([[
+      source_win = vim.api.nvim_get_current_win()
+      vim.api.nvim_open_win(0, false, {
+        win = 0,
+        split = 'left'
+      })
+    ]])
+
+    -- setup three windows for the tree buffer
+    exec_lua([[
+      vim.treesitter.start(0, 'lua')
+      vim.treesitter.inspect_tree()
+      tree_win = vim.api.nvim_get_current_win()
+      tree_win_copy_1 = vim.api.nvim_open_win(0, false, {
+        win = 0,
+        split = 'left'
+      })
+      tree_win_copy_2 = vim.api.nvim_open_win(0, false, {
+        win = 0,
+        split = 'left'
+      })
+    ]])
+
+    -- close original source window
+    exec_lua('vim.api.nvim_win_close(source_win, false)')
+
+    -- navigates correctly to the remaining source buffer window
+    feed('<CR>')
+    eq('', n.api.nvim_get_vvar('errmsg'))
+
+    -- close original tree window
+    exec_lua([[
+       vim.api.nvim_set_current_win(tree_win_copy_1)
+       vim.api.nvim_win_close(tree_win, false)
+    ]])
+
+    -- navigates correctly to the remaining source buffer window
+    feed('<CR>')
+    eq('', n.api.nvim_get_vvar('errmsg'))
+
+    -- close source buffer window and all remaining tree windows
+    t.pcall_err(exec_lua, 'vim.api.nvim_win_close(0, false)')
+
+    eq(false, exec_lua('return vim.api.nvim_win_is_valid(tree_win_copy_1)'))
+    eq(false, exec_lua('return vim.api.nvim_win_is_valid(tree_win_copy_2)'))
+  end)
 end)