diff --git a/lua/treesitter-context.lua b/lua/treesitter-context.lua index 32e9d431..380ee31a 100644 --- a/lua/treesitter-context.lua +++ b/lua/treesitter-context.lua @@ -15,25 +15,51 @@ local all_contexts = {} --- @param ms? number --- @return F local function throttle_by_id(f, ms) - ms = ms or 200 + ms = ms or 150 + local timers = {} --- @type table + local state = {} --- @type table local waiting = {} --- @type table - return function(id) + + local schedule = function(id, wrapper) + state[id] = "scheduled" + vim.schedule_wrap(wrapper)(id, wrapper) + end + + local on_throttle_finish = function(id, wrapper) + assert(state[id] == "throttled") + 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] == "scheduled") + state[id] = "running" + f(id) + assert(state[id] == "running") + state[id] = "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 - vim.schedule(function() f(id) end) -- 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 + -- Don't set 'waiting' for 'scheduled' state since the callback is about to start. + -- Consequently, there is no need to run it again after throttling is completed. + if state[id] ~= "scheduled" then + waiting[id] = true + end end end @@ -173,7 +199,6 @@ function M.enable() if config.multiwindow then table.insert(update_events, 'WinResized') - table.insert(update_events, 'WinLeave') end autocmd(update_events, update)