From 5cfa03f2bdf312213b69cae329645f48da904ea1 Mon Sep 17 00:00:00 2001
From: Stephan Seitz <stephan.seitz@fau.de>
Date: Mon, 20 Jul 2020 23:56:32 +0200
Subject: Textobjects: add swap feature

---
 lua/nvim-treesitter/configs.lua     |   4 +-
 lua/nvim-treesitter/textobjects.lua | 120 +++++++++++++++++++++++++++++++++++-
 lua/nvim-treesitter/ts_utils.lua    |  67 +++++++++++++++++---
 queries/c/textobjects.scm           |   3 +
 queries/cpp/textobjects.scm         |   3 +
 5 files changed, 185 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 _, _ , nod