change_choice: use cursor-restore system from update_dependents.

Code is essentially the same thing. Also allow passing the current cursor to ls.set/change_choice, which is useful for extras.select_choice, where the cursor-state is "destroyed" due to vim.input, and should be saved before that is even opened.
L3MON4D3 · Nov 7, 2024 · 1464b3f · 1464b3f
1 parent 14280e1
commit 1464b3f
Show file tree

Hide file tree

Showing 6 changed files with 142 additions and 95 deletions.
diff --git a/lua/luasnip/extras/select_choice.lua b/lua/luasnip/extras/select_choice.lua
@@ -1,24 +1,30 @@
 local session = require("luasnip.session")
 local ls = require("luasnip")
+local node_util = require("luasnip.nodes.util")
 
-local function set_choice_callback(_, indx)
-	if not indx then
-		return
+local function set_choice_callback(data)
+	return function(_, indx)
+		if not indx then
+			return
+		end
+		-- feed+immediately execute i to enter INSERT after vim.ui.input closes.
+		-- vim.api.nvim_feedkeys("i", "x", false)
+		ls.set_choice(indx, {cursor_restore_data = data})
 	end
-	-- feed+immediately execute i to enter INSERT after vim.ui.input closes.
-	vim.api.nvim_feedkeys("i", "x", false)
-	ls.set_choice(indx)
 end
 
 local function select_choice()
 	assert(
 		session.active_choice_nodes[vim.api.nvim_get_current_buf()],
 		"No active choiceNode"
 	)
+	local active = session.current_nodes[vim.api.nvim_get_current_buf()]
+
+	local restore_data = node_util.store_cursor_node_relative(active)
 	vim.ui.select(
 		ls.get_current_choices(),
 		{ kind = "luasnip" },
-		set_choice_callback
+		set_choice_callback(restore_data)
 	)
 end
 

diff --git a/lua/luasnip/init.lua b/lua/luasnip/init.lua
@@ -149,63 +149,15 @@ local function unlink_current()
 	unlink_set_adjacent_as_current_no_log(current.parent.snippet)
 end
 
-local store_id = 0
-local function store_cursor_node_relative(node)
-	local data = {}
-
-	local snippet_current_node = node
-
-	-- store for each snippet!
-	-- the innermost snippet may be destroyed, and we would have to restore the
-	-- cursor in a snippet above that.
-	while snippet_current_node do
-		local snip = snippet_current_node:get_snippet()
-
-		local snip_data = {}
-
-		snip_data.key = node.key
-		node.store_id = store_id
-		snip_data.store_id = store_id
-		snip_data.node = snippet_current_node
-
-		store_id = store_id + 1
-
-		local cursor_state = feedkeys.last_state()
-		snip_data.cursor_end_relative =
-			util.pos_sub(cursor_state.pos, node.mark:get_endpoint(1))
-
-		if cursor_state.pos_end then
-			snip_data.selection_other_end_end_relative = util.pos_sub(cursor_state.pos_end, node.mark:get_endpoint(1))
-		end
-
-		data[snip] = snip_data
-
-		snippet_current_node = snip:get_snippet().parent_node
-	end
-
-	return data
-end
-
 local function get_corresponding_node(parent, data)
 	return parent:find_node(function(test_node)
 		return (test_node.store_id == data.store_id)
 			or (data.key ~= nil and test_node.key == data.key)
 	end, {find_in_child_snippets = true})
 end
 
-local function restore_cursor_pos_relative(node, data)
-	if data.selection_other_end_end_relative then
-		-- is a selection => restore it.
-		local selection_from = util.pos_add(node.mark:get_endpoint(1), data.cursor_end_relative)
-		local selection_to = util.pos_add(node.mark:get_endpoint(1), data.selection_other_end_end_relative)
-		feedkeys.select_range(selection_from, selection_to)
-	else
-		feedkeys.move_to(util.pos_add(node.mark:get_endpoint(1), data.cursor_end_relative))
-	end
-end
-
 local function node_update_dependents_preserve_position(node, current, opts)
