From d46d931cd9ecd7ce2a852212314dc9c8b7dd7149 Mon Sep 17 00:00:00 2001
From: Folke Lemaitre <folke.lemaitre@gmail.com>
Date: Wed, 10 May 2023 08:31:52 +0200
Subject: [PATCH] perf: added debounce for calculating context

---
 README.md                  |  1 +
 lua/treesitter-context.lua | 34 +++++++++++-----------------------
 2 files changed, 12 insertions(+), 23 deletions(-)

diff --git a/README.md b/README.md
index c9545b28..cb0a65c1 100644
--- a/README.md
+++ b/README.md
@@ -207,6 +207,7 @@ require'treesitter-context'.setup{
   -- When separator is set, the context will only show up when there are at least 2 lines above cursorline.
   separator = nil,
   zindex = 20, -- The Z-index of the context window
+  throttle_ms = 100, -- Throttle context updates in milliseconds
 }
 ```
 
diff --git a/lua/treesitter-context.lua b/lua/treesitter-context.lua
index b4235c1b..ae86ec8f 100644
--- a/lua/treesitter-context.lua
+++ b/lua/treesitter-context.lua
@@ -33,6 +33,7 @@ local defaultConfig = {
   zindex = 20,
   mode = 'cursor', -- Choices: 'cursor', 'topline'
   separator = nil,
+  throttle_ms = 100, -- Throttle context updates in milliseconds
 }
 
 --- @type Config
@@ -392,31 +393,18 @@ end
 
 --- @generic F: function
 --- @param fn F
+--- @param ms? number
 --- @return F
-local function throttle_fn(fn)
-  local recalc_after_cooldown = false
-  local cooling_down = false
-  local function wrapped()
-    if cooling_down then
-      recalc_after_cooldown = true
-    else
-      local start = vim.loop.hrtime()
-      fn()
-      local elapsed_ms = math.floor((vim.loop.hrtime() - start) / 1e6)
-      -- If this took < 2ms, we don't need a cooldown period. This prevents the context floats from flickering
-      if elapsed_ms > 2 then
-        cooling_down = true
-        vim.defer_fn(function()
-          cooling_down = false
-          if recalc_after_cooldown then
-            recalc_after_cooldown = false
-            wrapped()
-          end
-        end, 20)
-      end
-    end
+local function throttle_fn(fn, ms)
+  ms = ms
+  local timer = vim.loop.new_timer()
+  return function(...)
+    local argv = vim.F.pack_len(...)
+    timer:start(ms, 0, function()
+      timer:stop()
+      vim.schedule_wrap(fn)(vim.F.unpack_len(argv))
+    end)
   end
-  return wrapped
 end
 
 local function close()