nested selector

lua/treesj/format.lua

@@ -45,7 +45,7 @@ function M.get_node_at_cursor(root_lang_tree)
   )
 end
 
-function M._format(mode, override)
+function M._format(mode, override, selector)
   -- Tree reparsing is required, otherwise the tree may not be updated
   -- and each node will be processed only once (until
   -- the tree is updated). See issue #118
@@ -67,7 +67,7 @@ function M._format(mode, override)
 
   -- If the node is marked as "disabled", continue searching from its parent.
   while true do
-    found, tsn_data = pcall(search.get_configured_node, start_node)
+    found, tsn_data = pcall(selector or search.get_configured_node, start_node)
     if not found then
       notify.warn(tsn_data)
       return
@@ -169,4 +169,35 @@ function M._format(mode, override)
   pcall(vim.api.nvim_win_set_cursor, 0, new_cursor)
 end
 
+M.last_selected = nil
+function M._nested(selector, mode, preset)
+  local nodes
+  M._format(mode, preset, function(start_node)
+    if not nodes then
+      nodes = search.get_configured_nodes(start_node)
+    end
+    local selected
+    if settings.remember_selected and M.last_selected then
+      -- TODO: basic indexing is not ideal if the sticky cursor doesn't work?
+      selected = nodes[M.last_selected]
+    else
+      if type(selector) == 'string' then
+        selector = require('treesj.selectors.' .. selector).selector
+      end
+      local sel, index = selector(nodes)
+      if sel then
+        M.last_selected = index
+        selected = sel
+      end
+    end
+    if selected then
+      return selected
+    else
+      -- TODO: change the error messages
+      error(msg.no_chosen_node, 0)
+      return
+    end
+  end)
+end
+
 return M

lua/treesj/init.lua

@@ -11,32 +11,49 @@ M.setup = function(opts)
   settings._create_commands()
 end
 
-M.__format = function()
-  require('treesj.format')._format(M._mode_to_use, M._preset_to_use)
-end
-
-local function perform(mode, preset)
-  M._mode_to_use = mode
-  M._preset_to_use = preset
-
+local function repeatable(fn)
   if settings.settings.dot_repeat then
-    vim.opt.operatorfunc = "v:lua.require'treesj'.__format"
-    vim.api.nvim_feedkeys('g@l', 'nix', true)
+    M.__repeat = fn
+    vim.opt.operatorfunc = "v:lua.require'treesj'.__repeat"
+    vim.api.nvim_feedkeys(vim.v.count1 .. 'g@l', 'n', true)
   else
-    M.__format()
+    fn()
   end
 end
 
+M.format = function(mode, preset)
+  M.nested_format('count', mode, preset)
+end
+
 M.toggle = function(preset)
-  perform(nil, preset)
+  M.format(nil, preset)
 end
 
 M.join = function(preset)
-  perform("join", preset)
+  M.format('join', preset)
 end
 
 M.split = function(preset)
-  perform("split", preset)
+  M.format('split', preset)
+end
+
+M.nested_format = function(selector, mode, preset)
+  require('treesj.format').last_selected = nil
+  repeatable(function()
+    require('treesj.format')._nested(selector, mode, preset)
+  end)
+end
+
+M.nested_toggle = function(selector, preset)
+  M.nested_format(selector, nil, preset)
+end
+
+M.nested_join = function(selector, preset)
+  M.nested_format(selector, 'join', preset)
+end
+
+M.nested_split = function(selector, preset)
+  M.nested_format(selector, 'split', preset)
 end
 
 return M

lua/treesj/notify.lua

