feat: multiwindow support

doc/nvim-treesitter-context.txt

@@ -23,6 +23,9 @@ not required.
       -- Enable this plugin (Can be enabled/disabled later via commands)
       enable = true,
 
+      -- Enable multiwindow support.
+      multiwindow = false,
+
       -- How many lines the window should span. Values <= 0 mean no limit.
       max_lines = 0,
 

lua/treesitter-context.lua

@@ -1,6 +1,8 @@
 local api = vim.api
 
 local config = require('treesitter-context.config')
+local render = require('treesitter-context.render')
+local context = require('treesitter-context.context')
 
 local augroup = api.nvim_create_augroup
 local command = api.nvim_create_user_command
@@ -10,45 +12,8 @@ local enabled = false
 --- @type table<integer, Range4[]>
 local all_contexts = {}
 
---- @generic F: function
---- @param f F
---- @param ms? number
---- @return F
-local function throttle(f, ms)
-  ms = ms or 200
-  local timer = assert(vim.loop.new_timer())
-  local waiting = 0
-  return function()
-    if timer:is_active() then
-      waiting = waiting + 1
-      return
-    end
-    waiting = 0
-    f() -- first call, execute immediately
-    timer:start(ms, 0, function()
-      if waiting > 1 then
-        vim.schedule(f) -- only execute if there are calls waiting
-      end
-    end)
-  end
-end
-
-local had_open = false
-
-local function close()
-  if had_open then
-    require('treesitter-context.render').close()
-  end
-end
-
---- @param bufnr integer
---- @param winid integer
---- @param ctx_ranges Range4[]
---- @param ctx_lines string[]
-local function open(bufnr, winid, ctx_ranges, ctx_lines)
-  had_open = true
-  require('treesitter-context.render').open(bufnr, winid, ctx_ranges, ctx_lines)
-end
+--- @type table<integer, boolean>
+local win_update_scheduled = {}
 
 local attached = {} --- @type table<integer,true>
 
@@ -82,27 +47,71 @@ local function can_open(bufnr, winid)
   return true
 end
 
