Textobjects: add swap feature

3 files changed, 179 insertions, 12 deletions

diff --git a/lua/nvim-treesitter/configs.lua b/lua/nvim-treesitter/configs.lua
index 02a636dff..350a8cdad 100644
--- a/lua/nvim-treesitter/configs.lua
+++ b/lua/nvim-treesitter/configs.lua
@@ -85,7 +85,9 @@ local builtin_modules = {
     is_supported = function(lang)
       return has_some_textobject_mapping(lang) or queries.has_textobjects(lang)
     end,
-    keymaps = {}
+    keymaps = {},
+    swap_next_keymaps = {},
+    swap_previous_keymaps = {}
   }
 }
 
diff --git a/lua/nvim-treesitter/textobjects.lua b/lua/nvim-treesitter/textobjects.lua
index 7227b9e0e..9c312f44a 100644
--- a/lua/nvim-treesitter/textobjects.lua
+++ b/lua/nvim-treesitter/textobjects.lua
@@ -8,7 +8,7 @@ local ts_utils = require'nvim-treesitter.ts_utils'
 
 local M = {}
 
-function M.select_textobject(query_string)
+local function get_textobject_at_point(query_string)
   local bufnr = vim.api.nvim_get_current_buf()
   local lang = parsers.get_buf_lang(bufnr)
   if not lang then return end
@@ -66,11 +66,93 @@ function M.select_textobject(query_string)
     if smallest_range.start then
       local start_range = {smallest_range.start.node:range()}
       local node_range = {smallest_range.node:range()}
-      ts_utils.update_selection(bufnr, {start_range[1], start_range[2], node_range[3], node_range[4]})
+      return bufnr, {start_range[1], start_range[2], node_range[3], node_range[4]}, smallest_range.node
     else
-      ts_utils.update_selection(bufnr, smallest_range.node)
+      return bufnr, {smallest_range.node:range()}, smallest_range.node
+    end
+  end
+end
+
+function M.select_textobject(query_string)
+  local bufnr, textobject = get_textobject_at_point(query_string)
+  if textobject then
+    ts_utils.update_selection(bufnr, textobject)
+  end
+end
+
+local function swap_textobject(query_string, direction)
+  local bufnr, textobject_range, node = get_textobject_at_point(query_string)
+  local step = direction > 0 and 1 or -1
+  if not node then return end
+  for _ = 1, math.abs(direction), step do
+      if direction > 0 then
+        ts_utils.swap_nodes(textobject_range, M.next_textobject(node, query_string, true, bufnr), bufnr, "yes, set cursor!")
+      else
+        ts_utils.swap_nodes(textobject_range, M.previous_textobject(node, query_string, true, bufnr), bufnr, "yes, set cursor!")
+      end
+  end
+end
+
+function M.swap_textobject_next(query_string)
+  swap_textobject(query_string, 1)
+end
+
+function M.swap_textobject_previous(query_string)
+  swap_textobject(query_string, -1)
+end
+
+function M.next_textobject(node, query_string, same_parent, bufnr)
+  local bufnr = bufnr or api.nvim_get_current_buf()
+
+  local matches = queries.get_capture_matches(bufnr, query_string, 'textobjects')
+  local _, _ , node_end = node:end_()
+  local next_node
+  local next_node_start
+
+  for _, m in pairs(matches) do
+    local _, _, other_end = m.node:start()
+    if other_end > node_end then
+      if not same_parent or node:parent() == m.node:parent() then
+        if not next_node then
+          next_node = m
+          _, _, next_node_start = next_node.node:start()
+        end
+        if other_end < next_node_start then
+          next_node = m
+          _, _, next_node_start = next_node.node:start()
+        end
+      end
     end
   end
+
+  return next_node and next_node.node
+end
+
+function M.previous_textobject(node, query_string, same_parent, bufnr)
+  local bufnr = bufnr or api.nvim_get_current_buf()
+
+  local matches = queries.get_capture_matches(bufnr, query_string, 'textobjects')
+  local _, _ , node_start = node:start()
+  local previous_node
+  local previous_node_end
+
+  for _, m in pairs(matches) do
+    local _, _, other_end = m.node:end_()
+    if other_end < node_start then
+      if not same_parent or node:parent() == m.node:parent() then
+        if not previous_node then
+          previous_node = m
+          _, _, previous_node_end = previous_node.node:end_()
+        end
+        if other_end > previous_node_end then
+          previous_node = m
+          _, _, previous_node_end = previous_node.node:end_()
+        end
+      end
+    end
+  end
+
+  return previous_node and previous_node.node
 end
 
 function M.attach(bufnr, lang)
@@ -90,6 +172,28 @@ function M.attach(bufnr, lang)
       api.nvim_buf_set_keymap(buf, "v", mapping, cmd, {silent = true, noremap = true })
     end
   end
