refactor: make go_to_context() get contexts directly

- Removes the need for the internal `all_contexts` variable.
- Remove `bufnr` argument from `Context.get()` and make `winid`
  optional.

Author: lewis6991

lua/treesitter-context.lua

@@ -23,9 +23,6 @@ end
 --- @module 'treesitter-context.render'
 local Render = defer_require('treesitter-context.render')
 
---- @type table<integer, Range4[]>
-local all_contexts = {}
-
 --- Schedule a function to run on the next event loop iteration.
 --- If the function is called again within 150ms, it will be scheduled
 --- again to run on the next event loop iteration. This means that
@@ -110,17 +107,14 @@ local update_win = throttle_by_id(function(winid)
     return
   end
 
-  local bufnr = api.nvim_win_get_buf(winid)
-
-  local context_ranges, context_lines = require('treesitter-context.context').get(bufnr, winid)
-  all_contexts[bufnr] = context_ranges
+  local context_ranges, context_lines = require('treesitter-context.context').get(winid)
 
   if not context_ranges or #context_ranges == 0 then
     Render.close(winid)
     return
   end
 
-  Render.open(bufnr, winid, context_ranges, assert(context_lines))
+  Render.open(winid, context_ranges, assert(context_lines))
 end)
 
 local multiwindow_events = {
@@ -312,8 +306,7 @@ function M.go_to_context(depth)
   depth = depth or 1
   local line = api.nvim_win_get_cursor(0)[1]
   local context = nil
-  local bufnr = api.nvim_get_current_buf()
-  local contexts = all_contexts[bufnr] or {}
+  local contexts = require('treesitter-context.context').get() or {}
 
   for idx = #contexts, 1, -1 do
     local c = contexts[idx]
@@ -326,7 +319,7 @@ function M.go_to_context(depth)
     end
   end
 
-  if context == nil then
+  if not context then
     return
   end
 

lua/treesitter-context/context.lua

@@ -312,10 +312,12 @@ local function range_is_valid(range)
   return true
 end
 
---- @param bufnr integer
---- @param winid integer
+--- @param winid? integer
 --- @return Range4[]?, string[]?
-function M.get(bufnr, winid)
+function M.get(winid)
+  winid = winid or api.nvim_get_current_win()
+  local bufnr = api.nvim_win_get_buf(winid)
+
   -- vim.treesitter.get_parser() calls bufload(), but we don't actually want to load the buffer:
   -- this method is called during plugin init, before other plugins or the user's config
   -- have a chance to initialize.

lua/treesitter-context/render.lua

@@ -449,11 +449,11 @@ end
 
 local M = {}
 
---- @param bufnr integer
 --- @param winid integer
 --- @param ctx_ranges Range4[]
 --- @param ctx_lines string[]
-function M.open(bufnr, winid, ctx_ranges, ctx_lines)
+function M.open(winid, ctx_ranges, ctx_lines)
+  local bufnr = api.nvim_win_get_buf(winid)
   local gutter_width = get_gutter_width(winid)
   local win_width = math.max(1, api.nvim_win_get_width(winid) - gutter_width)
 

test/contexts_spec.lua

@@ -176,20 +176,15 @@ for _, lang in ipairs(langs_with_queries) do
       for cursor_row, context_rows in pairs(contexts) do
         it(('line %s in %s'):format(cursor_row, test_file), function()
           cmd('edit ' .. test_file)
-          local bufnr = api.nvim_get_current_buf()
           local winid = api.nvim_get_current_win()
           api.nvim_win_set_cursor(winid, { cursor_row + 1, 0 })
           assert(fn.getline('.'):match('{{CURSOR}}'))
           feed(string.format('zt%d<C-y>', #context_rows + 2))
 
           --- @type [integer,integer,integer,integer][]
-          local ranges = exec_lua(
-            [[
-          return require('treesitter-context.context').get(...)
-        ]],
-            bufnr,
-            winid
-          )
+          local ranges = exec_lua(function(...)
+            return assert(require('treesitter-context.context').get(...))
+          end, winid)
 
           local act_context_rows = {} --- @type integer[]
           for _, r in ipairs(ranges) do