chore(iterators): create Iterator module and migrate iterators to use it (#1392)

lua/nvim-tree/actions/collapse-all.lua

@@ -1,6 +1,7 @@
 local renderer = require "nvim-tree.renderer"
 local utils = require "nvim-tree.utils"
 local core = require "nvim-tree.core"
+local Iterator = require "nvim-tree.iterators.node-iterator"
 
 local M = {}
 
@@ -9,34 +10,29 @@ function M.fn(keep_buffers)
     return
   end
 
-  local buffer_paths = {}
-  for _, buffer in ipairs(vim.api.nvim_list_bufs()) do
-    table.insert(buffer_paths, vim.api.nvim_buf_get_name(buffer))
-  end
+  local buffer_paths = vim.tbl_map(function(buffer)
+    return vim.api.nvim_buf_get_name(buffer)
+  end, vim.api.nvim_list_bufs())
 
-  local function iter(nodes)
-    for _, node in pairs(nodes) do
-      if node.open then
-        local new_open = false
-
-        if keep_buffers == true then
-          for _, buffer_path in ipairs(buffer_paths) do
-            local matches = utils.str_find(buffer_path, node.absolute_path)
-            if matches then
-              new_open = true
-            end
+  Iterator.builder(core.get_explorer().nodes)
+    :hidden()
+    :applier(function(node)
+      node.open = false
+      if keep_buffers == true then
+        for _, buffer_path in ipairs(buffer_paths) do
+          local matches = utils.str_find(buffer_path, node.absolute_path)
+          if matches then
+            node.open = true
+            return
           end
         end
-
-        node.open = new_open
       end
-      if node.nodes then
-        iter(node.nodes)
-      end
-    end
-  end
+    end)
+    :recursor(function(n)
+      return n.nodes
+    end)
+    :iterate()
 
-  iter(core.get_explorer().nodes)
   renderer.draw()
 end
 

lua/nvim-tree/actions/expand-all.lua

@@ -1,6 +1,7 @@
 local core = require "nvim-tree.core"
 local renderer = require "nvim-tree.renderer"
 local utils = require "nvim-tree.utils"
+local Iterator = require "nvim-tree.iterators.node-iterator"
 
 local M = {}
 
@@ -20,34 +21,38 @@ local function expand(node)
   end
 end
 
+local function should_expand(expansion_count, node)
+  local should_halt = expansion_count >= M.MAX_FOLDER_DISCOVERY
+  local should_exclude = M.EXCLUDE[node.name]
+  return not should_halt and node.nodes and not node.open and not should_exclude
+end
+
 local function gen_iterator()
   local expansion_count = 0
 
-  local function iterate(parent)
-    if expansion_count >= M.MAX_FOLDER_DISCOVERY then
-      return true
-    end
-
+  return function(parent)
     if parent.parent and parent.nodes and not parent.open then
       expansion_count = expansion_count + 1
       expand(parent)
     end
 
-    for _, node in pairs(parent.nodes) do
-      if node.nodes and not node.open and not M.EXCLUDE[node.name] then
-        expansion_count = expansion_count + 1
-        expand(node)
-      end
-
-      if node.open then
-        if iterate(node) then
-          return true
+    Iterator.builder(parent.nodes)
+      :hidden()
+      :applier(function(node)
+        if should_expand(expansion_count, node) then
+          expansion_count = expansion_count + 1
+          expand(node)
         end
-      end
+      end)
+      :recursor(function(node)
+        return expansion_count < M.MAX_FOLDER_DISCOVERY and node.open and node.nodes
+      end)
+      :iterate()
+
+    if expansion_count >= M.MAX_FOLDER_DISCOVERY then
+      return true
     end
   end
-
-  return iterate
 end
 
 function M.fn(base_node)

lua/nvim-tree/actions/find-file.lua

@@ -4,6 +4,7 @@ local view = require "nvim-tree.view"
 local utils = require "nvim-tree.utils"
 local renderer = require "nvim-tree.renderer"
 local core = require "nvim-tree.core"
+local Iterator = require "nvim-tree.iterators.node-iterator"
 
 local M = {}
 
@@ -24,53 +25,31 @@ function M.fn(fname)
     return
   end
 
-  local i = core.get_nodes_starting_line() - 1
-  local tree_altered = false
+  local line = core.get_nodes_starting_line()
 
-  local function iterate_nodes(nodes)
-    for _, node in ipairs(nodes) do
-      if not node.hidden then
-        i = i + 1
+  local found = Iterator.builder(core.get_explorer().nodes)
+    :matcher(function(node)
+      return node.absolute_path == fname_real or node.link_to == fname_real
+    end)
+    :applier(function(node)
+      line = line + 1
+      local abs_match = vim.startswith(fname_real, node.absolute_path .. utils.path_separator)
+      local link_match = node.link_to and vim.startswith(fname_real, node.link_to .. utils.path_separator)
 
-        if not node.absolute_path or not uv.fs_stat(node.absolute_path) then
-          break
-        end
-
-        -- match against node absolute and link, as symlinks themselves will differ
-        if node.absolute_path == fname_real or node.link_to == fname_real then
-          return i
-        end
-        local abs_match = vim.startswith(fname_real, node.absolute_path .. utils.path_separator)
-        local link_match = node.link_to and vim.startswith(fname_real, node.link_to .. utils.path_separator)
-        local path_matches = node.nodes and (abs_match or link_match)
-        if path_matches then
-          if not node.open then
-            node.open = true
-            tree_altered = true
-          end
-
-          if #node.nodes == 0 then
-            core.get_explorer():expand(node)
-          end
-
-          if iterate_nodes(node.nodes) ~= nil then
-            return i
-          end
-          -- mandatory to iterate i
-        elseif node.open then
-          iterate_nodes(node.nodes)
+      if abs_match or link_match then
+        node.open = true
+        if #node.nodes == 0 then
+          core.get_explorer():expand(node)
         end
       end
-    end
+    end)
+    :iterate()
+
+  if found and view.is_visible() then
+    renderer.draw()
+    view.set_cursor { line, 0 }
   end
 
-  local index = iterate_nodes(core.get_explorer().nodes)
-  if tree_altered then
-    renderer.draw()
-  end
-  if index and view.is_visible() then
-    view.set_cursor { index, 0 }
-  end
   running[fname] = false
 
   log.profile_end(ps, "find file %s", fname)

lua/nvim-tree/git/init.lua

@@ -3,6 +3,7 @@ local utils = require "nvim-tree.utils"
 local git_utils = require "nvim-tree.git.utils"
 local Runner = require "nvim-tree.git.runner"
 local Watcher = require("nvim-tree.watcher").Watcher
+local Iterator = require "nvim-tree.iterators.node-iterator"
 
 local M = {
   config = {},
@@ -84,21 +85,19 @@ local function reload_tree_at(project_root)
   local project_files = project.files and project.files or {}
   local project_dirs = project.dirs and project.dirs or {}
 
-  local function iterate(n)
-    local parent_ignored = n.git_status == "!!"
-    for _, node in pairs(n.nodes) do
+  Iterator.builder(root_node.nodes)
+    :hidden()
+    :applier(function(node)
+      local parent_ignored = node.parent.git_status == "!!"
       node.git_status = project_dirs[node.absolute_path] or project_files[node.absolute_path]
       if not node.git_status and parent_ignored then
         node.git_status = "!!"
       end
-
-      if node.nodes and #node.nodes > 0 then
-        iterate(node)
-      end
-    end
-  end
-
-  iterate(root_node)
+    end)
+    :recursor(function(node)
+      return node.nodes and #node.nodes > 0 and node.nodes
+    end)
+    :iterate()
 
   require("nvim-tree.renderer").draw()
 end

lua/nvim-tree/iterators/node-iterator.lua

@@ -0,0 +1,67 @@
+local NodeIterator = {}
+NodeIterator.__index = NodeIterator
+
+function NodeIterator.builder(nodes)
+  return setmetatable({
+    nodes = nodes,
+    _filter_hidden = function(node)
+      return not node.hidden
+    end,
+    _apply_fn_on_node = function(_) end,
+    _match = function(_) end,
+    _recurse_with = function(node)
+      return node.nodes
+    end,
+  }, NodeIterator)
+end
+
+function NodeIterator:hidden()
+  self._filter_hidden = function(_)
+    return true
+  end
+  return self
+end
+
+function NodeIterator:matcher(f)
+  self._match = f
+  return self
+end
+
+function NodeIterator:applier(f)
+  self._apply_fn_on_node = f
+  return self
+end
+
+function NodeIterator:recursor(f)
+  self._recurse_with = f
+  return self
+end
+
+function NodeIterator:iterate()
+  local function iter(nodes)
+    local i = 1
+    for _, node in ipairs(nodes) do
+      if self._filter_hidden(node) then
+        if self._match(node) then
+          return node, i
+        end
+        self._apply_fn_on_node(node)
+        local children = self._recurse_with(node)
+        if children then
+          local n, idx = iter(children)
+          i = i + idx
+          if n then
+            return n, i
+          else
+            i = i + 1
+          end
+        end
+      end
+    end
+    return nil, i
+  end
+
+  return iter(self.nodes)
+end
+
+return NodeIterator

lua/nvim-tree/live-filter.lua

@@ -1,6 +1,7 @@
 local a = vim.api
 
 local view = require "nvim-tree.view"
+local Iterator = require "nvim-tree.iterators.node-iterator"
 
 local M = {
   filter = nil,
@@ -11,15 +12,13 @@ local function redraw()
 end
 
 local function reset_filter(node_)
-  local function iterate(n)
-    n.hidden = false
-    if n.nodes then
-      for _, node in pairs(n.nodes) do
-        iterate(node)
-      end
-    end
-  end
-  iterate(node_ or TreeExplorer)
+  node_ = node_ or TreeExplorer
+  Iterator.builder(node_.nodes)
+    :hidden()
+    :applier(function(node)
+      node.hidden = false
+    end)
+    :iterate()
 end
 
 local overlay_bufnr = nil
@@ -47,6 +46,8 @@ function M.apply_filter(node_)
     return
   end
 
+  -- TODO(kiyan): this iterator cannot yet be refactored with the Iterator module
+  -- since the node mapper is based on its children
   local function iterate(node)
     local filtered_nodes = 0
     local nodes = node.group_next and { node.group_next } or node.nodes

lua/nvim-tree/utils.lua

@@ -3,6 +3,8 @@ local has_notify, notify = pcall(require, "notify")
 local a = vim.api
 local uv = vim.loop
 
+local Iterator = require "nvim-tree.iterators.node-iterator"
+
 local M = {
   debouncers = {},
 }
@@ -103,27 +105,12 @@ end
 -- @param nodes list of node
 -- @param fn    function(node): boolean
 function M.find_node(nodes, fn)
-  local function iter(nodes_, fn_)
-    local i = 1
-    for _, node in ipairs(nodes_) do
-      if not node.hidden then
-        if fn_(node) then
-          return node, i
-        end
-        if node.open and #node.nodes > 0 then
-          local n, idx = iter(node.nodes, fn_)
-          i = i + idx
-          if n then
-            return n, i
-          end
-        else
-          i = i + 1
-        end
-      end
-    end
-    return nil, i
-  end
-  local node, i = iter(nodes, fn)
+  local node, i = Iterator.builder(nodes)
+    :matcher(fn)
+    :recursor(function(node)
+      return node.open and #node.nodes > 0 and node.nodes
+    end)
+    :iterate()
   i = require("nvim-tree.view").is_root_folder_visible() and i or i - 1
   i = require("nvim-tree.live-filter").filter and i + 1 or i
   return node, i
@@ -137,27 +124,20 @@ function M.get_node_from_path(path)
     return explorer
   end
 
-  local function iterate(nodes)
-    for _, node in pairs(nodes) do
-      if node.absolute_path == path or node.link_to == path then
-        return node
-      end
+  return Iterator.builder(explorer.nodes)
+    :hidden()
+    :matcher(function(node)
+      return node.absolute_path == path or node.link_to == path
+    end)
+    :recursor(function(node)
       if node.nodes then
-        local res = iterate(node.nodes)
-        if res then
-          return res
-        end
+        return node.nodes
       end
       if node.group_next then
-        local res = iterate { node.group_next }
-        if res then
-          return res
-        end
+        return { node.group_next }
       end
-    end
-  end
-
-  return iterate(explorer.nodes)
+    end)
+    :iterate()
 end
 
 -- get the highest parent of grouped nodes
@@ -176,21 +156,17 @@ end
 function M.get_nodes_by_line(nodes_all, line_start)
   local nodes_by_line = {}
   local line = line_start
-  local function iter(nodes)
-    for _, node in ipairs(nodes) do
-      if not node.hidden then
-        nodes_by_line[line] = node
-        line = line + 1
-        if node.open == true then
-          local child = iter(node.nodes)
-          if child ~= nil then
-            return child
-          end
-        end
-      end
-    end
-  end
-  iter(nodes_all)
+
+  Iterator.builder(nodes_all)
+    :applier(function(node)
+      nodes_by_line[line] = node
+      line = line + 1
+    end)
+    :recursor(function(node)
+      return node.open == true and node.nodes
+    end)
+    :iterate()
+
   return nodes_by_line
 end