Implement config option to specify movements should stay in scope

Author: aaronik

lua/treewalker/init.lua

@@ -13,6 +13,7 @@ Treewalker.opts = {
   jumplist = true,
   select = false,
   notifications = true,
+  scope_confined = false,
 }
 
 -- This does not need to be called for Treewalker to work. The defaults are preinitialized and aim to be sane.

lua/treewalker/movement.lua

@@ -24,6 +24,24 @@ local function add_jumplist_for_move(command)
   end
 end
 
+---@param current_node TSNode
+---@param candidate TSNode
+---@return boolean
+local function should_confine_vertical_move(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
+
 ---@return nil
 function M.move_out()
   -- Add to jumplist at original cursor position before normalizing
@@ -55,6 +73,10 @@ function M.move_up()
   local target, row = targets.up(current_node, current_row)
   if not target or not row then return end
 
+  if should_confine_vertical_move(current_node, target) then
+    return
+  end
+
   local is_neighbor = nodes.have_neighbor_srow(current_node, target)
 
   if not is_neighbor then
@@ -70,6 +92,10 @@ function M.move_down()
   local target, row = targets.down(current_node, current_row)
   if not target or not row then return end
 
+  if should_confine_vertical_move(current_node, target) then
+    return
+  end
+
   local is_neighbor = nodes.have_neighbor_srow(current_node, target)
 
   if not is_neighbor then

lua/treewalker/nodes.lua

@@ -207,6 +207,27 @@ function M.get_descendants(node)
   return descendants
 end
 
+---@param parent TSNode
+---@param node TSNode
+---@return boolean
+function M.is_descendant_of(parent, node)
+  local iter = node
+  while iter do
+    if iter == parent then
+      return true
+    end
+    iter = iter:parent()
+  end
+  return false
+end
+
+---@param node TSNode
+---@return TSNode|nil
+function M.scope_parent(node)
+  local anchor = M.get_highest_row_coincident(node)
+  return anchor:parent()
+end
+
 -- Take row, give next row / node with same indentation
 ---@param current_row integer
 ---@param dir "up" | "down"

lua/treewalker/options.lua

@@ -7,7 +7,8 @@ local M = {}
 ---    jumplist: boolean | 'left',
 ---    highlight_group: string,
 ---    select: boolean,
----    notifications: boolean }
+---    notifications: boolean,
+---    scope_confined: boolean }
 
 ---@param opts Opts
 ---@return boolean, table<string>
@@ -41,6 +42,10 @@ function M.validate_opts(opts)
     table.insert(errors, "`notifications` should be boolean or nil")
   end
 
+  if type(opts.scope_confined) ~= "boolean" and opts.scope_confined ~= nil then
+    table.insert(errors, "`scope_confined` should be boolean or nil")
+  end
+
   if #errors == 0 then
     return true, {}
   else

tests/treewalker/c_sharp_spec.lua

@@ -103,4 +103,17 @@ describe("In a C Sharp file", function()
     assert.same(first_block, lines.get_lines(135, 138))
     h.assert_cursor_at(135, 13)
   end)
+  describe("scope_confined", function()
+    before_each(function()
+      tw.setup({ scope_confined = true })
+    end)
+
+    it("confines move_down", function()
+      h.assert_confined_by_parent(22, 1, 'down')
+    end)
+
+    it("confines move_up", function()
+      h.assert_confined_by_parent(20, 1, 'up')
+    end)
+  end)
 end)

tests/treewalker/c_spec.lua

@@ -70,4 +70,17 @@ describe("In a c file:", function()
     tw.move_up()
     h.assert_highlighted(26, 1, 33, 1)
   end)
+  describe("scope_confined", function()
+    before_each(function()
+      tw.setup({ scope_confined = true })
+    end)
+
+    it("confines move_down", function()
+      h.assert_confined_by_parent(23, 1, 'down')
+    end)
+
+    it("confines move_up", function()
+      h.assert_confined_by_parent(10, 1, 'up')
+    end)
+  end)
 end)

tests/treewalker/go_spec.lua

@@ -47,4 +47,18 @@ describe("Movement in Go file with tab indentation:", function()
     local row = vim.fn.line('.')
     assert.equals(19, row, "Should move from main to helper (line 19)")
   end)
+
+  describe("scope_confined", function()
+    before_each(function()
+      tw.setup({ scope_confined = true })
+    end)
+
+    it("confines move_down", function()
+      h.assert_confined_by_parent(17, 1, 'down')
+    end)
+
+    it("confines move_up", function()
+      h.assert_confined_by_parent(16, 1, 'up')
+    end)
+  end)
 end)

tests/treewalker/haskell_spec.lua

@@ -28,6 +28,17 @@ describe("In a haskell file: ", function()
   --   tw.move_out()
   --   h.assert_cursor_at(19, 1)
   -- end)
-end)
+  describe("scope_confined", function()
+    before_each(function()
+      tw.setup({ scope_confined = true })
+    end)
 
+    it("confines move_down", function()
+      h.assert_confined_by_parent(10, 1, 'down')
+    end)
 
+    it("confines move_up", function()
+      h.assert_confined_by_parent(11, 1, 'up')
+    end)
+  end)
+end)