Merge pull request #171 from ilyaZar/codex/pr-159-visual-context-bounds

fix visual context bounds

2 files changed, 111 insertions, 6 deletions

diff --git a/lua/99/prompt-settings.lua b/lua/99/prompt-settings.lua
index fb219db..a688a0e 100644
--- a/lua/99/prompt-settings.lua
+++ b/lua/99/prompt-settings.lua
@@ -1,7 +1,14 @@
----@param buffer number
+---@param range _99.Range
+---@param n number lines of context above and below the selection
 ---@return string
-local function get_file_contents(buffer)
-  local lines = vim.api.nvim_buf_get_lines(buffer, 0, -1, false)
+local function get_surrounding_context(range, n)
+  local start_row, _ = range.start:to_vim()
+  local end_row, _ = range.end_:to_vim()
+  local line_count = vim.api.nvim_buf_line_count(range.buffer)
+  local from = math.max(start_row - n, 0)
+  local to = math.min(end_row + 1 + n, line_count)
+  local lines = vim.api.nvim_buf_get_lines(range.buffer, from, to, false)
+
   return table.concat(lines, "\n")
 end
 
@@ -148,13 +155,13 @@ consider the context of the selection and what you are suppose to be implementin
 <SELECTION_CONTENT>
 %s
 </SELECTION_CONTENT>
-<FILE_CONTAINING_SELECTION>
+<SURROUNDING_CONTEXT>
 %s
-</FILE_CONTAINING_SELECTION>
+</SURROUNDING_CONTEXT>
 ]],
       range:to_string(),
       range:to_text(),
-      get_file_contents(range.buffer)
+      get_surrounding_context(range, 100)
     )
   end,
   read_tmp = function()
diff --git a/lua/99/test/visual_context_bounds_spec.lua b/lua/99/test/visual_context_bounds_spec.lua
new file mode 100644
index 0000000..5b89ad4
--- /dev/null
+++ b/lua/99/test/visual_context_bounds_spec.lua
@@ -0,0 +1,98 @@
+-- luacheck: globals describe it assert after_each
+local prompt_settings = require("99.prompt-settings")
+local test_utils = require("99.test.test_utils")
+local geo = require("99.geo")
+
+local eq = assert.are.same
+local Range = geo.Range
+local Point = geo.Point
+
+local function make_lines(count, overrides)
+  local out = {}
+  for i = 1, count do
+    out[i] = overrides[i] or string.format("line %04d", i)
+  end
+  return out
+end
+
+local function range_for_line(buffer, row)
+  return Range:new(
+    buffer,
+    Point:from_1_based(row, 1),
+    Point:from_1_based(row, 1)
+  )
+end
+
+local function prompt_context_lines(range)
+  local prompt = prompt_settings.prompts.visual_selection(range)
+  local context = prompt:match(
+    "<SURROUNDING_CONTEXT>\n(.-)\n</SURROUNDING_CONTEXT>"
+  )
+
+  assert.is_not_nil(context)
+  if context == "" then
+    return {}
+  end
+
+  return vim.split(context, "\n", { plain = true })
+end
+
+describe("visual context", function()
+  after_each(function()
+    test_utils.clean_files()
+  end)
+
+  it("should handle selections at the top of a file", function()
+    local buffer = test_utils.create_file(
+      make_lines(1000, {
+        [1] = "selected line",
+        [101] = "last line in context",
+        [901] = "line that should not appear",
+        [1000] = "bottom of file",
+      }),
+      "lua",
+      1,
+      0
+    )
+
+    local context = prompt_context_lines(range_for_line(buffer, 1))
+
+    eq(101, #context)
+    eq("selected line", context[1])
+    eq("last line in context", context[#context])
+    assert.is_false(vim.tbl_contains(context, "line that should not appear"))
+  end)
+
+  it("should keep context around middle selections", function()
+    local buffer = test_utils.create_file(
+      make_lines(1000, {
+        [400] = "first line in context",
+        [500] = "selected line",
+        [600] = "last line in context",
+      }),
+      "lua",
+      500,
+      0
+    )
+
+    local context = prompt_context_lines(range_for_line(buffer, 500))
+
+    eq(201, #context)
+    eq("first line in context", context[1])
+    eq("selected line", context[101])
+    eq("last line in context", context[#context])
+  end)
+
+  it("should keep small files whole", function()
+    local content = {
+      "first available line",
+      "selected middle line",
+      "last available line",
+    }
+    local buffer = test_utils.create_file(content, "lua", 2, 0)
+
+    local context = prompt_context_lines(range_for_line(buffer, 2))
+
+    eq(content, context)
+  end)
+end)