feat: Improve move_cursor. (#334)

Adds "sticky" option for `move_cursor`, making the cursor "stick" to the text as the buffer gets modified.

Author: kylechui

doc/nvim-surround.txt

@@ -613,16 +613,29 @@ configured separately. The default highlight group used is `Visual`:
 --------------------------------------------------------------------------------
 3.6. Cursor                                 *nvim-surround.config.move_cursor*
 
-By default, when a surround action is performed, the cursor moves to the
-beginning of the action.
+By default (or when `move_cursor = "begin"`), when a surround action is
+performed, the cursor moves to the beginning of the action.
 
     Old text                    Command         New text ~
     some_t*ext                  ysiw[           *[ some_text ]
     another { sample *}         ds{             another *sample
     (hello* world)              csbB            *{hello world}
 
-This behavior can be disabled by setting `move_cursor = false` in one of the
-setup functions.
+If `move_cursor` is set to `"sticky"`, the cursor will "stick" to the current
+character, and move with the text as the buffer changes.
+
+    Old text                    Command         New text ~
+    some_t*ext                  ysiw[           [ some_t*ext ]
+    another { sample *}         ds{             another sampl*e
+    (hello* world)              csbffoo<CR>     foo(hello* world)
+
+If `move_cursor` is set to `false`, the cursor won't move at all, regardless
+of how the buffer changes.
+
+    Old text                    Command         New text ~
+    some_t*ext                  ysiw[           [ some_*text ]
+    another { *sample }         ds{             another sa*mple
+    (hello* world)              csbffoo<CR>     foo(he*llo world)
 
 --------------------------------------------------------------------------------
 3.7. Indentation                           *nvim-surround.config.indent_lines*

lua/nvim-surround/annotations.lua

@@ -35,7 +35,7 @@
 ---@field surrounds table<string, surround>
 ---@field aliases table<string, string|string[]>
 ---@field highlight { duration: integer }
----@field move_cursor false|"begin"|"end"
+---@field move_cursor false|"begin"|"sticky"
 ---@field indent_lines function
 
 --[====================================================================================================================[
@@ -58,5 +58,5 @@
 ---@field surrounds? table<string, false|user_surround>
 ---@field aliases? table<string, false|string|string[]>
 ---@field highlight? { duration: false|integer }
----@field move_cursor? false|"begin"|"end"
+---@field move_cursor? false|"begin"|"sticky"
 ---@field indent_lines? false|function

lua/nvim-surround/buffer.lua

@@ -2,6 +2,11 @@ local config = require("nvim-surround.config")
 
 local M = {}
 
+M.namespace = {
+    highlight = vim.api.nvim_create_namespace("nvim-surround-highlight"),
+    extmark = vim.api.nvim_create_namespace("nvim-surround-extmark"),
+}
+
 --[====================================================================================================================[
                                                 Cursor helper functions
 --]====================================================================================================================]
@@ -24,11 +29,12 @@ M.set_curpos = function(pos)
 end
 
 -- Move the cursor to a location in the buffer, depending on the `move_cursor` setting.
----@param pos { first_pos: position, old_pos: position } Various positions in the buffer.
+---@param pos { first_pos: position, sticky_pos: position, old_pos: position } Various positions in the buffer.
 M.restore_curpos = function(pos)
-    -- TODO: Add a `last_pos` field for if `move_cursor` is set to "end"
     if config.get_opts().move_cursor == "begin" then
         M.set_curpos(pos.first_pos)
+    elseif config.get_opts().move_cursor == "sticky" then
+        M.set_curpos(pos.sticky_pos)
     elseif not config.get_opts().move_cursor then
         M.set_curpos(pos.old_pos)
     end
@@ -117,6 +123,29 @@ M.set_operator_marks = function(motion)
     M.set_mark(">", visual_marks[2])
 end
 
+-- Gets extmark position for the current buffer.
+---@param extmark integer The extmark ID number.
+---@return position @The position of the extmark in the buffer.
+---@nodiscard
+M.get_extmark = function(extmark)
+    local pos = vim.api.nvim_buf_get_extmark_by_id(0, M.namespace.extmark, extmark, {})
+    return { pos[1] + 1, pos[2] + 1 }
+end
+
+-- Creates an extmark for the given position.
+---@param pos position The position in the buffer.
+---@return integer @The extmark ID.
+---@nodiscard
+M.set_extmark = function(pos)
+    return vim.api.nvim_buf_set_extmark(0, M.namespace.extmark, pos[1] - 1, pos[2] - 1, {})
+end
+
+-- Deletes an extmark from the buffer.
+---@param extmark integer The extmark ID number.
+M.del_extmark = function(extmark)
+    vim.api.nvim_buf_del_extmark(0, M.namespace.extmark, extmark)
+end
+
 --[====================================================================================================================[
                                              Byte indexing helper functions
 --]====================================================================================================================]
@@ -257,11 +286,10 @@ M.highlight_selection = function(selection)
     if not selection then
         return
     end
-    local namespace = vim.api.nvim_create_namespace("NvimSurround")
 
     vim.highlight.range(
         0,
-        namespace,
+        M.namespace.highlight,
         "NvimSurroundHighlight",
         { selection.first_pos[1] - 1, selection.first_pos[2] - 1 },
         { selection.last_pos[1] - 1, selection.last_pos[2] - 1 },
@@ -273,8 +301,7 @@ end
 
 -- Clears all nvim-surround highlights for the buffer.
 M.clear_highlights = function()
-    local namespace = vim.api.nvim_create_namespace("NvimSurround")
-    vim.api.nvim_buf_clear_namespace(0, namespace, 0, -1)
+    vim.api.nvim_buf_clear_namespace(0, M.namespace.highlight, 0, -1)
     -- Force the screen to clear the highlight immediately
     vim.cmd.redraw()
 end

lua/nvim-surround/init.lua

@@ -64,12 +64,16 @@ M.normal_surround = function(args)
     local first_pos = args.selection.first_pos
     local last_pos = { args.selection.last_pos[1], args.selection.last_pos[2] + 1 }
 
+    local sticky_mark = buffer.set_extmark(M.normal_curpos)
     buffer.insert_text(last_pos, args.delimiters[2])
     buffer.insert_text(first_pos, args.delimiters[1])
+
     buffer.restore_curpos({
         first_pos = first_pos,
+        sticky_pos = buffer.get_extmark(sticky_mark),
         old_pos = M.normal_curpos,
     })
+    buffer.del_extmark(sticky_mark)
 
     if args.line_mode then
         config.get_opts().indent_lines(first_pos[1], last_pos[1] + #args.delimiters[1] + #args.delimiters[2] - 2)
@@ -92,6 +96,7 @@ M.visual_surround = function(args)
         return
     end
 
+    local sticky_mark = buffer.set_extmark(args.curpos)
     if vim.fn.visualmode() == "\22" then -- Visual block mode case (add delimiters to every line)
         if vim.o.selection == "exclusive" then
             last_pos[2] = last_pos[2] - 1
@@ -144,8 +149,10 @@ M.visual_surround = function(args)
     config.get_opts().indent_lines(first_pos[1], last_pos[1] + #delimiters[1] + #delimiters[2] - 2)
     buffer.restore_curpos({
         first_pos = first_pos,
+        sticky_pos = buffer.get_extmark(sticky_mark),
         old_pos = args.curpos,
     })
+    buffer.del_extmark(sticky_mark)
 end
 
 -- Delete a surrounding delimiter pair, if it exists.
@@ -165,17 +172,21 @@ M.delete_surround = function(args)
     local selections = utils.get_nearest_selections(args.del_char, "delete")
 
     if selections then
+        local sticky_mark = buffer.set_extmark(args.curpos)
         -- Delete the right selection first to ensure selection positions are correct
         buffer.delete_selection(selections.right)
         buffer.delete_selection(selections.left)
+
         config.get_opts().indent_lines(
             selections.left.first_pos[1],
             selections.left.first_pos[1] + selections.right.first_pos[1] - selections.left.last_pos[1]
         )
         buffer.restore_curpos({
             first_pos = selections.left.first_pos,
+            sticky_pos = buffer.get_extmark(sticky_mark),
             old_pos = args.curpos,
         })
+        buffer.del_extmark(sticky_mark)
     end
 
     cache.set_callback("v:lua.require'nvim-surround'.delete_callback")
@@ -221,13 +232,16 @@ M.change_surround = function(args)
             selections.right.first_pos[2] = space_end + 1
         end
 
+        local sticky_mark = buffer.set_extmark(args.curpos)
         -- Change the right selection first to ensure selection positions are correct
         buffer.change_selection(selections.right, delimiters[2])
         buffer.change_selection(selections.left, delimiters[1])
         buffer.restore_curpos({
             first_pos = selections.left.first_pos,
+            sticky_pos = buffer.get_extmark(sticky_mark),
             old_pos = args.curpos,
         })
+        buffer.del_extmark(sticky_mark)
 
         if args.line_mode then
             local first_pos = selections.left.first_pos