Skip to content

Commit b36bf49

Browse files
authored
refactor: Extmark management. (#339)
* refactor: Simplify `visual_surround`. * refactor: Use `with_extmark` to prevent leaking extmarks.
1 parent 687ea2f commit b36bf49

File tree

2 files changed

+74
-73
lines changed

2 files changed

+74
-73
lines changed

lua/nvim-surround/buffer.lua

+13
Original file line numberDiff line numberDiff line change
@@ -146,6 +146,19 @@ M.del_extmark = function(extmark)
146146
vim.api.nvim_buf_del_extmark(0, M.namespace.extmark, extmark)
147147
end
148148

149+
-- Runs a callback function with an extmark.
150+
---@param pos position The initial position of the extmark.
151+
---@param callback fun(): nil
152+
---@return position
153+
---@nodiscard
154+
M.with_extmark = function(pos, callback)
155+
local extmark = M.set_extmark(pos)
156+
callback()
157+
local extmark_pos = M.get_extmark(extmark)
158+
M.del_extmark(extmark)
159+
return extmark_pos
160+
end
161+
149162
--[====================================================================================================================[
150163
Byte indexing helper functions
151164
--]====================================================================================================================]

lua/nvim-surround/init.lua

+61-73
Original file line numberDiff line numberDiff line change
@@ -64,16 +64,15 @@ M.normal_surround = function(args)
6464
local first_pos = args.selection.first_pos
6565
local last_pos = { args.selection.last_pos[1], args.selection.last_pos[2] + 1 }
6666

67-
local sticky_mark = buffer.set_extmark(M.normal_curpos)
68-
buffer.insert_text(last_pos, args.delimiters[2])
69-
buffer.insert_text(first_pos, args.delimiters[1])
70-
67+
local sticky_pos = buffer.with_extmark(M.normal_curpos, function()
68+
buffer.insert_text(last_pos, args.delimiters[2])
69+
buffer.insert_text(first_pos, args.delimiters[1])
70+
end)
7171
buffer.restore_curpos({
7272
first_pos = first_pos,
73-
sticky_pos = buffer.get_extmark(sticky_mark),
73+
sticky_pos = sticky_pos,
7474
old_pos = M.normal_curpos,
7575
})
76-
buffer.del_extmark(sticky_mark)
7776

7877
if args.line_mode then
7978
config.get_opts().indent_lines(first_pos[1], last_pos[1] + #args.delimiters[1] + #args.delimiters[2] - 2)
@@ -84,7 +83,6 @@ end
8483
-- Add delimiters around a visual selection.
8584
---@param args { line_mode: boolean, curpos: position, curswant: number }
8685
M.visual_surround = function(args)
87-
-- Get a character and selection from the user
8886
local ins_char = input.get_char()
8987

9088
if vim.fn.visualmode() == "V" then
@@ -96,63 +94,59 @@ M.visual_surround = function(args)
9694
return
9795
end
9896

99-
local sticky_mark = buffer.set_extmark(args.curpos)
100-
if vim.fn.visualmode() == "\22" then -- Visual block mode case (add delimiters to every line)
101-
if vim.o.selection == "exclusive" then
102-
last_pos[2] = last_pos[2] - 1
103-
end
104-
-- Get (visually) what columns the start and end are located at
105-
local first_disp = vim.fn.strdisplaywidth(buffer.get_line(first_pos[1]):sub(1, first_pos[2] - 1)) + 1
106-
local last_disp = vim.fn.strdisplaywidth(buffer.get_line(last_pos[1]):sub(1, last_pos[2] - 1)) + 1
107-
-- Find the min/max for some variables, since visual blocks can either go diagonally or anti-diagonally
108-
local mn_disp, mx_disp = math.min(first_disp, last_disp), math.max(first_disp, last_disp)
109-
local mn_lnum, mx_lnum = math.min(first_pos[1], last_pos[1]), math.max(first_pos[1], last_pos[1])
110-
-- Check if $ was used in creating the block selection
111-
local surround_to_end_of_line = args.curswant == vim.v.maxcol
112-
-- Surround each line with the delimiter pair, last to first (for indexing reasons)
113-
for lnum = mx_lnum, mn_lnum, -1 do
114-
local line = buffer.get_line(lnum)
115-
if surround_to_end_of_line then
116-
buffer.insert_text({ lnum, #buffer.get_line(lnum) + 1 }, delimiters[2])
117-
else
118-
local index = buffer.get_last_byte({ lnum, 1 })[2]
119-
-- The current display count should be >= the desired one
120-
while vim.fn.strdisplaywidth(line:sub(1, index)) < mx_disp and index <= #line do
121-
index = buffer.get_last_byte({ lnum, index + 1 })[2]
97+
if vim.o.selection == "exclusive" then
98+
last_pos[2] = last_pos[2] - 1
99+
end
100+
local sticky_pos = buffer.with_extmark(args.curpos, function()
101+
if vim.fn.visualmode() == "\22" then -- Visual block mode case (add delimiters to every line)
102+
-- Get (visually) what columns the start and end are located at
103+
local first_disp = vim.fn.strdisplaywidth(buffer.get_line(first_pos[1]):sub(1, first_pos[2] - 1)) + 1
104+
local last_disp = vim.fn.strdisplaywidth(buffer.get_line(last_pos[1]):sub(1, last_pos[2] - 1)) + 1
105+
-- Find the min/max for some variables, since visual blocks can either go diagonally or anti-diagonally
106+
local mn_disp, mx_disp = math.min(first_disp, last_disp), math.max(first_disp, last_disp)
107+
local mn_lnum, mx_lnum = math.min(first_pos[1], last_pos[1]), math.max(first_pos[1], last_pos[1])
108+
-- Check if $ was used in creating the block selection
109+
local surround_to_end_of_line = args.curswant == vim.v.maxcol
110+
-- Surround each line with the delimiter pair, last to first (for indexing reasons)
111+
for lnum = mx_lnum, mn_lnum, -1 do
112+
local line = buffer.get_line(lnum)
113+
if surround_to_end_of_line then
114+
buffer.insert_text({ lnum, #buffer.get_line(lnum) + 1 }, delimiters[2])
115+
else
116+
local index = buffer.get_last_byte({ lnum, 1 })[2]
117+
-- The current display count should be >= the desired one
118+
while vim.fn.strdisplaywidth(line:sub(1, index)) < mx_disp and index <= #line do
119+
index = buffer.get_last_byte({ lnum, index + 1 })[2]
120+
end
121+
-- Go to the end of the current character
122+
index = buffer.get_last_byte({ lnum, index })[2]
123+
buffer.insert_text({ lnum, index + 1 }, delimiters[2])
122124
end
123-
-- Go to the end of the current character
124-
index = buffer.get_last_byte({ lnum, index })[2]
125-
buffer.insert_text({ lnum, index + 1 }, delimiters[2])
126-
end
127125

128-
local index = 1
129-
-- The current display count should be <= the desired one
130-
while vim.fn.strdisplaywidth(line:sub(1, index - 1)) + 1 < mn_disp and index <= #line do
131-
index = buffer.get_last_byte({ lnum, index })[2] + 1
132-
end
133-
if vim.fn.strdisplaywidth(line:sub(1, index - 1)) + 1 > mn_disp then
134-
-- Go to the beginning of the previous character
135-
index = buffer.get_first_byte({ lnum, index - 1 })[2]
126+
local index = 1
127+
-- The current display count should be <= the desired one
128+
while vim.fn.strdisplaywidth(line:sub(1, index - 1)) + 1 < mn_disp and index <= #line do
129+
index = buffer.get_last_byte({ lnum, index })[2] + 1
130+
end
131+
if vim.fn.strdisplaywidth(line:sub(1, index - 1)) + 1 > mn_disp then
132+
-- Go to the beginning of the previous character
133+
index = buffer.get_first_byte({ lnum, index - 1 })[2]
134+
end
135+
buffer.insert_text({ lnum, index }, delimiters[1])
136136
end
137-
buffer.insert_text({ lnum, index }, delimiters[1])
137+
else -- Regular visual mode case
138+
last_pos = buffer.get_last_byte(last_pos)
139+
buffer.insert_text({ last_pos[1], last_pos[2] + 1 }, delimiters[2])
140+
buffer.insert_text(first_pos, delimiters[1])
138141
end
139-
else -- Regular visual mode case
140-
if vim.o.selection == "exclusive" then
141-
last_pos[2] = last_pos[2] - 1
142-
end
143-
144-
last_pos = buffer.get_last_byte(last_pos)
145-
buffer.insert_text({ last_pos[1], last_pos[2] + 1 }, delimiters[2])
146-
buffer.insert_text(first_pos, delimiters[1])
147-
end
142+
end)
148143

149144
config.get_opts().indent_lines(first_pos[1], last_pos[1] + #delimiters[1] + #delimiters[2] - 2)
150145
buffer.restore_curpos({
151146
first_pos = first_pos,
152-
sticky_pos = buffer.get_extmark(sticky_mark),
147+
sticky_pos = sticky_pos,
153148
old_pos = args.curpos,
154149
})
155-
buffer.del_extmark(sticky_mark)
156150
end
157151

158152
-- Delete a surrounding delimiter pair, if it exists.
@@ -168,25 +162,21 @@ M.delete_surround = function(args)
168162
return "g@l"
169163
end
170164

171-
-- Get the selections to delete
172165
local selections = utils.get_nearest_selections(args.del_char, "delete")
173-
174166
if selections then
175-
local sticky_mark = buffer.set_extmark(args.curpos)
176-
-- Delete the right selection first to ensure selection positions are correct
177-
buffer.delete_selection(selections.right)
178-
buffer.delete_selection(selections.left)
179-
167+
local sticky_pos = buffer.with_extmark(args.curpos, function()
168+
buffer.delete_selection(selections.right)
169+
buffer.delete_selection(selections.left)
170+
end)
180171
config.get_opts().indent_lines(
181172
selections.left.first_pos[1],
182173
selections.left.first_pos[1] + selections.right.first_pos[1] - selections.left.last_pos[1]
183174
)
184175
buffer.restore_curpos({
185176
first_pos = selections.left.first_pos,
186-
sticky_pos = buffer.get_extmark(sticky_mark),
177+
sticky_pos = sticky_pos,
187178
old_pos = args.curpos,
188179
})
189-
buffer.del_extmark(sticky_mark)
190180
end
191181

192182
cache.set_callback("v:lua.require'nvim-surround'.delete_callback")
@@ -232,22 +222,20 @@ M.change_surround = function(args)
232222
selections.right.first_pos[2] = space_end + 1
233223
end
234224

235-
local sticky_mark = buffer.set_extmark(args.curpos)
236-
-- Change the right selection first to ensure selection positions are correct
237-
buffer.change_selection(selections.right, delimiters[2])
238-
buffer.change_selection(selections.left, delimiters[1])
225+
local sticky_pos = buffer.with_extmark(args.curpos, function()
226+
buffer.change_selection(selections.right, delimiters[2])
227+
buffer.change_selection(selections.left, delimiters[1])
228+
end)
239229
buffer.restore_curpos({
240230
first_pos = selections.left.first_pos,
241-
sticky_pos = buffer.get_extmark(sticky_mark),
231+
sticky_pos = sticky_pos,
242232
old_pos = args.curpos,
243233
})
244-
buffer.del_extmark(sticky_mark)
245234

246235
if args.line_mode then
247-
local first_pos = selections.left.first_pos
248-
local last_pos = selections.right.last_pos
249-
250-
config.get_opts().indent_lines(first_pos[1], last_pos[1] + #delimiters[1] + #delimiters[2] - 2)
236+
local first_line = selections.left.first_pos[1]
237+
local last_line = selections.right.last_pos[1]
238+
config.get_opts().indent_lines(first_line, last_line + #delimiters[1] + #delimiters[2] - 2)
251239
end
252240
end
253241

0 commit comments

Comments
 (0)