-local update = throttle(function()
-  local bufnr = api.nvim_get_current_buf()
-  local winid = api.nvim_get_current_win()
+---@param winid integer
+local update_single_context = function(winid)
+  -- Since the update is performed asynchronously, the window may be closed at this moment.
+  -- Therefore, we need to check if it is still valid.
+  if not api.nvim_win_is_valid(winid) then
+    return
+  end
+
+  local bufnr = api.nvim_win_get_buf(winid)
 
   if not can_open(bufnr, winid) then
-    close()
+    render.close(winid)
     return
   end
 
-  local context, context_lines = require('treesitter-context.context').get(bufnr, winid)
-  all_contexts[bufnr] = context
+  local context_ranges, context_lines = context.get(bufnr, winid)
+  all_contexts[bufnr] = context_ranges
 
-  if not context or #context == 0 then
-    close()
+  if not context_ranges or #context_ranges == 0 then
+    render.close(winid)
     return
   end
 
   assert(context_lines)
 
-  open(bufnr, winid, context, context_lines)
-end)
+  render.open(bufnr, winid, context_ranges, context_lines)
+end
+
+---@param winid integer
+local schedule_single_context_update = function(winid)
+  if not win_update_scheduled[winid] then
+    win_update_scheduled[winid] = true
+    vim.schedule(function()
+      win_update_scheduled[winid] = nil
+      update_single_context(winid)
+    end)
+  end
+end
+
+---@param args table
+local on_update_event = function(args)
+  local multiwindow_events = { "WinResized", "OptionSet", "User" }
+  if config.multiwindow and vim.tbl_contains(multiwindow_events, args.event) then
+    for _, winid in pairs(api.nvim_list_wins()) do
+      schedule_single_context_update(winid)
+    end
+  else
+    schedule_single_context_update(api.nvim_get_current_win())
+  end
+end
+
+---@param args table
+local function on_close_event(args)
+  if args.event == 'WinClosed' then
+    render.close(tonumber(args.match))
+  else
+    render.close(api.nvim_get_current_win())
+  end
+end
+
+local function close_all()
+  for _, winid in pairs(api.nvim_list_wins()) do
+    render.close(winid)
+  end
+end
 
 local M = {
   config = config,
@@ -125,7 +134,20 @@ function M.enable()
 
   attached[cbuf] = true
 
-  autocmd({ 'WinScrolled', 'BufEnter', 'WinEnter', 'VimResized' }, update)
+  local update_events = {
+    'WinScrolled',
+    'BufEnter',
+    'WinEnter',
+    'VimResized',
+    'CursorMoved',
+  }
+
+  if config.multiwindow then
+    table.insert(update_events, 'WinResized')
+    table.insert(update_events, 'WinLeave')
+  end
+
+  autocmd(update_events, on_update_event)
 
   autocmd('BufReadPost', function(args)
     attached[args.buf] = nil
@@ -138,27 +160,31 @@ function M.enable()
     attached[args.buf] = nil
   end)
 
-  autocmd('CursorMoved', update)
-
   autocmd('OptionSet', function(args)
     if args.match == 'number' or args.match == 'relativenumber' then
-      update()
+      on_update_event(args)
     end
   end)
 
-  autocmd({ 'BufLeave', 'WinLeave' }, close)
+  if config.multiwindow then
+    autocmd({ 'WinClosed' }, on_close_event)
+  else
+    autocmd({ 'BufLeave', 'WinLeave' }, on_close_event)
+  end
 
-  autocmd('User', close, { pattern = 'SessionSavePre' })
-  autocmd('User', update, { pattern = 'SessionSavePost' })
+  autocmd('User', close_all, { pattern = 'SessionSavePre' })
+  autocmd('User', on_update_event, { pattern = 'SessionSavePost' })
 
-  update()
+  for _, winid in pairs(api.nvim_list_wins()) do
+    update_single_context(winid)
+  end
   enabled = true
 end
 
 function M.disable()
   augroup('treesitter_context_update', {})
   attached = {}
-  close()
+  close_all()
   enabled = false
 end
 
@@ -210,7 +236,7 @@ end
 function M.go_to_context(depth)
   depth = depth or 1
   local line = api.nvim_win_get_cursor(0)[1]
-  local context = nil
+  local context_ranges = nil
   local bufnr = api.nvim_get_current_buf()
   local contexts = all_contexts[bufnr] or {}
 
@@ -220,17 +246,17 @@ function M.go_to_context(depth)
       break
     end
     if c[1] + 1 < line then
-      context = c
+      context_ranges = c
       depth = depth - 1
     end
   end
 
-  if context == nil then
+  if context_ranges == nil then
     return
   end
 
   vim.cmd([[ normal! m' ]]) -- add current cursor position to the jump list
-  api.nvim_win_set_cursor(0, { context[1] + 1, context[2] })
+  api.nvim_win_set_cursor(0, { context_ranges[1] + 1, context_ranges[2] })
 end
 
 return M

lua/treesitter-context/config.lua

@@ -1,6 +1,6 @@
-
 --- @class (exact) TSContext.Config
 --- @field enable boolean
+--- @field multiwindow boolean
 --- @field max_lines integer
 --- @field min_window_height integer
 --- @field line_numbers boolean
@@ -16,6 +16,9 @@
 --- Enable this plugin (Can be enabled/disabled later via commands)
 --- @field enable? boolean
 ---
+--- Enable multiwindow support.
+--- @field multiwindow? boolean
+---
 --- How many lines the window should span. Values <= 0 mean no limit.
 --- @field max_lines? integer
 ---
@@ -45,6 +48,7 @@
 --- @type TSContext.Config
 local default_config = {
   enable = true,
+  multiwindow = false,
   max_lines = 0, -- no limit
   min_window_height = 0,
   line_numbers = true,

lua/treesitter-context/context.lua

@@ -64,8 +64,7 @@ end
 --- @param node TSNode
 --- @param query vim.treesitter.Query
 --- @return Range4?
-local context_range = cache.memoize(function(node, query)
-  local bufnr = api.nvim_get_current_buf()
+local context_range = cache.memoize(function(node, query, bufnr)
   local range = { node:range() } --- @type Range4
   range[3] = range[1] + 1
   range[4] = 0
@@ -145,16 +144,17 @@ local function trim_contexts(context_ranges, context_lines, trim, top)
 end
 
 --- @param range Range4
+--- @param bufnr integer
 --- @return Range4, string[]
-local function get_text_for_range(range)
+local function get_text_for_range(range, bufnr)
   local start_row, end_row, end_col = range[1], range[3], range[4]
 
   if end_col == 0 then
     end_row = end_row - 1
     end_col = -1
   end
 
-  local lines = api.nvim_buf_get_text(0, start_row, 0, end_row, -1, {})
+  local lines = api.nvim_buf_get_text(bufnr, start_row, 0, end_row, -1, {})
 
   -- Strip any empty lines from the node
   while #lines > 0 do
@@ -316,9 +316,9 @@ function M.get(bufnr, winid)
 
         -- Only process the parent if it is not in view.
         if parent_start_row < contexts_end_row then
-          local range0 = context_range(parent, query)
+          local range0 = context_range(parent, query, bufnr)
           if range0 and range_is_valid(range0) then
-            local range, lines = get_text_for_range(range0)
+            local range, lines = get_text_for_range(range0, bufnr)
             if range_is_valid(range) then
               local last_context = context_ranges[#context_ranges]
               if last_context and parent_start_row == last_context[1] then