Implement scope confinement for swap commands

Author: aaronik

lua/treewalker/swap.lua

@@ -38,6 +38,24 @@ local function is_supported_ft()
   return not unsupported_filetypes[ft]
 end
 
+---@param current_node TSNode
+---@param candidate TSNode
+---@return boolean
+local function should_confine_swap(current_node, candidate)
+  local opts = require('treewalker').opts
+  if opts.scope_confined ~= true then
+    return false
+  end
+
+  local current_parent = nodes.scope_parent(current_node)
+  if not current_parent then
+    return false
+  end
+
+  local candidate_anchor = nodes.get_highest_row_coincident(candidate)
+  return not nodes.is_descendant_of(current_parent, candidate_anchor)
+end
+
 function M.swap_down()
   vim.cmd("normal! ^")
   if not is_supported_ft() then return end
@@ -50,6 +68,10 @@ function M.swap_down()
   local target = targets.down(current, row)
   if not target then return end
 
+  if should_confine_swap(current, target) then
+    return
+  end
+
   local current_augments = augment.get_node_augments(current)
   local current_all = { current, unpack(current_augments) }
   local current_srow = nodes.get_srow(current)
@@ -82,6 +104,10 @@ function M.swap_up()
   local target = targets.up(current, row)
   if not target then return end
 
+  if should_confine_swap(current, target) then
+    return
+  end
+
   local current_augments = augment.get_node_augments(current)
   local current_all = { current, unpack(current_augments) }
   local current_srow = nodes.get_srow(current)
@@ -121,6 +147,10 @@ function M.swap_right()
 
   if not current or not target then return end
 
+  if should_confine_swap(current, target) then
+    return
+  end
+
   -- set a mark to track where the target started, so we may later go there after the swap
   local ns_id = vim.api.nvim_create_namespace("treewalker#swap_right")
   local ext_id = vim.api.nvim_buf_set_extmark(
@@ -158,6 +188,10 @@ function M.swap_left()
 
   if not current or not target then return end
 
+  if should_confine_swap(current, target) then
+    return
+  end
+
   operations.swap_nodes(target, current)
 
   -- Place cursor

tests/treewalker/c_sharp_spec.lua

@@ -115,5 +115,21 @@ describe("In a C Sharp file", function()
     it("confines move_up", function()
       h.assert_confined_by_parent(20, 1, 'up')
     end)
+
+    it("confines swap_down", function()
+      h.assert_swap_confined_by_parent(22, 1, 'down')
+    end)
+
+    it("confines swap_up", function()
+      h.assert_swap_confined_by_parent(20, 1, 'up')
+    end)
+
+    it("confines swap_right", function()
+      h.assert_swap_confined_by_parent(22, 1, 'right')
+    end)
+
+    it("confines swap_left", function()
+      h.assert_swap_confined_by_parent(20, 1, 'left')
+    end)
   end)
 end)

tests/treewalker/c_spec.lua

@@ -82,5 +82,21 @@ describe("In a c file:", function()
     it("confines move_up", function()
       h.assert_confined_by_parent(10, 1, 'up')
     end)
+
+    it("confines swap_down", function()
+      h.assert_swap_confined_by_parent(23, 1, 'down')
+    end)
+
+    it("confines swap_up", function()
+      h.assert_swap_confined_by_parent(10, 1, 'up')
+    end)
+
+    it("confines swap_right", function()
+      h.assert_swap_confined_by_parent(23, 1, 'right')
+    end)
+
+    it("confines swap_left", function()
+      h.assert_swap_confined_by_parent(10, 1, 'left')
+    end)
   end)
 end)

tests/treewalker/go_spec.lua

@@ -60,5 +60,21 @@ describe("Movement in Go file with tab indentation:", function()
     it("confines move_up", function()
       h.assert_confined_by_parent(16, 1, 'up')
     end)
+
+    it("confines swap_down", function()
+      h.assert_swap_confined_by_parent(17, 1, 'down')
+    end)
+
+    it("confines swap_up", function()
+      h.assert_swap_confined_by_parent(16, 1, 'up')
+    end)
+
+    it("confines swap_right", function()
+      h.assert_swap_confined_by_parent(17, 1, 'right')
+    end)
+
+    it("confines swap_left", function()
+      h.assert_swap_confined_by_parent(16, 1, 'left')
+    end)
   end)
 end)

tests/treewalker/haskell_spec.lua