@@ -7,6 +7,7 @@ local M = {}
 
 M.msg = {
   no_detect_node = 'No detected node at cursor',
+  no_chosen_node = 'Node choice aborted',
   no_configured_lang = 'Language "%s" is not configured',
   contains_error = 'The node "%s" or its descendants contain a syntax error and cannot be %s',
   no_configured_node = 'Node "%s" for lang "%s" is not configured',

lua/treesj/search.lua

@@ -173,6 +173,46 @@ function M.get_configured_node(node)
 
   return data
 end
+---Return the all configured nodes
+---@param node TSNode|nil TSNode instance
+---@return table
+function M.get_configured_nodes(node)
+  if not node then
+    error(msg.node_not_received, 0)
+  end
+
+  local lang = get_node_lang(node)
+  if not langs[lang] then
+    error(msg.no_configured_lang:format(lang), 0)
+  end
+  local start_node_type = node:type()
+
+  local nodes = {}
+  local done = {}
+  while node do
+    local data = search_node(node, lang)
+
+    if not data or not data.tsnode then
+      error(msg.no_configured_node:format(start_node_type, lang), 0)
+      break
+    end
+
+    local id = table.concat({ data.tsnode:range() }, '-')
+    if done[id] then
+      break
+    else
+      done[id] = true
+    end
+
+    nodes[#nodes + 1] = data
+    node = data.tsnode:parent()
+  end
+
+  if #nodes == 0 then
+    error(msg.no_configured_node:format(start_node_type, lang), 0)
+  end
+  return nodes
+end
 
 ---Return the preset for current node if it no contains field 'target_nodes'
 ---@param tsn_type string TSNode type

lua/treesj/selectors/count.lua

@@ -0,0 +1,8 @@
+local M = {}
+local meta_M = {}
+M = setmetatable(M, meta_M)
+function M.selector(nodes)
+  local c = vim.v.count1
+  return nodes[c], c
+end
+return M

lua/treesj/selectors/flash.lua

@@ -0,0 +1,97 @@
+local M = {}
+local flash = require('flash')
+local Pos = require('flash.search.pos')
+
+local opts = { mode = 'treesj' }
+function M.setup(config)
+  opts = config
+end
+
+function M.selector(nodes)
+  local currwin = vim.api.nvim_get_current_win()
+  local buf = vim.api.nvim_get_current_buf()
+  local line_count = vim.api.nvim_buf_line_count(buf)
+  local state = flash.jump(vim.tbl_deep_extend('force', {
+    matcher = function(win, _)
+      if win ~= currwin then
+        return {}
+      end
+      local ret = {}
+      local done = {}
+      -- https://github.com/folke/flash.nvim/blob/967117690bd677cb7b6a87f0bc0077d2c0be3a27/lua/flash/plugins/treesitter.lua#L52
+      for i, node in ipairs(nodes) do
+        local tsn = node.tsnode
+        local range = { tsn:range() }
+        local match = {
+          win = win,
+          node = node,
+          pos = { range[1] + 1, range[2] },
+          end_pos = { range[3] + 1, range[4] - 1 },
+          index = i,
+        }
+
+        -- If the match is at the end of the buffer,
+        -- then move it to the last character of the last line.
+        if match.end_pos[1] > line_count then
+          match.end_pos[1] = line_count
+          match.end_pos[2] = #vim.api.nvim_buf_get_lines(
+            buf,
+            match.end_pos[1] - 1,
+            match.end_pos[1],
+            false
+          )[1]
+        elseif match.end_pos[2] == -1 then
+          -- If the end points to the start of the next line, move it to the
+          -- end of the previous line.
+          -- Otherwise operations include the first character of the next line
+          local line = vim.api.nvim_buf_get_lines(
+            buf,
+            match.end_pos[1] - 2,
+            match.end_pos[1] - 1,
+            false
+          )[1]
+          match.end_pos[1] = match.end_pos[1] - 1
+          match.end_pos[2] = #line
+        end
+        local id =
+          table.concat(vim.tbl_flatten({ match.pos, match.end_pos }), '.')
+        if not done[id] then
+          done[id] = true
+          ret[#ret + 1] = match
+        end
+      end
+      for m, match in ipairs(ret) do
+        match.pos = Pos(match.pos)
+        match.end_pos = Pos(match.end_pos)
+        match.depth = #ret - m
+      end
+      return ret
+    end,
+    action = function(match, state)
+      state.final_match = match
+    end,
+    search = {
+      multi_window = false,
+      wrap = true,
+      incremental = false,
+      max_length = 0,
+    },
+    label = {
+      before = true,
+      after = true,
+    },
+    highlight = {
+      matches = false,
+    },
+    actions = {
+      -- TODO: incremental preview/operations
+    },
+    jump = { autojump = true },
+  }, opts or {}))
+  local m = state.final_match
+  if m then
+    return m.node, m.index
+  end
+end
+
+return M

lua/treesj/settings.lua

@@ -25,6 +25,8 @@ local DEFAULT_SETTINGS = {
   dot_repeat = true,
   ---@type nil|function Callback for treesj error handler. func (err_text, level, ...)
   on_error = nil,
+  ---@type boolean Whether dot repeat on nested operations should remember the selected level of nesting
+  remember_selected = true,
 }
 
 local commands = {