Rewrite recursive expand to use plenary.async

Author: ghostbuster91

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

@@ -2,6 +2,7 @@
 
 local vim = vim
 local fs_actions = require("neo-tree.sources.filesystem.lib.fs_actions")
+local fs = require("neo-tree.sources.filesystem")
 local utils = require("neo-tree.utils")
 local renderer = require("neo-tree.ui.renderer")
 local events = require("neo-tree.events")
@@ -10,6 +11,7 @@ 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")
 
 ---Gets the node parent folder
 ---@param state table to look for nodes
@@ -108,35 +110,23 @@ 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
-  end
-  --state.explicitly_opened_directories = state.explicitly_opened_directories or {}
-
-  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
+---Expand all nodes
+---@param state table The state of the source
+---@param node table A node to expand
+M.expand_all_nodes = function(state, node)
+  log.debug("Expanding all nodes under " .. node:get_id())
+  local task = function ()
+    fs.expand_directory(state, node)
+  end
+  async.run(
+      task,
+      function ()
+        log.debug("All nodes expanded - redrawing")
+        vim.schedule_wrap(function()
+          renderer.redraw(state)
+        end)
       end
-    end
-  end
-
-  for _, node in ipairs(state.tree:get_nodes()) do
-    expand_node(node)
-  end
-  renderer.redraw(state)
+    )
 end
 
 M.close_node = function(state, callback)

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

@@ -68,10 +68,7 @@ M.delete_visual = function(state, selected_nodes)
 end
 
 M.expand_all_nodes = function(state)
-  local toggle_dir_no_redraw = function(_state, node)
-    fs.toggle_directory(_state, node, nil, true, true)
-  end
-  cc.expand_all_nodes(state, toggle_dir_no_redraw)
+  cc.expand_all_nodes(state, state.tree:get_node(state.path))
 end
 
 ---Shows the filter input, which will filter the tree.

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

@@ -10,6 +10,7 @@ local log = require("neo-tree.log")
 local manager = require("neo-tree.sources.manager")
 local git = require("neo-tree.git")
 local glob = require("neo-tree.sources.filesystem.lib.globtopattern")
