feat: Repeat operations for adding delimiters. (#360)

* feat: Very jank way of supporting numeric prefixes for normal_surround.

* fix: Slightly smarter way of handling numeric prefixes for normal_surround.

* feat: Repeat visual surrounds.

* tests: Repeating delimiters.

Author: kylechui

lua/nvim-surround/cache.lua

@@ -2,7 +2,7 @@ local M = {}
 
 -- These variables hold cache values for dot-repeating the three actions
 
----@type { delimiters: string[][]|nil, line_mode: boolean }
+---@type { delimiters: string[][]|nil, line_mode: boolean, count: integer }
 M.normal = {}
 ---@type { char: string }
 M.delete = {}

lua/nvim-surround/init.lua

@@ -50,23 +50,34 @@ M.normal_surround = function(args)
     local config = require("nvim-surround.config")
     local buffer = require("nvim-surround.buffer")
     local cache = require("nvim-surround.cache")
+    local utils = require("nvim-surround.utils")
     -- Call the operatorfunc if it has not been called yet
     if not args.selection then
-        -- Clear the normal cache (since it was user-called)
-        cache.normal = { line_mode = args.line_mode }
+        -- Clear the normal cache's delimiters (since it was user-called)
+        cache.normal = { line_mode = args.line_mode, count = vim.v.count1 }
         M.normal_curpos = buffer.get_curpos()
         M.pending_surround = true
 
         vim.go.operatorfunc = "v:lua.require'nvim-surround'.normal_callback"
-        return "g@"
+
+        -- Very jank way of resetting v:count to 1 before getting the motion, to ensure that the count
+        -- does not multiply against the motion
+        local del_str = ""
+        local n = vim.v.count1
+        while n > 1 do
+            del_str = del_str .. "<Del>"
+            n = n / 10
+        end
+        return del_str .. "g@"
     end
 
     local first_pos = args.selection.first_pos
     local last_pos = { args.selection.last_pos[1], args.selection.last_pos[2] + 1 }
 
+    local delimiters = utils.repeat_delimiters(args.delimiters, cache.normal.count)
     local sticky_pos = buffer.with_extmark(M.normal_curpos, function()
-        buffer.insert_text(last_pos, args.delimiters[2])
-        buffer.insert_text(first_pos, args.delimiters[1])
+        buffer.insert_text(last_pos, delimiters[2])
+        buffer.insert_text(first_pos, delimiters[1])
     end)
     buffer.restore_curpos({
         first_pos = first_pos,
@@ -75,7 +86,7 @@ M.normal_surround = function(args)
     })
 
     if args.line_mode then
-        config.get_opts().indent_lines(first_pos[1], last_pos[1] + #args.delimiters[1] + #args.delimiters[2] - 2)
+        config.get_opts().indent_lines(first_pos[1], last_pos[1] + #delimiters[1] + #delimiters[2] - 2)
     end
     M.pending_surround = false
 end
@@ -86,17 +97,19 @@ M.visual_surround = function(args)
     local config = require("nvim-surround.config")
     local buffer = require("nvim-surround.buffer")
     local input = require("nvim-surround.input")
+    local utils = require("nvim-surround.utils")
     local ins_char = input.get_char()
 
     if vim.fn.visualmode() == "V" then
         args.line_mode = true
     end
-    local delimiters = config.get_delimiters(ins_char, args.line_mode)
     local first_pos, last_pos = buffer.get_mark("<"), buffer.get_mark(">")
-    if not delimiters or not first_pos or not last_pos then
+    local raw_delimiters = config.get_delimiters(ins_char, args.line_mode)
+    if not raw_delimiters or not first_pos or not last_pos then
         return
     end
 
+    local delimiters = utils.repeat_delimiters(raw_delimiters, vim.v.count1)
     if vim.o.selection == "exclusive" then
         last_pos[2] = last_pos[2] - 1
     end

lua/nvim-surround/utils.lua

@@ -7,6 +7,22 @@ local M = {}
 -- Do nothing.
 M.NOOP = function() end
 
+-- Repeats a delimiter pair n times.
+---@param delimiters delimiter_pair The delimiters to be repeated.
+---@param n integer The number of times to repeat the delimiters.
+---@return delimiter_pair @The repeated delimiters.
+---@nodiscard
+M.repeat_delimiters = function(delimiters, n)
+    local acc = { { "" }, { "" } }
+    for _ = 1, n do
+        acc[1][#acc[1]] = acc[1][#acc[1]] .. delimiters[1][1]
+        vim.list_extend(acc[1], delimiters[1], 2)
+        acc[2][#acc[2]] = acc[2][#acc[2]] .. delimiters[2][1]
+        vim.list_extend(acc[2], delimiters[2], 2)
+    end
+    return acc
+end
+
 -- Gets the nearest two selections for the left and right surrounding pair.
 ---@param char string|nil A character representing what kind of surrounding pair is to be selected.
 ---@param action "delete"|"change" A string representing what action is being performed.

tests/basics_spec.lua

@@ -85,7 +85,7 @@ describe("nvim-surround", function()
         set_lines({ "here", "we", "have", "several", "lines" })
         set_curpos({ 2, 2 })
         vim.cmd("normal 3ySS`")
-        check_lines({ "here", "`", "we", "have", "several", "`", "lines" })
+        check_lines({ "here", "`", "`", "`", "we", "`", "`", "`", "have", "several", "lines" })
     end)
 
     it("can surround empty lines", function()
@@ -769,4 +769,101 @@ describe("nvim-surround", function()
             "{hello world}",
         })
     end)
+
+    it("can handle number prefixing for normal surround", function()
+        set_lines({
+            "some more placeholder text",
+            "some more lines",
+            "hello world",
+        })
+        set_curpos({ 2, 7 })
+        vim.cmd("normal 2ysiw*")
+        check_lines({
+            "some more placeholder text",
+            "some **more** lines",
+            "hello world",
+        })
+
+        vim.cmd("normal 3yssa")
+        check_lines({
+            "some more placeholder text",
+            "<<<some **more** lines>>>",
+            "hello world",
+        })
+
+        set_curpos({ 1, 6 })
+        vim.cmd("normal 2ys2wb")
+        check_lines({
+            "some ((more placeholder)) text",
+            "<<<some **more** lines>>>",
+            "hello world",
+        })
+
+        set_curpos({ 3, 1 })
+        vim.cmd("normal 20ysw'")
+        check_lines({
+            "some ((more placeholder)) text",
+            "<<<some **more** lines>>>",
+            "''''''''''''''''''''hello'''''''''''''''''''' world",
+        })
+    end)
+
+    it("can handle number prefixing for visual surround", function()
+        set_lines({
+            "some more placeholder text",
+            "some more lines",
+            "hello world",
+        })
+        set_curpos({ 1, 6 })
+        vim.cmd("normal! v")
+        set_curpos({ 2, 11 })
+        vim.cmd("normal 3Sr")
+        check_lines({
+            "some [[[more placeholder text",
+            "some more l]]]ines",
+            "hello world",
+        })
+
+        set_curpos({ 3, 8 })
+        vim.cmd("normal! v")
+        set_curpos({ 2, 3 })
+        vim.cmd("normal 2Sa")
+        check_lines({
+            "some [[[more placeholder text",
+            "so<<me more l]]]ines",
+            "hello wo>>rld",
+        })
+    end)
+
+    it("can handle number prefixing for visual surround (line mode)", function()
+        set_lines({
+            "some more placeholder text",
+        })
+        set_curpos({ 1, 6 })
+        vim.cmd("normal V3Sa")
+        check_lines({
+            "<",
+            "<",
+            "<",
+            "some more placeholder text",
+            ">",
+            ">",
+            ">",
+        })
+    end)
+
+    it("can handle number prefixing for visual surround (block mode)", function()
+        set_lines({
+            "some more placeholder text",
+            "a short line",
+            "a slightly longer line",
+        })
+        set_curpos({ 1, 6 })
+        vim.cmd("normal " .. ctrl_v .. "jj2w3Sa")
+        check_lines({
+            "some <<<more placehold>>>er text",
+            "a sho<<<rt line>>>",
+            "a sli<<<ghtly longer l>>>ine",
+        })
+    end)
 end)