+  for mapping, query in pairs(config.swap_next_keymaps) do
+    if type(query) == 'table' then
+      query = query[lang]
+    elseif not queries.get_query(lang, 'textobjects') then
+      query = nil
+    end
+    if query then
+      local cmd = ":lua require'nvim-treesitter.textobjects'.swap_textobject_next('"..query.."')<CR>"
+      api.nvim_buf_set_keymap(buf, "n", mapping, cmd, {silent = true})
+    end
+  end
+  for mapping, query in pairs(config.swap_previous_keymaps) do
+    if type(query) == 'table' then
+      query = query[lang]
+    elseif not queries.get_query(lang, 'textobjects') then
+      query = nil
+    end
+    if query then
+      local cmd = ":lua require'nvim-treesitter.textobjects'.swap_textobject_previous('"..query.."')<CR>"
+      api.nvim_buf_set_keymap(buf, "n", mapping, cmd, {silent = true})
+    end
+  end
 end
 
 function M.detach(bufnr)
@@ -108,6 +212,16 @@ function M.detach(bufnr)
       api.nvim_buf_del_keymap(buf, "v", mapping)
     end
   end
+  for mapping, query in pairs(config.swap_next_keymaps) or pairs(config.swap_previous_keymaps) do
+    if type(query) == 'table' then
+      query = query[lang]
+    elseif not queries.get_query(lang, 'textobjects') then
+      query = nil
+    end
+    if query then
+      api.nvim_buf_del_keymap(buf, "n", mapping)
+    end
+  end
 end
 
 return M
diff --git a/lua/nvim-treesitter/ts_utils.lua b/lua/nvim-treesitter/ts_utils.lua
index 0a5cbc608..eddd0f56b 100644
--- a/lua/nvim-treesitter/ts_utils.lua
+++ b/lua/nvim-treesitter/ts_utils.lua
@@ -13,7 +13,7 @@ function M.get_node_text(node, bufnr)
   if not node then return {} end
 
   -- We have to remember that end_col is end-exclusive
-  local start_row, start_col, end_row, end_col = node:range()
+  local start_row, start_col, end_row, end_col = M.get_node_range(node)
 
   if start_row ~= end_row then
     local lines = api.nvim_buf_get_lines(bufnr, start_row, end_row+1, false)
@@ -131,12 +131,7 @@ end
 
 -- Set visual selection to node
 function M.update_selection(buf, node)
-  local start_row, start_col, end_row, end_col
-  if type(node) == 'table' then
-    start_row, start_col, end_row, end_col = unpack(node)
-  else
-    start_row, start_col, end_row, end_col = node:range()
-  end
+  local start_row, start_col, end_row, end_col = M.get_node_range(node)
 
   if end_row == vim.fn.line('$') then
     end_col = #vim.fn.getline('$')
@@ -187,8 +182,16 @@ function M.is_in_node_range(node, line, col)
   end
 end
 
+function M.get_node_range(node_or_range)
+  if type(node_or_range) == 'table' then
+    return unpack(node_or_range)
+  else
+    return node_or_range:range()
+  end
+end
+
 function M.node_to_lsp_range(node)
-  local start_line, start_col, end_line, end_col = node:range()
+  local start_line, start_col, end_line, end_col = M.get_node_range(node)
   local rtn = {}
   rtn.start = { line = start_line, character = start_col }
   rtn['end'] = { line = end_line, character = end_col }
@@ -225,4 +228,52 @@ function M.memoize_by_buf_tick(fn, bufnr_fn)
   end
 end
 
+function M.swap_nodes(node_or_range1, node_or_range2, bufnr, cursor_to_second)
+  if not node_or_range1 or not node_or_range2 then return end
+  local range1 = M.node_to_lsp_range(node_or_range1)
+  local range2 = M.node_to_lsp_range(node_or_range2)
+
+  local text1 = M.get_node_text(node_or_range1)
+  local text2 = M.get_node_text(node_or_range2)
+
+  local edit1 = { range = range1, newText = table.concat(text2, '\n') }
+  local edit2 = { range = range2, newText = table.concat(text1, '\n') }
+  vim.lsp.util.apply_text_edits({edit1, edit2}, bufnr)
+
+  if cursor_to_second then
+    -- Set the item in jump list
+    vim.cmd "normal! m'"
+
+    local char_delta = 0
+    local line_delta = 0
+    if range1["end"].line < range2.start.line
+       or (range1["end"].line == range2.start.line and range1["end"].character < range2.start.character) then
+      line_delta = #text2 - #text1
+    end
+
+    if range1["end"].line == range2.start.line and range1["end"].character < range2.start.character then
+      if line_delta ~= 0 then
+        --- why?
+        --correction_after_line_change =  -range2.start.character
+        --text_now_before_range2 = #(text2[#text2])
+        --space_between_ranges = range2.start.character - range1["end"].character
+        --char_delta = correction_after_line_change + text_now_before_range2 + space_between_ranges
+        --- Equivalent to:
+        char_delta = #(text2[#text2]) - range1["end"].character
+
+        -- add range1.start.character if last line of range1 (now text2) does not start at 0
+        if range1.start.line == range2.start.line + line_delta then
+            char_delta = char_delta + range1.start.character
+        end
+      else
+        char_delta = #(text2[#text2]) - #(text1[#text1])
+      end
+    end
+
+    api.nvim_win_set_cursor(api.nvim_get_current_win(),
+                            {range2.start.line + 1 + line_delta,
+                             range2.start.character + char_delta})
+  end
+end
+
 return M