-	local restore_data = store_cursor_node_relative(current)
+	local restore_data = node_util.store_cursor_node_relative(current)
 
 	-- update all nodes that depend on this one.
 	local ok, res =
@@ -230,7 +182,7 @@ local function node_update_dependents_preserve_position(node, current, opts)
 		if not opts.no_move and opts.restore_position then
 			-- node is visible: restore position.
 			local active_snippet = current:get_snippet()
-			restore_cursor_pos_relative(current, restore_data[active_snippet])
+			node_util.restore_cursor_pos_relative(current, restore_data[active_snippet])
 		end
 
 		return { jump_done = false, new_current = current }
@@ -275,7 +227,7 @@ local function node_update_dependents_preserve_position(node, current, opts)
 
 			if not opts.no_move and opts.restore_position then
 				-- node is visible: restore position
-				restore_cursor_pos_relative(new_current, snip_restore_data)
+				node_util.restore_cursor_pos_relative(new_current, snip_restore_data)
 			end
 
 			return { jump_done = false, new_current = new_current }
@@ -621,23 +573,33 @@ local function safe_choice_action(snip, ...)
 		return session.current_nodes[vim.api.nvim_get_current_buf()]
 	end
 end
-local function change_choice(val)
+local function change_choice(val, opts)
 	local active_choice =
 		session.active_choice_nodes[vim.api.nvim_get_current_buf()]
+
 	assert(active_choice, "No active choiceNode")
+
+	-- if the active choice exists current_node still does.
+	local current_node = session.current_nodes[vim.api.nvim_get_current_buf()]
+	local restore_data = opts and opts.cursor_restore_data or node_util.store_cursor_node_relative(current_node)
+
 	local new_active = no_region_check_wrap(
 		safe_choice_action,
 		active_choice.parent.snippet,
 		active_choice.change_choice,
 		active_choice,
 		val,
-		session.current_nodes[vim.api.nvim_get_current_buf()]
+		session.current_nodes[vim.api.nvim_get_current_buf()],
+		restore_data
 	)
 	session.current_nodes[vim.api.nvim_get_current_buf()] = new_active
 	active_update_dependents()
 end
 
-local function set_choice(choice_indx)
+local function set_choice(choice_indx, opts)
+	local current_node = session.current_nodes[vim.api.nvim_get_current_buf()]
+	local restore_data = opts and opts.cursor_restore_data or node_util.store_cursor_node_relative(current_node)
+
 	local active_choice =
 		session.active_choice_nodes[vim.api.nvim_get_current_buf()]
 	assert(active_choice, "No active choiceNode")
@@ -649,7 +611,8 @@ local function set_choice(choice_indx)
 		active_choice.set_choice,
 		active_choice,
 		choice,
-		session.current_nodes[vim.api.nvim_get_current_buf()]
+		current_node,
+		restore_data
 	)
 	session.current_nodes[vim.api.nvim_get_current_buf()] = new_active
 	active_update_dependents()

diff --git a/lua/luasnip/nodes/choiceNode.lua b/lua/luasnip/nodes/choiceNode.lua
@@ -227,17 +227,12 @@ end
 -- used to uniquely identify this change-choice-action.
 local change_choice_id = 0
 
-function ChoiceNode:set_choice(choice, current_node)
+function ChoiceNode:set_choice(choice, current_node, cursor_restore_data)
 	change_choice_id = change_choice_id + 1
 	-- to uniquely identify this node later (storing the pointer isn't enough
 	-- because this is supposed to work with restoreNodes, which are copied).
 	current_node.change_choice_id = change_choice_id
 
-	local insert_pre_cc = vim.fn.mode() == "i"
-	-- is byte-indexed! Doesn't matter here, but important to be aware of.
-	local cursor_node_offset =
-		util.pos_offset(current_node.mark:pos_begin(), util.get_cursor_0ind())
-
 	self.active_choice:store()
 
 	-- tear down current choice.
@@ -289,30 +284,22 @@ function ChoiceNode:set_choice(choice, current_node)
 			-- and this choiceNode, then set the cursor.
 
 			node_util.refocus(self, target_node)