+local async = require("plenary.async")
 
 local M = {
   name = "filesystem",
@@ -428,4 +429,76 @@ M.toggle_directory = function(state, node, path_to_reveal, skip_redraw, recursiv
   end
 end
 
+--- 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)
+    local function rec(_node, to_load)
+      if _node.loaded == false then
+        table.insert(to_load, _node)
+      else
+        if not _node:is_expanded() then
+          _node:expand()
+          state.explicitly_opened_directories[_node:get_id()] = true
+        end
+        local children = state.tree:get_nodes(_node:get_id())
+        log.debug("Expanding childrens of " .. _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
+--- loading nodes if necessary.
+--- asyn method
+---@param node table a node to expand
+---@param state table current state of the source
+local function expand_and_load(node, state)
+    local function rec(to_load, progress)
+      local to_load_current = expand_loaded(node, state)
+      for _,v in ipairs(to_load_current) do
+        table.insert(to_load, v)
+      end
+      if progress <= #to_load then
+        M.expand_directory(state, to_load[progress])
+        rec(to_load, progress + 1)
+      end
+    end
+    rec({}, 1)
+end
+
+---Expands given node recursively loading all descendant nodes if needed
+---async method
+---@param state table current state of the source
+---@param node table a node to expand
+M.expand_directory = function(state, node)
+  log.debug("Expanding directory " .. node:get_id())
+  if node.type ~= "directory" then
+    return
+  end
+  state.explicitly_opened_directories = state.explicitly_opened_directories or {}
+  if node.loaded == false then
+    local id = node:get_id()
+    state.explicitly_opened_directories[id] = true
+    renderer.position.set(state, nil)
+    fs_scan.get_dir_items_async(state, id, true)
+    -- ignore results as we know here that all descendant nodes have been already loaded
+    expand_loaded(node ,state)
+  else
+    expand_and_load(node, state)
+  end
+end
+
 return M

lua/neo-tree/sources/filesystem/lib/fs_scan.lua

@@ -106,6 +106,26 @@ local render_context = function(context)
   context = nil
 end
 
+local job_complete_async = function(context)
+  local state = context.state
+  local parent_id = context.parent_id
+  if #context.all_items == 0 then
+    log.info("No items, skipping git ignored/status lookups")
+  elseif state.filtered_items.hide_gitignored or state.enable_git_status then
+    local mark_ignored_async = async.wrap(function (_state, _all_items, _callback)
+      git.mark_ignored(_state, _all_items, _callback)
+    end, 3)
+    local all_items = mark_ignored_async(state, context.all_items)
+
+    if parent_id then
+      vim.list_extend(state.git_ignored, all_items)
+    else
+      state.git_ignored = all_items
+    end
+  end
+  return context
+end
+
 local job_complete = function(context)
   local state = context.state
   local parent_id = context.parent_id
@@ -135,8 +155,8 @@ local job_complete = function(context)
         state.git_ignored = all_items
       end
     end
+    render_context(context)
   end
-  render_context(context)
 end
 
 local function create_node(context, node)
@@ -199,26 +219,34 @@ local function scan_dir_sync(context, path)
   end
 end
 
-local function scan_dir_async(context, path, callback)
-  get_children_async(path, function(children)
-    for _, child in ipairs(children) do
-      create_node(context, child)
-      if child.type == "directory" then
-        local grandchild_nodes = get_children_sync(child.path)
-        if
-          grandchild_nodes == nil
-          or #grandchild_nodes == 0
-          or #grandchild_nodes == 1 and grandchild_nodes[1].type == "directory"
-        then
-          scan_dir_sync(context, child.path)
-        end
+local function scan_dir_async(context, path)
+  log.debug("scan_dir_async - start " .. path)
+
+  local get_children = async.wrap(function (callback)
+    return get_children_async(path, callback)
+  end, 1)
+
+  local children = get_children()
+  for _, child in ipairs(children) do
+    create_node(context, child)
+    if child.type == "directory" then
+      local grandchild_nodes = get_children_sync(child.path)
+      if
+        grandchild_nodes == nil
+        or #grandchild_nodes == 0
+        or #grandchild_nodes == 1 and grandchild_nodes[1].type == "directory"
+      then
+        scan_dir_sync(context, child.path)
       end
     end
-    process_node(context, path)
-    callback(path)
-  end)
+  end
+
+  log.debug("scan_dir_async - finish " .. path)
+  process_node(context, path)
+  return path
 end
 
+
 -- async_scan scans all the directories in context.paths_to_load
 -- and adds them as items to render in the UI.
 local function async_scan(context, path)
@@ -229,7 +257,9 @@ local function async_scan(context, path)
     local scan_tasks = {}
     for _, p in ipairs(context.paths_to_load) do
       local scan_task = async.wrap(function(callback)
-        scan_dir_async(context, p, callback)
+        async.run(function ()
+            scan_dir_async(context, p)
+        end, callback)
       end, 1)
       table.insert(scan_tasks, scan_task)
     end
@@ -471,7 +501,6 @@ M.get_items = function(state, parent_id, path_to_reveal, callback, async, recurs
   context.path_to_reveal = path_to_reveal
   context.recursive = recursive
   context.callback = callback
-
   -- Create root folder
   local root = file_items.create_item(context, parent_id or state.path, "directory")
   root.name = vim.fn.fnamemodify(root.path, ":~")
@@ -490,6 +519,62 @@ M.get_items = function(state, parent_id, path_to_reveal, callback, async, recurs
   end
 end
 
+-- async method
+M.get_dir_items_async = function(state, parent_id, recursive)
+  local context = file_items.create_context()
+  context.state = state
+  context.parent_id = parent_id
+  context.path_to_reveal = nil
+  context.recursive = recursive
+  context.callback = nil
+  context.paths_to_load = {}
+
+  -- Create root folder
+  local root = file_items.create_item(context, parent_id or state.path, "directory")
+  root.name = vim.fn.fnamemodify(root.path, ":~")
+  root.loaded = true
+  root.search_pattern = state.search_pattern
+  context.root = root
+  context.folders[root.path] = root
+  state.default_expanded_nodes = state.force_open_folders or { state.path }
+
+  local filtered_items = state.filtered_items or {}
+  context.is_a_never_show_file = function(fname)
+    if fname then
+      local _, name = utils.split_path(fname)
+      if name then
+        if filtered_items.never_show and filtered_items.never_show[name] then
+          return true
+        end
+        if utils.is_filtered_by_pattern(filtered_items.never_show_by_pattern, fname, name) then
+          return true
+        end
+      end
+    end
+    return false
+  end
+  table.insert(context.paths_to_load, parent_id)
+ 
+  local scan_tasks = {}
+  for _, p in ipairs(context.paths_to_load) do
+     local scan_task = function ()
+       scan_dir_async(context, p)
+     end
+     table.insert(scan_tasks, scan_task)
+  end
+  async.util.join(scan_tasks)
+
+  job_complete_async(context)
+
+  local finalize =  async.wrap(function (_context, _callback)
+    vim.schedule(function ()
+        render_context(_context)
+        _callback()
+    end)
+  end, 2)
+  finalize(context)
+end
+
 M.stop_watchers = function(state)
   if state.use_libuv_file_watcher and state.tree then
     -- We are loaded a new root or refreshing, unwatch any folders that were