refactor(filesystem): fix recursive expand of nodes (#957)

Co-authored-by: ghostbuster91 <ghostbuster91@users.noreply.github.com>

lua/neo-tree/git/ignored.lua

@@ -118,7 +118,7 @@ M.mark_ignored = function(state, items, callback)
         on_exit = function(self, code, _)
           local result
           if code ~= 0 then
-            log.debug("Failed to load ignored files for", state.path, ":", self:stderr_result())
+            log.debug("Failed to load ignored files for", folder, ":", self:stderr_result())
             result = {}
           else
             result = self:result()

lua/neo-tree/sources/common/commands.lua

@@ -10,6 +10,8 @@ local popups = require("neo-tree.ui.popups")
 local log = require("neo-tree.log")
 local help = require("neo-tree.sources.common.help")
 local Preview = require("neo-tree.sources.common.preview")
+local async = require("plenary.async")
+local node_expander = require("neo-tree.sources.common.node_expander")
 
 ---Gets the node parent folder
 ---@param state table to look for nodes
@@ -108,34 +110,28 @@ M.add_directory = function(state, callback)
   fs_actions.create_directory(in_directory, callback, using_root_directory)
 end
 
-M.expand_all_nodes = function(state, toggle_directory)
-  if toggle_directory == nil then
-    toggle_directory = function(_, node)
-      node:expand()
-    end
+---Expand all nodes
+---@param state table The state of the source
+---@param node table A node to expand
+---@param prefetcher table an object with two methods `prefetch(state, node)` and `should_prefetch(node) => boolean`
+M.expand_all_nodes = function(state, node, prefetcher)
+  log.debug("Expanding all nodes under " .. node:get_id())
+  if prefetcher == nil then
+    prefetcher = node_expander.default_prefetcher
   end
 
-  local expand_node
-  expand_node = function(node)
-    local id = node:get_id()
-    if node.type == "directory" and not node:is_expanded() then
-      toggle_directory(state, node)
-      node = state.tree:get_node(id)
-    end
-    local children = state.tree:get_nodes(id)
-    if children then
-      for _, child in ipairs(children) do
-        if child.type == "directory" then
-          expand_node(child)
-        end
-      end
-    end
-  end
+  renderer.position.set(state, nil)
 
-  for _, node in ipairs(state.tree:get_nodes()) do
-    expand_node(node)
+  local task = function ()
+    node_expander.expand_directory_recursively(state, node, prefetcher)
   end
-  renderer.redraw(state)
+  async.run(
+      task,
+      function ()
+        log.debug("All nodes expanded - redrawing")
+        renderer.redraw(state)
+      end
+    )
 end
 
 M.close_node = function(state, callback)

lua/neo-tree/sources/common/node_expander.lua

@@ -0,0 +1,82 @@
+local log = require("neo-tree.log")
+
+local M = {}
+
+--- Recursively expand all loaded nodes under the given node
+--- returns table with all discovered nodes that need to be loaded 
+---@param node table a node to expand
+---@param state table current state of the source
+---@return table discovered nodes that need to be loaded
+local function expand_loaded(node, state, prefetcher)
+    local function rec(current_node, to_load)
+      if prefetcher.should_prefetch(current_node) then
+        log.trace("Node " .. current_node:get_id() .. "not loaded, saving for later")
+        table.insert(to_load, current_node)
+      else
+        if not current_node:is_expanded() then
+          current_node:expand()
+          state.explicitly_opened_directories[current_node:get_id()] = true
+        end
+        local children = state.tree:get_nodes(current_node:get_id())
+        log.debug("Expanding childrens of " .. current_node:get_id())
+        for _, child in ipairs(children) do
+          if child.type == "directory" then
+             rec(child, to_load)
+          else
+            log.trace("Child: " .. child.name .. " is not a directory, skipping")
+          end
+        end
+      end
+    end
+
+    local to_load = {}
+    rec(node, to_load)
+    return to_load
+end
+
+--- Recursively expands all nodes under the given node collecting all unloaded nodes
+--- Then run prefetcher on all unloaded nodes. Finally, expand loded nodes. 
+--- async method
+---@param node table a node to expand
+---@param state table current state of the source
+local function expand_and_load(node, state, prefetcher)
+      local to_load = expand_loaded(node, state, prefetcher)
+      for _, _node in ipairs(to_load) do
+        prefetcher.prefetch(state, _node)
+        -- no need to handle results as prefetch is recursive
+        expand_loaded(_node, state, prefetcher)
+      end
+end
+
+--- Expands given node recursively loading all descendant nodes if needed
+--- Nodes will be loaded using given prefetcher
+--- async method
+---@param state table current state of the source
+---@param node table a node to expand
+---@param prefetcher table an object with two methods `prefetch(state, node)` and `should_prefetch(node) => boolean`
+M.expand_directory_recursively = function(state, node, prefetcher)
+  log.debug("Expanding directory " .. node:get_id())
+  if node.type ~= "directory" then
+    return
+  end
+  state.explicitly_opened_directories = state.explicitly_opened_directories or {}
+  if prefetcher.should_prefetch(node) then
+    local id = node:get_id()
+    state.explicitly_opened_directories[id] = true
+    prefetcher.prefetch(state, node)
+    expand_loaded(node, state, prefetcher)
+  else
+    expand_and_load(node, state, prefetcher)
+  end
+end
+
+M.default_prefetcher = {
+  prefetch = function (state, node)
+    log.debug("Default expander prefetch does nothing")
+  end,
+  should_prefetch = function (node)
+    return false
+  end
+}
+
+return M

lua/neo-tree/sources/filesystem/commands.lua

@@ -67,11 +67,11 @@ M.delete_visual = function(state, selected_nodes)
   cc.delete_visual(state, selected_nodes, utils.wrap(refresh, state))
 end
 
-M.expand_all_nodes = function(state)
-  local toggle_dir_no_redraw = function(_state, node)
-    fs.toggle_directory(_state, node, nil, true, true)
+M.expand_all_nodes = function(state, node)
+  if node == nil then
+    node = state.tree:get_node(state.path)
   end
-  cc.expand_all_nodes(state, toggle_dir_no_redraw)
+  cc.expand_all_nodes(state, node, fs.prefetcher)
 end
 
 ---Shows the filter input, which will filter the tree.

lua/neo-tree/sources/filesystem/init.lua

@@ -393,7 +393,7 @@ M.setup = function(config, global_config)
 end
 
 ---Expands or collapses the current node.
-M.toggle_directory = function(state, node, path_to_reveal, skip_redraw, recursive)
+M.toggle_directory = function(state, node, path_to_reveal, skip_redraw, recursive, callback)
   local tree = state.tree
   if not node then
     node = tree:get_node()
@@ -406,7 +406,7 @@ M.toggle_directory = function(state, node, path_to_reveal, skip_redraw, recursiv
     local id = node:get_id()
     state.explicitly_opened_directories[id] = true
     renderer.position.set(state, nil)
-    fs_scan.get_items(state, id, path_to_reveal, nil, false, recursive)
+    fs_scan.get_items(state, id, path_to_reveal, callback, false, recursive)
   elseif node:has_children() then
     local updated = false
     if node:is_expanded() then
@@ -428,4 +428,14 @@ M.toggle_directory = function(state, node, path_to_reveal, skip_redraw, recursiv
   end
 end
 
+M.prefetcher = {
+  prefetch = function (state, node)
+    log.debug("Running fs prefetch for: " .. node:get_id())
+    fs_scan.get_dir_items_async(state, node:get_id(), true)
+  end,
+  should_prefetch = function (node)
+    return not node.loaded
+  end
+}
+
 return M