reworked tree init

3 files changed, 81 insertions, 65 deletions

diff --git a/src/file/order.rs b/src/file/order.rs
index c982147..5004e02 100644
--- a/src/file/order.rs
+++ b/src/file/order.rs
@@ -6,19 +6,19 @@ use std::cmp::Ordering;
 pub type FileComparator = dyn Fn(&File, &File) -> Ordering;
 
 /// Yields function pointer to the appropriate `File` comparator.
-pub fn comparator(Context { sort, dir_order, .. }: &Context) -> Option<Box<FileComparator>> {
+pub fn comparator(sort: Sort, dir_order: DirOrder) -> Option<Box<FileComparator>> {
     if matches!(sort, Sort::None) {
         return None;
     }
 
     match dir_order {
         DirOrder::First => {
-            Some(Box::new(move |a, b| dir_first_comparator(a, b, base_comparator(*sort))))
+            Some(Box::new(move |a, b| dir_first_comparator(a, b, base_comparator(sort))))
         },
         DirOrder::Last => {
-            Some(Box::new(move |a, b| dir_last_comparator(a, b, base_comparator(*sort))))
+            Some(Box::new(move |a, b| dir_last_comparator(a, b, base_comparator(sort))))
         },
-        DirOrder::None => Some(base_comparator(*sort)),
+        DirOrder::None => Some(base_comparator(sort)),
     }
 }
 
diff --git a/src/file/tree/mod.rs b/src/file/tree/mod.rs
index 6a6888a..f296fe6 100644
--- a/src/file/tree/mod.rs
+++ b/src/file/tree/mod.rs
@@ -43,72 +43,88 @@ impl Tree {
             root_id,
         } = Self::load(ctx)?;
 
-        let mut dfs_stack = vec![root_id];
+        let mut dir_stack = vec![root_id];
         let mut inode_set = HashSet::default();
 
-        'outer: while let Some(node_id) = dfs_stack.last() {
-            let current_id = *node_id;
+        // Keeps track of which directory entry we're at for each directory while doing depth first
+        // traversal. The key is the NodeID of a directory with the value being an index into a
+        // slice of the directory's children.
+        let mut dirent_cursor_map = HashMap::<NodeId, usize>::default();
 
-            let current_node_path = arena[current_id].get().path();
+        // Map of dynamically computed directory sizes.
+        let mut dirsize_map = HashMap::<NodeId, u64>::default();
 
-            let Some(children) = branches.get_mut(current_node_path) else {
-                dfs_stack.pop();
+        'outer: while let Some(node_id) = dir_stack.last() {
+            let current_dir = *node_id;
+
+            let current_node_path = arena[current_dir].get().path();
+
+            let Some(dirents) = branches.get_mut(current_node_path) else {
+                dir_stack.pop();
                 continue;
             };
 
-            while let Some(child_node_id) = children.pop() {
-                current_id.append(child_node_id, &mut arena);
-
-                let (child_size, child_is_dir, child_inode) = {
-                    let child_node = arena[child_node_id].get();
-                    let is_dir = child_node.file_type().is_some_and(|f| f.is_dir());
-                    let size = child_node.size().value();
-                    let inode = match child_node.inode() {
-                        Ok(value) => {
-                            #[cfg(unix)]
-                            column_metadata.update_inode_attr_widths(&value);
-                            value
-                        },
-                        Err(err) => {
-                            log::warn!(
-                                "Failed to query inode of {} which may affect disk usage report: {err}",
-                                child_node.path().display(),
-                            );
-                            continue;
-                        },
-                    };
-                    (size, is_dir, inode)
-                };
+            let current_dirsize = dirsize_map.entry(current_dir).or_insert(0);
+            let dirent_cursor = dirent_cursor_map.entry(current_dir).or_insert(0);
 
-                if child_is_dir {
-                    dfs_stack.push(child_node_id);
+            for dirent_node_id in &dirents[*dirent_cursor..] {
+                *dirent_cursor += 1;
+                let dirent_node_id = *dirent_node_id;
+                let dirent_node = arena[dirent_node_id].get();
+
+                if dirent_node.is_dir() {
+                    dir_stack.push(dirent_node_id);
                     continue 'outer;
                 }
 
-                if inode_set.insert(child_inode) {
-                    *arena[current_id].get_mut().size_mut() += child_size;
-                }
+                match dirent_node.inode() {
+                    Ok(inode) => {
+                        #[cfg(unix)]
+                        column_metadata.update_inode_attr_widths(&inode);
 
-                *arena[current_id].get_mut().size_mut() += inode_set
-                    .insert(child_inode)
-                    .then_some(child_size)
-                    .unwrap_or(0);
+                        *current_dirsize += inode_set
+                            .insert(inode)
+                            .then(|| dirent_node.size().value())
+                            .unwrap_or(0);
+                    },
+                    Err(err) => {
+                        log::warn!(
+                            "Failed to query inode of {} which may affect disk usage report: {err}",
+                            dirent_node.path().display(),
+                        );
+                    },
+                };
             }
 
-            dfs_stack.pop();
+            dir_stack.pop();
+
+            // To play nicely with aliasing rules around mutable refs
+            let current_dirsize = *current_dirsize;
 
-            if let Some(parent_id) = current_id.ancestors(&arena).nth(1) {
-                let current_dir_size = { arena[current_id].get().size().value() };
-                *arena[parent_id].get_mut().size_mut() += current_dir_size;
+            if let Some(parent_dir) = dir_stack.last() {
+                if let Some(parent_dirsize) = dirsize_map.get_mut(parent_dir) {
+                    *parent_dirsize += current_dirsize;
+                }
             }
         }
 
-        column_metadata.update_size_width(arena[root_id].get(), ctx);
+        for (dir_id, dirsize) in dirsize_map.into_iter() {
+            let dir = arena[dir_id].get_mut();
+            *dir.size_mut() += dirsize;
 
-        if let Some(comparator) = order::comparator(ctx) {
-            Self::tree_sort(root_id, &mut arena, comparator);
+            if let Some(dirents) = branches.remove(dir.path()) {
+                for dirent_id in dirents {
+                    dir_id.append(dirent_id, &mut arena);
+                }
+            }
         }
 
+        column_metadata.update_size_width(arena[root_id].get(), ctx);
+
+        //if let Some(comparator) = order::comparator(ctx.sort, ctx.dir_order) {
+            //Self::tree_sort(root_id, &mut arena, comparator);
+        //}
+
         let tree = Self { root_id, arena };
 
         Ok((tree, column_metadata))
@@ -122,25 +138,25 @@ impl Tree {
             root_id,
         } = Self::load(ctx)?;
 
-        let mut dfs_stack = vec![root_id];
+        let mut dir_stack = vec![root_id];
 
-        'outer: while let Some(node_id) = dfs_stack.last() {
-            let current_id = *node_id;
+        'outer: while let Some(node_id) = dir_stack.last() {
+            let current_dir = *node_id;
 
-            let current_node_path = arena[current_id].get().path();
+            let current_node_path = arena[current_dir].get().path();
 
-            let Some(children) = branches.get_mut(current_node_path) else {
-                dfs_stack.pop();
+            let Some(dirents) = branches.get_mut(current_node_path) else {
+                dir_stack.pop();
                 continue;
             };
 
-            while let Some(child_node_id) = children.pop() {
-                current_id.append(child_node_id, &mut arena);
+            while let Some(dirent_node_id) = dirents.pop() {
+                current_dir.append(dirent_node_id, &mut arena);
 
-                let child_node = arena[child_node_id].get();
+                let dirent_node = arena[dirent_node_id].get();
 
                 #[cfg(unix)]
-                match child_node.inode() {
+                match dirent_node.inode() {
                     Ok(value) => {
                         column_metadata.update_inode_attr_widths(&value);
                         value
@@ -148,23 +164,23 @@ impl Tree {
                     Err(err) => {
                         log::warn!(
                             "Failed to query inode of {}: {err}",
-                            child_node.path().display(),
+                            dirent_node.path().display(),
                         );
                         continue;
                     },
                 };
 
-                if child_node.file_type().is_some_and(|f| f.is_dir()) {
-                    dfs_stack.push(child_node_id);
+                if dirent_node.file_type().is_some_and(|f| f.is_dir()) {
+                    dir_stack.push(dirent_node_id);
                     continue 'outer;
                 }
             }
 
-            dfs_stack.pop();
+            dir_stack.pop();
         }
 
         if !matches!(ctx.sort, Sort::Size | Sort::Rsize) {
-            if let Some(comparator) = order::comparator(ctx) {
+            if let Some(comparator) = order::comparator(ctx.sort, ctx.dir_order) {
                 Self::tree_sort(root_id, &mut arena, comparator);
             }
         }
diff --git a/src/render/mod.rs b/src/render/mod.rs
index 24a5358..8b66539 100644
--- a/src/render/mod.rs
+++ b/src/render/mod.rs
@@ -168,7 +168,7 @@ pub fn inverted_tree(file_tree: &file::Tree, ctx: &Context) -> Result<String> {
             log::warn!("{e}");
         }
 
-        if utils::node_is_dir(node) && depth < max_depth {
+        if node.is_dir() && depth < max_depth {
             if is_last_sibling(node_id, depth) {
                 inherited_prefix_components.push(SEP);
             } else {