fix(lsp): support v:count in selection_range() #34551

Co-authored-by: Yi Ming <ofseed@foxmail.com>

runtime/doc/lsp.txt

@@ -1839,11 +1839,11 @@ rename({new_name}, {opts})                              *vim.lsp.buf.rename()*
 
 selection_range({direction})                   *vim.lsp.buf.selection_range()*
     Perform an incremental selection at the cursor position based on ranges
-    given by the LSP. The `direction` parameter specifies whether the
+    given by the LSP. The `direction` parameter specifies the number of times
-    selection should head inward or outward.
+    to expand the selection. Negative values will shrink the selection.
 
     Parameters: ~
-      • {direction}  (`'inner'|'outer'`)
+      • {direction}  (`integer`)
 
 signature_help({config})                        *vim.lsp.buf.signature_help()*
     Displays signature information about the symbol under the cursor in a

runtime/doc/news.txt

@@ -24,7 +24,8 @@ EXPERIMENTS
 
 LSP
 
-• todo
+• |vim.lsp.buf.selection_range()| now accepts an integer which specifies its
+  direction, rather than a string `'outer'` or `'inner'`.
 
 OPTIONS
 

runtime/lua/vim/_defaults.lua

@@ -210,12 +210,12 @@ do
     end, { desc = 'vim.lsp.buf.implementation()' })
 
     vim.keymap.set('x', 'an', function()
-      vim.lsp.buf.selection_range('outer')
+      vim.lsp.buf.selection_range(vim.v.count1)
-    end, { desc = "vim.lsp.buf.selection_range('outer')" })
+    end, { desc = 'vim.lsp.buf.selection_range(vim.v.count1)' })
 
     vim.keymap.set('x', 'in', function()
-      vim.lsp.buf.selection_range('inner')
+      vim.lsp.buf.selection_range(-vim.v.count1)
-    end, { desc = "vim.lsp.buf.selection_range('inner')" })
+    end, { desc = 'vim.lsp.buf.selection_range(-vim.v.count1)' })
 
     vim.keymap.set('n', 'gO', function()
       vim.lsp.buf.document_symbol()

runtime/lua/vim/lsp/buf.lua

@@ -1413,20 +1413,16 @@ local function is_empty(range)
 end
 
 --- Perform an incremental selection at the cursor position based on ranges given by the LSP. The
---- `direction` parameter specifies whether the selection should head inward or outward.
+--- `direction` parameter specifies the number of times to expand the selection. Negative values
+--- will shrink the selection.
 ---
---- @param direction 'inner' | 'outer'
+--- @param direction integer
 function M.selection_range(direction)
-  validate('direction', direction, function(v)
+  validate('direction', direction, 'number')
-    return v == 'inner' or v == 'outer'
-  end)
 
   if selection_ranges then
-    local offset = direction == 'outer' and 1 or -1
+    local new_index = selection_ranges.index + direction
-    local new_index = selection_ranges.index + offset
+    selection_ranges.index = math.min(#selection_ranges.ranges, math.max(1, new_index))
-    if new_index <= #selection_ranges.ranges and new_index >= 1 then
-      selection_ranges.index = new_index
-    end
 
     select_range(selection_ranges.ranges[selection_ranges.index])
     return
@@ -1498,8 +1494,9 @@ function M.selection_range(direction)
       })
 
       if #ranges > 0 then
-        selection_ranges = { index = 1, ranges = ranges }
+        local index = math.min(#ranges, math.max(1, direction))
-        select_range(ranges[1])
+        selection_ranges = { index = index, ranges = ranges }
+        select_range(ranges[index])
       end
     end
   )

test/functional/plugin/lsp/selection_range_spec.lua

@@ -0,0 +1,122 @@
+local t = require('test.testutil')
+local n = require('test.functional.testnvim')()
+local t_lsp = require('test.functional.plugin.lsp.testutil')
+local Screen = require('test.functional.ui.screen')
+
+local dedent = t.dedent
+local exec_lua = n.exec_lua
+local insert = n.insert
+
+local clear_notrace = t_lsp.clear_notrace
+local create_server_definition = t_lsp.create_server_definition
+
+describe('vim.lsp.selection_range', function()
+  local text = dedent([[
+  hello
+  hello
+  hello
+  hello
+  hello]])
+
+  --- @type test.functional.ui.screen
+  local screen
+
+  before_each(function()
+    clear_notrace()
+    screen = Screen.new(50, 9)
+
+    exec_lua(create_server_definition)
+    exec_lua(function()
+      _G.server = _G._create_server({
+        capabilities = {
+          selectionRangeProvider = true,
+        },
+        handlers = {
+          ['textDocument/selectionRange'] = function(_, _, callback)
+            callback(nil, {
+              {
+                range = {
+                  start = { line = 2, character = 0 },
+                  ['end'] = { line = 2, character = 5 },
+                },
+                parent = {
+                  range = {
+                    start = { line = 1, character = 0 },
+                    ['end'] = { line = 3, character = 5 },
+                  },
+                  parent = {
+                    range = {
+                      start = { line = 0, character = 0 },
+                      ['end'] = { line = 5, character = 5 },
+                    },
+                    parent = nil,
+                  },
+                },
+              },
+            })
+          end,
+        },
+      })
+
+      return vim.lsp.start({ name = 'dummy', cmd = _G.server.cmd })
+    end)
+
+    insert(text)
+  end)
+
+  it('selects ranges', function()
+    -- Initial range
+    exec_lua(function()
+      local win = vim.api.nvim_get_current_win()
+      vim.api.nvim_win_set_cursor(win, { 3, 0 })
+      vim.lsp.buf.selection_range(1)
+    end)
+
+    screen:expect([[
+      hello                                             |*2
+      {17:hell}^o                                             |
+      hello                                             |*2
+      {1:~                                                 }|*3
+      {5:-- VISUAL --}                                      |
+    ]])
+
+    -- Outermost range
+    exec_lua(function()
+      vim.lsp.buf.selection_range(99)
+    end)
+
+    screen:expect([[
+      {17:hello}                                             |*4
+      {17:hell}^o                                             |
+      {1:~                                                 }|*3
+      {5:-- VISUAL --}                                      |
+    ]])
+
+    -- Back to innermost
+    exec_lua(function()
+      vim.lsp.buf.selection_range(-99)
+    end)
+
+    screen:expect([[
+      hello                                             |*2
+      {17:hell}^o                                             |
+      hello                                             |*2
+      {1:~                                                 }|*3
+      {5:-- VISUAL --}                                      |
+    ]])
+
+    -- Middle range
+    exec_lua(function()
+      vim.lsp.buf.selection_range(1)
+    end)
+
+    screen:expect([[
+      hello                                             |
+      {17:hello}                                             |*2
+      {17:hell}^o                                             |
+      hello                                             |
+      {1:~                                                 }|*3
+      {5:-- VISUAL --}                                      |
+    ]])
+  end)
+end)