diff --git a/lua/nvim-surround/cache.lua b/lua/nvim-surround/cache.lua index 7383513..201a8ea 100644 --- a/lua/nvim-surround/cache.lua +++ b/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 = {} diff --git a/lua/nvim-surround/init.lua b/lua/nvim-surround/init.lua index 0732bde..41896e3 100644 --- a/lua/nvim-surround/init.lua +++ b/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 .. "" + 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 diff --git a/lua/nvim-surround/utils.lua b/lua/nvim-surround/utils.lua index e44f920..13da16d 100644 --- a/lua/nvim-surround/utils.lua +++ b/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. diff --git a/tests/basics_spec.lua b/tests/basics_spec.lua index 9d312cf..5201f59 100644 --- a/tests/basics_spec.lua +++ b/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", + "<<>>", + "hello world", + }) + + set_curpos({ 1, 6 }) + vim.cmd("normal 2ys2wb") + check_lines({ + "some ((more placeholder)) text", + "<<>>", + "hello world", + }) + + set_curpos({ 3, 1 }) + vim.cmd("normal 20ysw'") + check_lines({ + "some ((more placeholder)) text", + "<<>>", + "''''''''''''''''''''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<>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 <<>>er text", + "a sho<<>>", + "a sli<<>>ine", + }) + end) end)