@@ -40,5 +40,21 @@ describe("In a haskell file: ", function()
     it("confines move_up", function()
       h.assert_confined_by_parent(11, 1, 'up')
     end)
+
+    it("confines swap_down", function()
+      h.assert_swap_confined_by_parent(10, 1, 'down')
+    end)
+
+    it("confines swap_up", function()
+      h.assert_swap_confined_by_parent(11, 1, 'up')
+    end)
+
+    it("confines swap_right", function()
+      h.assert_swap_confined_by_parent(10, 1, 'right')
+    end)
+
+    it("confines swap_left", function()
+      h.assert_swap_confined_by_parent(11, 1, 'left')
+    end)
   end)
 end)

tests/treewalker/helpers.lua

@@ -157,4 +157,39 @@ function M.assert_confined_by_parent(row, col, direction)
 	assert(vim.treesitter.is_ancestor(scope_parent, after_anchor))
 end
 
+---
+--- Swap should also respect scope_confined: when swapping at the edge of the current
+--- scope, it should not swap with a node outside that scope.
+---
+---@param row integer
+---@param col integer
+---@param direction 'up'|'down'|'left'|'right'
+function M.assert_swap_confined_by_parent(row, col, direction)
+	local nodes = require("treewalker.nodes")
+
+	vim.fn.cursor(row, col)
+
+	local before = nodes.get_highest_row_coincident(nodes.get_highest_node_at_current_row())
+	local scope_parent = before:parent()
+
+	assert(scope_parent ~= nil, "test must start within a non-top-level scope")
+
+	local original = lines.get_lines(0, -1)
+
+	if direction == "up" then
+		tw.swap_up()
+	elseif direction == "down" then
+		tw.swap_down()
+	elseif direction == "left" then
+		tw.swap_left()
+	elseif direction == "right" then
+		tw.swap_right()
+	else
+		error("invalid direction: " .. tostring(direction))
+	end
+
+	-- Should not have swapped across the scope boundary.
+	assert.same(original, lines.get_lines(0, -1))
+end
+
 return M

tests/treewalker/html_spec.lua

@@ -65,5 +65,21 @@ describe("In an html file", function()
     it("confines move_up", function()
       h.assert_confined_by_parent(15, 1, 'up')
     end)
+
+    it("confines swap_down", function()
+      h.assert_swap_confined_by_parent(15, 1, 'down')
+    end)
+
+    it("confines swap_up", function()
+      h.assert_swap_confined_by_parent(15, 1, 'up')
+    end)
+
+    it("confines swap_right", function()
+      h.assert_swap_confined_by_parent(15, 1, 'right')
+    end)
+
+    it("confines swap_left", function()
+      h.assert_swap_confined_by_parent(15, 1, 'left')
+    end)
   end)
 end)

tests/treewalker/java_spec.lua

@@ -154,5 +154,21 @@ describe("In a java file:", function()
     it("confines move_up", function()
       h.assert_confined_by_parent(29, 1, 'up')
     end)
+
+    it("confines swap_down", function()
+      h.assert_swap_confined_by_parent(24, 1, 'down')
+    end)
+
+    it("confines swap_up", function()
+      h.assert_swap_confined_by_parent(29, 1, 'up')
+    end)
+
+    it("confines swap_right", function()
+      h.assert_swap_confined_by_parent(24, 1, 'right')
+    end)
+
+    it("confines swap_left", function()
+      h.assert_swap_confined_by_parent(29, 1, 'left')
+    end)
   end)
 end)

tests/treewalker/javascript_spec.lua

@@ -28,5 +28,21 @@ describe("Swapping in a typescript file:", function()
     it("confines move_up", function()
       h.assert_confined_by_parent(7, 1, 'up')
     end)
+
+    it("confines swap_down", function()
+      h.assert_swap_confined_by_parent(5, 3, 'down')
+    end)
+
+    it("confines swap_up", function()
+      h.assert_swap_confined_by_parent(7, 1, 'up')
+    end)
+
+    it("confines swap_right", function()
+      h.assert_swap_confined_by_parent(5, 3, 'right')
+    end)
+
+    it("confines swap_left", function()
+      h.assert_swap_confined_by_parent(7, 1, 'left')
+    end)
   end)
 end)

tests/treewalker/markdown_spec.lua

@@ -382,5 +382,21 @@ describe("Swapping in a markdown file with h2s at the top:", function()
     it("confines move_up", function()
       h.assert_confined_by_parent(19, 1, 'up')
     end)
+
+    it("confines swap_down", function()
+      h.assert_swap_confined_by_parent(33, 1, 'down')
+    end)
+
+    it("confines swap_up", function()
+      h.assert_swap_confined_by_parent(19, 1, 'up')
+    end)
+
+    it("confines swap_right", function()
+      h.assert_swap_confined_by_parent(33, 1, 'right')
+    end)
+
+    it("confines swap_left", function()
+      h.assert_swap_confined_by_parent(19, 1, 'left')
+    end)
   end)
 end)