fix(lsp): handle adjacent snippet tabstops

1 files changed, 35 insertions, 2 deletions

diff --git a/runtime/lua/vim/snippet.lua b/runtime/lua/vim/snippet.lua
index 4c62b5c076..7e37d84393 100644
--- a/runtime/lua/vim/snippet.lua
+++ b/runtime/lua/vim/snippet.lua
@@ -1,6 +1,7 @@
 local G = vim.lsp._snippet_grammar
 local snippet_group = vim.api.nvim_create_augroup('vim/snippet', {})
 local snippet_ns = vim.api.nvim_create_namespace('vim/snippet')
+local hl_group = 'SnippetTabstop'
 
 --- Returns the 0-based cursor position.
 ---
@@ -117,11 +118,11 @@ local Tabstop = {}
 --- @return vim.snippet.Tabstop
 function Tabstop.new(index, bufnr, range, choices)
   local extmark_id = vim.api.nvim_buf_set_extmark(bufnr, snippet_ns, range[1], range[2], {
-    right_gravity = false,
+    right_gravity = true,
     end_right_gravity = true,
     end_line = range[3],
     end_col = range[4],
-    hl_group = 'SnippetTabstop',
+    hl_group = hl_group,
   })
 
   local self = setmetatable(
@@ -161,6 +162,21 @@ function Tabstop:set_text(text)
   vim.api.nvim_buf_set_text(self.bufnr, range[1], range[2], range[3], range[4], text_to_lines(text))
 end
 
+--- Sets the right gravity of the tabstop's extmark.
+---
+--- @package
+--- @param right_gravity boolean
+function Tabstop:set_right_gravity(right_gravity)
+  local range = self:get_range()
+  self.extmark_id = vim.api.nvim_buf_set_extmark(self.bufnr, snippet_ns, range[1], range[2], {
+    right_gravity = right_gravity,
+    end_right_gravity = true,
+    end_line = range[3],
+    end_col = range[4],
+    hl_group = hl_group,
+  })
+end
+
 --- @class vim.snippet.Session
 --- @field bufnr integer
 --- @field extmark_id integer
@@ -218,6 +234,17 @@ function Session:get_dest_index(direction)
   end
 end
 
+--- Sets the right gravity of the tabstop group with the given index.
+---
+--- @package
+--- @param index integer
+--- @param right_gravity boolean
+function Session:set_group_gravity(index, right_gravity)
+  for _, tabstop in ipairs(self.tabstops[index]) do
+    tabstop:set_right_gravity(right_gravity)
+  end
+end
+
 --- @class vim.snippet.Snippet
 --- @field private _session? vim.snippet.Session
 local M = { session = nil }
@@ -560,9 +587,15 @@ function M.jump(direction)
   -- Clear the autocommands so that we can move the cursor freely while selecting the tabstop.
   vim.api.nvim_clear_autocmds({ group = snippet_group, buffer = M._session.bufnr })
 
+  -- Deactivate expansion of the current tabstop.
+  M._session:set_group_gravity(M._session.current_tabstop.index, true)
+
   M._session.current_tabstop = dest
   select_tabstop(dest)
 
+  -- Activate expansion of the destination tabstop.
+  M._session:set_group_gravity(dest.index, false)
+
   -- Restore the autocommands.
   setup_autocmds(M._session.bufnr)
 end