+			node_util.restore_cursor_pos_relative(target_node, cursor_restore_data[target_node.parent.snippet])
 
-			if insert_pre_cc then
-				feedkeys.move_to(
-					util.pos_from_offset(
-						target_node.mark:pos_begin(),
-						cursor_node_offset
-					)
-				)
-			else
-				node_util.select_node(target_node)
-			end
 			return target_node
 		end
 	end
 
 	return self.active_choice:jump_into(1)
 end
 
-function ChoiceNode:change_choice(dir, current_node)
+function ChoiceNode:change_choice(dir, current_node, cursor_restore_data)
 	-- stylua: ignore
 	return self:set_choice(
 		dir == 1 and self.active_choice.next_choice
 		          or self.active_choice.prev_choice,
-		current_node )
+		current_node,
+		cursor_restore_data)
 end
 
 function ChoiceNode:copy()

diff --git a/lua/luasnip/nodes/util.lua b/lua/luasnip/nodes/util.lua
@@ -862,6 +862,59 @@ local function str_args(args)
 	end, args)
 end
 
+local store_id = 0
+local function store_cursor_node_relative(node)
+	local data = {}
+
+	local snippet_current_node = node
+
+	local cursor_state = feedkeys.last_state()
+
+	-- store for each snippet!
+	-- the innermost snippet may be destroyed, and we would have to restore the
+	-- cursor in a snippet above that.
+	while snippet_current_node do
+		local snip = snippet_current_node:get_snippet()
+
+		local snip_data = {}
+
+		snip_data.key = node.key
+		node.store_id = store_id
+		snip_data.store_id = store_id
+		snip_data.node = snippet_current_node
+
+		store_id = store_id + 1
+
+		snip_data.cursor_start_relative =
+			util.pos_offset(node.mark:get_endpoint(-1), cursor_state.pos)
+
+		snip_data.mode = cursor_state.mode
+
+		if cursor_state.pos_v then
+			snip_data.selection_end_start_relative = util.pos_offset(node.mark:get_endpoint(-1), cursor_state.pos_v)
+		end
+
+		data[snip] = snip_data
+
+		snippet_current_node = snip:get_snippet().parent_node
+	end
+
+	return data
+end
+
+local function restore_cursor_pos_relative(node, data)
+	if data.mode == "i" then
+		feedkeys.insert_at(util.pos_from_offset(node.mark:get_endpoint(-1), data.cursor_start_relative))
+	elseif data.mode == "s" then
+		-- is a selection => restore it.
+		local selection_from = util.pos_from_offset(node.mark:get_endpoint(-1), data.cursor_start_relative)
+		local selection_to = util.pos_from_offset(node.mark:get_endpoint(-1), data.selection_end_start_relative)
+		feedkeys.select_range(selection_from, selection_to)
+	else
+		feedkeys.move_to_normal(util.pos_from_offset(node.mark:get_endpoint(-1), data.cursor_start_relative))
+	end
+end
+
 return {
 	subsnip_init_children = subsnip_init_children,
 	init_child_positions_func = init_child_positions_func,
@@ -887,5 +940,7 @@ return {
 	find_node_dependents = find_node_dependents,
 	collect_dependents = collect_dependents,
 	node_subtree_do = node_subtree_do,
-	str_args = str_args
+	str_args = str_args,
+	store_cursor_node_relative = store_cursor_node_relative,
+	restore_cursor_pos_relative = restore_cursor_pos_relative
 }
diff --git a/lua/luasnip/util/feedkeys.lua b/lua/luasnip/util/feedkeys.lua
@@ -99,7 +99,7 @@ end
 
 function M.select_range(b, e)
 	local id = next_id()
-	enqueued_cursor_state = {pos = vim.deepcopy(b), pos_end = vim.deepcopy(e), id = id}
+	enqueued_cursor_state = {pos = vim.deepcopy(b), pos_v = vim.deepcopy(e), mode = "s", id = id}
 	enqueue_action(function()
 		-- stylua: ignore
 		_feedkeys_insert(id,
@@ -132,7 +132,7 @@ end
 -- move the cursor to a position and enter insert-mode (or stay in it).
 function M.insert_at(pos)
 	local id = next_id()
-	enqueued_cursor_state = {pos = pos, id = id}
+	enqueued_cursor_state = {pos = pos, mode = "i", id = id}
 
 	enqueue_action(function()
 		-- if current and target mode is INSERT, there's no reason to leave it.
@@ -152,13 +152,18 @@ function M.insert_at(pos)
 end
 
 -- move, without changing mode.
-function M.move_to(pos)
+function M.move_to_normal(pos)
 	local id = next_id()
-	enqueued_cursor_state = {pos = pos, id = id}
+	-- preserve mode.
+	enqueued_cursor_state = {pos = pos, mode = "n", id = id}
 
 	enqueue_action(function()
-		util.set_cursor_0ind(pos)
-		M.confirm(id)
+		if vim.fn.mode():sub(1,1) == "n" then
+			util.set_cursor_0ind(pos)
+			M.confirm(id)
+		else
+			_feedkeys_insert(id, "<Esc>" .. cursor_set_keys(pos))
+		end
 	end, id)
 end
 
@@ -181,6 +186,7 @@ end
 function M.last_state()
 	if enqueued_cursor_state then
 		local state = vim.deepcopy(enqueued_cursor_state)
+		-- remove internal data.
 		state.id = nil
 		return state
 	end
@@ -190,14 +196,14 @@ function M.last_state()
 	local getposdot = vim.fn.getpos(".")
 	state.pos = {getposdot[2]-1, getposdot[3]-1}
 
-	-- only re-enter select for now.
-	if vim.fn.mode() == "s" then
-		local getposv = vim.fn.getpos("v")
-		-- store selection-range with end-position one column after the cursor
-		-- at the end (so -1 to make getpos-position 0-based, +1 to move it one
-		-- beyond the last character of the range)
-		state.pos_end = {getposv[2]-1, getposv[3]+1}
-	end
+	local getposv = vim.fn.getpos("v")
+	-- store selection-range with end-position one column after the cursor
+	-- at the end (so -1 to make getpos-position 0-based, +1 to move it one
+	-- beyond the last character of the range)
+	state.pos_v = {getposv[2]-1, getposv[3]}
+
+	-- only store first component.
+	state.mode = vim.fn.mode():sub(1,1)
 
 	return state
 end

diff --git a/tests/integration/choice_spec.lua b/tests/integration/choice_spec.lua
@@ -501,6 +501,36 @@ screen:expect({
     end                                               |
     {2:-- INSERT --}                                      |
   ]]
+})
+	end)
+
+	it("select_choice works.", function()
+		exec_lua[=[
+			ls.snip_expand(s("for", {
+				t"for ", c(1, {
+					sn(nil, {i(1, "k"), t", ", i(2, "v"), t" in ", c(3, {{t"pairs(",i(1),t")"}, {t"ipairs(",i(1),t")"}, i(nil)}, {restore_cursor = true}) }),
+					sn(nil, {i(1, "val"), t" in ", i(2) }),
+					sn(nil, {i(1, "i"), t" = ", i(2), t", ", i(3) }),
+					fmt([[{} in vim.gsplit({})]], {i(1, "str"), i(2)})
+				}, {restore_cursor = true}), t{" do", "\t"}, isn(2, {dl(1, l.LS_SELECT_DEDENT)}, "$PARENT_INDENT\t"), t{"", "end"}
+			}))
+		]=]
+		feed("<cmd>lua require('luasnip.extras.select_choice')()<Cr>2<Cr>")
+screen:expect({
+  grid = [[
+    for ^v{3:al} in  do                                    |
+                                                      |
+    {2:-- SELECT --}                                      |
+  ]]
+})
+		-- re-selecting correctly highlights text again (test by editing so the test does not pass immediately, without any changes!)
+		feed("<cmd>lua require('luasnip.extras.select_choice')()<Cr>2<Cr>val")
+screen:expect({
+  grid = [[
+    for val^ in  do                                    |
+                                                      |
+    {2:-- INSERT --}                                      |
+  ]]
 })
 	end)
 end)