fix: throttler

lua/treesitter-context.lua

@@ -15,25 +15,52 @@ local all_contexts = {}
 --- @param ms? number
 --- @return F
 local function throttle_by_id(f, ms)
-  ms = ms or 200
+  ms = ms or 150
+
+  local State = {
+    Scheduled = {},
+    Running = {},
+    Throttled = {},
+  }
+
   local timers = {} --- @type table<any,uv.uv_timer_t>
-  local waiting = {} --- @type table<any,boolean>
-  return function(id)
+  local state = {} --- @type table<any,table>
+  local waiting = {} --- @type table<any, boolean>
+
+  local schedule = function(id, wrapper)
+    state[id] = State.Scheduled
+    vim.schedule_wrap(wrapper)(id, wrapper)
+  end
+
+  local on_throttle_finish = function (id, wrapper)
+    if waiting[id] == nil then
+      timers[id] = nil
+      state[id] = nil
+      return
+    end
+    waiting[id] = nil
+    schedule(id, wrapper)
+  end
+
+  local wrapper = function(id, wrapper)
+    assert(state[id] == State.Scheduled)
+    state[id] = State.Running
+    f(id)
+    assert(state[id] == State.Running)
+    state[id] = State.Throttled
+
     if timers[id] == nil then
       timers[id] = assert(vim.loop.new_timer())
-    else
-      waiting[id] = true
-      return
     end
+    timers[id]:start(ms, 0, function() on_throttle_finish(id, wrapper) end)
+  end
 
-    f(id) -- first call, execute immediately
-    timers[id]:start(ms, 0, function()
-      if waiting[id] then
-        f(id) -- only execute if there are calls waiting
-      end
-      waiting[id] = nil
-      timers[id] = nil
-    end)
+  return function(id)
+    if state[id] == nil then
+      schedule(id, wrapper)
+      return
+    end
+    waiting[id] = true
   end
 end
 
@@ -67,10 +94,16 @@ local function cannot_open(bufnr, winid)
     or api.nvim_win_get_height(winid) < config.min_window_height
 end
 
--- Always use deferred execution to prevent from segfaults in neovim, like this one:
--- https://github.com/neovim/neovim/issues/31236
 ---@param winid integer
-local update_single_context = throttle_by_id(vim.schedule_wrap(function(winid)
+local update_single_context = throttle_by_id(function(winid)
+  -- Remove leaked contexts firstly.
+  local current_win = api.nvim_get_current_win()
+  if config.multiwindow then
+    require('treesitter-context.render').close_leaked_contexts()
+  else
+    require('treesitter-context.render').close_other_contexts(current_win)
+  end
+
   -- 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) or vim.fn.getcmdtype() ~= '' then
@@ -79,7 +112,7 @@ local update_single_context = throttle_by_id(vim.schedule_wrap(function(winid)
 
   local bufnr = api.nvim_win_get_buf(winid)
 
-  if cannot_open(bufnr, winid) or not config.multiwindow and winid ~= api.nvim_get_current_win() then
+  if cannot_open(bufnr, winid) or not config.multiwindow and winid ~= current_win then
     require('treesitter-context.render').close(winid)
     return
   end
@@ -95,7 +128,7 @@ local update_single_context = throttle_by_id(vim.schedule_wrap(function(winid)
   assert(context_lines)
 
   require('treesitter-context.render').open(bufnr, winid, context_ranges, context_lines)
-end))
+end)
 
 ---@param args table
 local function update(args)