From f0ab5e504b160d4bc60f52a02e8d2453052420d3 Mon Sep 17 00:00:00 2001
From: Liam Dyer <liamcdyer@gmail.com>
Date: Thu, 12 Dec 2024 13:12:56 -0500
Subject: [PATCH] feat: set cursor position for additional text edits

Closes #223
---
 lua/blink/cmp/completion/accept/init.lua |  8 ++--
 lua/blink/cmp/lib/text_edits.lua         | 49 +++++++++++++++++++++---
 lua/blink/cmp/lib/utils.lua              | 12 ++++++
 3 files changed, 59 insertions(+), 10 deletions(-)

diff --git a/lua/blink/cmp/completion/accept/init.lua b/lua/blink/cmp/completion/accept/init.lua
index 114c0041..60bea9c9 100644
--- a/lua/blink/cmp/completion/accept/init.lua
+++ b/lua/blink/cmp/completion/accept/init.lua
@@ -50,17 +50,15 @@ local function accept(ctx, item, callback)
         local temp_text_edit = vim.deepcopy(item.textEdit)
         temp_text_edit.newText = ''
         table.insert(all_text_edits, temp_text_edit)
-        text_edits_lib.apply(all_text_edits)
+        text_edits_lib.apply(all_text_edits, 0)
 
         -- Expand the snippet
         require('blink.cmp.config').snippets.expand(item.textEdit.newText)
 
-      -- OR Normal: Apply the text edit and move the cursor
+      -- Normal (non-snippet)
       else
         table.insert(all_text_edits, item.textEdit)
-        text_edits_lib.apply(all_text_edits)
-        -- TODO: should move the cursor only by the offset since text edit handles everything else?
-        ctx.set_cursor({ ctx.get_cursor()[1], item.textEdit.range.start.character + #item.textEdit.newText + offset })
+        text_edits_lib.apply(all_text_edits, offset)
       end
 
       -- Let the source execute the item itself
diff --git a/lua/blink/cmp/lib/text_edits.lua b/lua/blink/cmp/lib/text_edits.lua
index c9cae45c..acf07688 100644
--- a/lua/blink/cmp/lib/text_edits.lua
+++ b/lua/blink/cmp/lib/text_edits.lua
@@ -1,13 +1,20 @@
 local config = require('blink.cmp.config')
 local context = require('blink.cmp.completion.trigger.context')
+local utils = require('blink.cmp.lib.utils')
 
 local text_edits = {}
 
 --- Applies one or more text edits to the current buffer, assuming utf-8 encoding
 --- @param edits lsp.TextEdit[]
-function text_edits.apply(edits)
+--- @param offset number
+function text_edits.apply(edits, offset)
   local mode = context.get_mode()
-  if mode == 'default' then return vim.lsp.util.apply_text_edits(edits, vim.api.nvim_get_current_buf(), 'utf-8') end
+  local cursor = text_edits.get_post_apply_cursor(edits, offset)
+  if mode == 'default' then
+    vim.lsp.util.apply_text_edits(edits, vim.api.nvim_get_current_buf(), 'utf-8')
+    vim.api.nvim_win_set_cursor(0, cursor)
+    return
+  end
 
   assert(mode == 'cmdline', 'Unsupported mode for text edits: ' .. mode)
   assert(#edits == 1, 'Cmdline mode only supports one text edit. Contributions welcome!')
@@ -17,9 +24,41 @@ function text_edits.apply(edits)
   local edited_line = line:sub(1, edit.range.start.character)
     .. edit.newText
     .. line:sub(edit.range['end'].character + 1)
-  -- FIXME: for some reason, we have to set the cursor here, instead of later,
-  -- because this will override the cursor position set later
-  vim.fn.setcmdline(edited_line, edit.range.start.character + #edit.newText + 1)
+  vim.fn.setcmdline(edited_line, cursor[2])
+end
+
+--- Gets the cursor position after applying the text edits
+--- @param edits lsp.TextEdit[]
+--- @param offset number
+--- @return number[]
+function text_edits.get_post_apply_cursor(edits, offset)
+  local mode = context.get_mode()
+  local cursor = context.get_cursor()
+  local cursor_row = cursor[1] - 1 -- convert to 0-indexed
+  local cursor_col = cursor[2] -- already 0-indexed
+  if mode == 'default' then
+    -- Get the first edit that intersects with the cursor
+    local edit = utils.find(edits, function(edit)
+      if edit.range.start.line == cursor_row then
+        if edit.range.start.line ~= edit.range['end'].line then return true end
+        return edit.range.start.character >= cursor_col and edit.range['end'].character <= cursor_col
+      end
+      if edit.range['end'].line + 1 == cursor_row then return edit.range['end'].character >= cursor_col end
+      return edit.range.start.line + 1 >= cursor_row and edit.range['end'].line + 1 <= cursor_row
+    end)
+    if not edit then return cursor end
+
+    -- Move the cursor to the end of the edit
+    local lines = vim.split(edit.newText, '\n')
+    return { edit.range.start.line + #lines, edit.range.start.character + #lines[#lines] + offset }
+  end
+
+  assert(mode == 'cmdline', 'Unsupported mode for text edits: ' .. mode)
+  assert(#edits == 1, 'Cmdline mode only supports one text edit. Contributions welcome!')
+
+  -- TODO: support multiple edits in cmdline mode
+  local edit = edits[1]
+  return { 1, edit.range.start.character + #edit.newText + 1 + offset }
 end
 
 ------- Undo -------
diff --git a/lua/blink/cmp/lib/utils.lua b/lua/blink/cmp/lib/utils.lua
index e047b7cc..15710ad4 100644
--- a/lua/blink/cmp/lib/utils.lua
+++ b/lua/blink/cmp/lib/utils.lua
@@ -89,6 +89,18 @@ function utils.find_idx(arr, predicate)
   return nil
 end
 
+--- Finds an item in an array using a predicate function
+--- @generic T
+--- @param arr T[]
+--- @param predicate fun(item: T): boolean
+--- @return T | nil
+function utils.find(arr, predicate)
+  for _, v in ipairs(arr) do
+    if predicate(v) then return v end
+  end
+  return nil
+end
+
 --- Slices an array
 --- @generic T
 --- @param arr T[]