feat(ts_utils): allow starting selection after last character (#3233)

2 files changed, 69 insertions, 2 deletions

diff --git a/lua/nvim-treesitter/ts_utils.lua b/lua/nvim-treesitter/ts_utils.lua
index 8200e78ed..f5ad3354e 100644
--- a/lua/nvim-treesitter/ts_utils.lua
+++ b/lua/nvim-treesitter/ts_utils.lua
@@ -259,14 +259,14 @@ function M.update_selection(buf, node, selection_mode)
   selection_mode = selection_mode or "charwise"
   local start_row, start_col, end_row, end_col = M.get_vim_range({ M.get_node_range(node) }, buf)
 
-  vim.fn.setpos(".", { buf, start_row, start_col, 0 })
-
   -- Start visual selection in appropriate mode
   local v_table = { charwise = "v", linewise = "V", blockwise = "<C-v>" }
   ---- Call to `nvim_replace_termcodes()` is needed for sending appropriate
   ---- command to enter blockwise mode
   local mode_string = vim.api.nvim_replace_termcodes(v_table[selection_mode] or selection_mode, true, true, true)
   vim.cmd("normal! " .. mode_string)
+  vim.fn.setpos(".", { buf, start_row, start_col, 0 })
+  vim.cmd "normal! o"
   vim.fn.setpos(".", { buf, end_row, end_col, 0 })
 end
 
diff --git a/tests/unit/ts_utils_spec.lua b/tests/unit/ts_utils_spec.lua
index 1825a3b68..a3b334413 100644
--- a/tests/unit/ts_utils_spec.lua
+++ b/tests/unit/ts_utils_spec.lua
@@ -39,3 +39,70 @@ describe("is_in_node_range", function()
     assert.is_false(test_is_in_node_range(3, 0))
   end)
 end)
+
+describe("update_selection", function()
+  local function get_updated_selection(case)
+    vim.api.nvim_buf_set_lines(0, 0, -1, false, case.lines)
+    tsutils.update_selection(0, case.node, case.selection_mode)
+    vim.cmd "normal! y"
+    return vim.fn.getreg '"'
+  end
+
+  it("charwise1", function()
+    assert.equal(
+      get_updated_selection {
+        lines = { "foo", "", "bar" },
+        node = { 0, 0, 2, 1 },
+        selection_mode = "v",
+      },
+      "foo\n\nb"
+    )
+    it("charwise2", function() end)
+    assert.equal(
+      get_updated_selection {
+        lines = { "foo", "", "bar" },
+        node = { 0, 1, 2, 1 },
+        selection_mode = "v",
+      },
+      "oo\n\nb"
+    )
+    it("charwise3", function() end)
+    assert.equal(
+      get_updated_selection {
+        lines = { "foo", "", "bar" },
+        node = { 0, 2, 2, 1 },
+        selection_mode = "v",
+      },
+      "o\n\nb"
+    )
+    it("charwise4", function() end)
+    assert.equal(
+      get_updated_selection {
+        lines = { "foo", "", "bar" },
+        node = { 0, 3, 2, 1 },
+        selection_mode = "v",
+      },
+      "\n\nb"
+    )
+  end)
+  it("linewise", function()
+    assert.equal(
+      get_updated_selection {
+        lines = { "foo", "", "bar" },
+        node = { 0, 3, 2, 1 },
+        selection_mode = "V",
+      },
+      "foo\n\nbar\n"
+    )
+  end)
+  it("blockwise", function()
+    assert.equal(
+      get_updated_selection {
+        lines = { "foo", "", "bar" },
+        node = { 0, 3, 2, 1 },
+        selection_mode = "<C-v>",
+      },
+      "foo\n\nbar"
+    )
+  end)
+end)