feat(snippet): add default keymaps during snippet session

3 files changed, 91 insertions, 5 deletions

diff --git a/runtime/doc/news.txt b/runtime/doc/news.txt
index 439316f62e..f1b402c92e 100644
--- a/runtime/doc/news.txt
+++ b/runtime/doc/news.txt
@@ -85,6 +85,12 @@ DEFAULTS
   - |gra| in Normal and Visual mode maps to |vim.lsp.buf.code_action()|
   - CTRL-S in Insert mode maps to |vim.lsp.buf.signature_help()|
 
+• Snippet:
+  - `<Tab>` in Insert and Select mode maps to |vim.snippet.jump({ direction = 1 })|
+    when a snippet is active and jumpable forwards.
+  - `<S-Tab>` in Insert and Select mode maps to |vim.snippet.jump({ direction = -1 })|
+    when a snippet is active and jumpable backwards.
+
 EDITOR
 
 • TODO
diff --git a/runtime/lua/vim/snippet.lua b/runtime/lua/vim/snippet.lua
index 3d8f73f362..1ec5235d7b 100644
--- a/runtime/lua/vim/snippet.lua
+++ b/runtime/lua/vim/snippet.lua
@@ -2,6 +2,8 @@ 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'
+local jump_forward_key = '<tab>'
+local jump_backward_key = '<s-tab>'
 
 --- Returns the 0-based cursor position.
 ---
@@ -182,6 +184,8 @@ end
 --- @field extmark_id integer
 --- @field tabstops table<integer, vim.snippet.Tabstop[]>
 --- @field current_tabstop vim.snippet.Tabstop
+--- @field tab_keymaps { i: table<string, any>?, s: table<string, any>? }
+--- @field shift_tab_keymaps { i: table<string, any>?, s: table<string, any>? }
 local Session = {}
 
 --- Creates a new snippet session in the current buffer.
@@ -197,6 +201,8 @@ function Session.new(bufnr, snippet_extmark, tabstop_data)
     extmark_id = snippet_extmark,
     tabstops = {},
     current_tabstop = Tabstop.new(0, bufnr, { 0, 0, 0, 0 }),
+    tab_keymaps = { i = nil, s = nil },
+    shift_tab_keymaps = { i = nil, s = nil },
   }, { __index = Session })
 
   -- Create the tabstops.
@@ -207,9 +213,64 @@ function Session.new(bufnr, snippet_extmark, tabstop_data)
     end
   end
 
+  self:set_keymaps()
+
   return self
 end
 
+--- Sets the snippet navigation keymaps.
+---
+--- @package
+function Session:set_keymaps()
+  local function maparg(key, mode)
+    local map = vim.fn.maparg(key, mode, false, true) --[[ @as table ]]
+    if not vim.tbl_isempty(map) and map.buffer == 1 then
+      return map
+    else
+      return nil
+    end
+  end
+
+  local function set(jump_key, direction)
+    vim.keymap.set({ 'i', 's' }, jump_key, function()
+      return vim.snippet.active({ direction = direction })
+          and '<cmd>lua vim.snippet.jump(' .. direction .. ')<cr>'
+        or jump_key
+    end, { expr = true, silent = true, buffer = self.bufnr })
+  end
+
+  self.tab_keymaps = {
+    i = maparg(jump_forward_key, 'i'),
+    s = maparg(jump_forward_key, 's'),
+  }
+  self.shift_tab_keymaps = {
+    i = maparg(jump_backward_key, 'i'),
+    s = maparg(jump_backward_key, 's'),
+  }
+  set(jump_forward_key, 1)
+  set(jump_backward_key, -1)
+end
+
+--- Restores/deletes the keymaps used for snippet navigation.
+---
+--- @package
+function Session:restore_keymaps()
+  local function restore(keymap, lhs, mode)
+    if keymap then
+      vim.api.nvim_buf_call(self.bufnr, function()
+        vim.fn.mapset(keymap)
+      end)
+    else
+      vim.api.nvim_buf_del_keymap(self.bufnr, mode, lhs)
+    end
+  end
+
+  restore(self.tab_keymaps.i, jump_forward_key, 'i')
+  restore(self.tab_keymaps.s, jump_forward_key, 's')
+  restore(self.shift_tab_keymaps.i, jump_backward_key, 'i')
+  restore(self.shift_tab_keymaps.s, jump_backward_key, 's')
+end
+
 --- Returns the destination tabstop index when jumping in the given direction.
 ---
 --- @package
@@ -619,6 +680,8 @@ function M.stop()
     return
   end
 
+  M._session:restore_keymaps()
+
   vim.api.nvim_clear_autocmds({ group = snippet_group, buffer = M._session.bufnr })
   vim.api.nvim_buf_clear_namespace(M._session.bufnr, snippet_ns, 0, -1)
 
diff --git a/test/functional/lua/snippet_spec.lua b/test/functional/lua/snippet_spec.lua
index 413aa93994..bca0a59cb4 100644
--- a/test/functional/lua/snippet_spec.lua
+++ b/test/functional/lua/snippet_spec.lua
@@ -1,3 +1,5 @@
+---@diagnostic disable: no-unknown
+
 local t = require('test.testutil')
 local n = require('test.functional.testnvim')()
 
@@ -16,11 +18,6 @@ local retry = t.retry
 describe('vim.snippet', function()
   before_each(function()
     clear()
-
-    exec_lua([[
-      vim.keymap.set({ 'i', 's' }, '<Tab>', function() vim.snippet.jump(1) end, { buffer = true })
-      vim.keymap.set({ 'i', 's' }, '<S-Tab>', function() vim.snippet.jump(-1) end, { buffer = true })
-    ]])
   end)
   after_each(clear)
 
@@ -286,4 +283,24 @@ describe('vim.snippet', function()
     ]]
     )
   end)
+
+  it('restores snippet navigation keymaps', function()
+    -- Create a buffer keymap in insert mode that deletes all lines.
+    local curbuf = api.nvim_get_current_buf()
+    exec_lua('vim.api.nvim_buf_set_keymap(..., "i", "<Tab>", "<cmd>normal ggdG<cr>", {})', curbuf)
+
+    test_expand_success({ 'var $1 = $2' }, { 'var  = ' })
+
+    -- While the snippet is active, <Tab> should navigate between tabstops.
+    feed('x')
+    poke_eventloop()
+    feed('<Tab>0')
+    eq({ 'var x = 0' }, buf_lines(0))
+
+    exec_lua('vim.snippet.stop()')
+
+    -- After exiting the snippet, the buffer keymap should be restored.
+    feed('<Esc>O<cr><Tab>')
+    eq({ '' }, buf_lines(0))
+  end)
 end)