diff --git a/Cargo.lock b/Cargo.lock index 125fd416..22e88ec4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -183,6 +183,7 @@ dependencies = [ "crossbeam", "filesize", "ignore", + "indextree", "indoc", "lscolors", "num_cpus", @@ -284,6 +285,12 @@ dependencies = [ "winapi-util", ] +[[package]] +name = "indextree" +version = "4.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c40411d0e5c63ef1323c3d09ce5ec6d84d71531e18daed0743fccea279d7deb6" + [[package]] name = "indoc" version = "2.0.1" diff --git a/Cargo.toml b/Cargo.toml index 9a2b686f..ddc766fd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -29,6 +29,7 @@ clap_complete = "4.1.1" crossbeam = "0.8.2" filesize = "0.2.0" ignore = "0.4.2" +indextree = "4.6.0" lscolors = { version = "0.13.0", features = ["ansi_term"] } num_cpus = "1.15.0" once_cell = "1.17.0" diff --git a/src/render/context/mod.rs b/src/render/context/mod.rs index f4968016..4f1c1955 100644 --- a/src/render/context/mod.rs +++ b/src/render/context/mod.rs @@ -141,23 +141,6 @@ impl Context { // user arguments. let mut args = vec![OsString::from("--")]; - // Used to pick either from config or user args. - let mut pick_args_from = |id: &str, matches: &ArgMatches| { - if let Ok(Some(raw)) = matches.try_get_raw(id) { - let kebap = id.replace("_", "-"); - - let raw_args = raw - .map(OsStr::to_owned) - .map(|s| vec![OsString::from(format!("--{}", kebap)), s]) - .filter(|pair| pair[1] != "false") - .flatten() - .filter(|s| s != "true") - .collect::>(); - - args.extend(raw_args); - } - }; - let mut ids = user_args.ids().map(Id::as_str).collect::>(); ids.extend(config_args.ids().map(Id::as_str).collect::>()); @@ -165,21 +148,27 @@ impl Context { ids = crate::utils::uniq(ids); for id in ids { - // Don't look at me... my shame.. if id == "Context" { continue; + } else if id == "dir" { + if let Ok(Some(raw)) = user_args.try_get_raw(id) { + let raw_args = raw.map(OsStr::to_owned).collect::>(); + + args.extend(raw_args); + continue; + } } if let Some(user_arg) = user_args.value_source(id) { match user_arg { // prioritize the user arg if user provided a command line argument - ValueSource::CommandLine => pick_args_from(id, &user_args), + ValueSource::CommandLine => Self::pick_args_from(id, &user_args, &mut args), // otherwise prioritize argument from the config - _ => pick_args_from(id, &config_args), + _ => Self::pick_args_from(id, &config_args, &mut args), } } else { - pick_args_from(id, &config_args) + Self::pick_args_from(id, &config_args, &mut args) } } @@ -241,6 +230,23 @@ impl Context { builder.build() } + + /// Used to pick either from config or user args when constructing [Context]. + fn pick_args_from(id: &str, matches: &ArgMatches, args: &mut Vec) { + if let Ok(Some(raw)) = matches.try_get_raw(id) { + let kebap = id.replace("_", "-"); + + let raw_args = raw + .map(OsStr::to_owned) + .map(|s| vec![OsString::from(format!("--{}", kebap)), s]) + .filter(|pair| pair[1] != "false") + .flatten() + .filter(|s| s != "true") + .collect::>(); + + args.extend(raw_args); + } + } } #[derive(Debug)] diff --git a/src/render/disk_usage.rs b/src/render/disk_usage.rs index 5c3fb863..5aa3eed4 100644 --- a/src/render/disk_usage.rs +++ b/src/render/disk_usage.rs @@ -53,7 +53,7 @@ pub enum SiPrefix { } /// Represents either logical or physical size and handles presentation. -#[derive(Debug)] +#[derive(Clone, Debug)] pub struct FileSize { pub bytes: u64, #[allow(dead_code)] @@ -92,9 +92,9 @@ impl FileSize { } } -impl AddAssign<&Self> for FileSize { - fn add_assign(&mut self, rhs: &Self) { - self.bytes += rhs.bytes; +impl AddAssign for FileSize { + fn add_assign(&mut self, rhs: u64) { + self.bytes += rhs; } } diff --git a/src/render/tree/mod.rs b/src/render/tree/mod.rs index 45be38d3..23998db1 100644 --- a/src/render/tree/mod.rs +++ b/src/render/tree/mod.rs @@ -1,7 +1,8 @@ use crate::render::{context::Context, disk_usage::FileSize, order::Order}; use crossbeam::channel::{self, Sender}; use error::Error; -use ignore::{WalkBuilder, WalkParallel, WalkState}; +use ignore::{WalkBuilder, WalkParallel}; +use indextree::{Arena, NodeId}; use node::Node; use std::{ collections::{HashMap, HashSet}, @@ -9,9 +10,9 @@ use std::{ fmt::{self, Display, Formatter}, fs, path::PathBuf, - slice::Iter, thread, }; +use visitor::{BranchVisitorBuilder, TraversalState}; /// Errors related to traversal, [Tree] construction, and the like. pub mod error; @@ -25,150 +26,190 @@ pub mod node; /// [ui::LS_COLORS] initialization and ui theme for [Tree]. pub mod ui; +/// Custom visitor that operates on each thread during filesystem traversal. +mod visitor; + /// In-memory representation of the root-directory and its contents which respects `.gitignore` and /// hidden file rules depending on [WalkParallel] config. #[derive(Debug)] pub struct Tree { - root: Node, + inner: Arena, + root: NodeId, ctx: Context, } pub type TreeResult = Result; -pub type Branches = HashMap>; -pub type TreeComponents = (Node, Branches); impl Tree { /// Constructor for [Tree]. - pub fn new(root: Node, ctx: Context) -> Self { - Self { root, ctx } + pub fn new(inner: Arena, root: NodeId, ctx: Context) -> Self { + Self { inner, root, ctx } } /// Initiates file-system traversal and [Tree construction]. pub fn init(ctx: Context) -> TreeResult { - let root = Self::traverse(&ctx)?; - - Ok(Self::new(root, ctx)) - } + let (inner, root) = Self::traverse(&ctx)?; - /// Returns a reference to the root [Node]. - fn root(&self) -> &Node { - &self.root + Ok(Self::new(inner, root, ctx)) } - /// Maximum depth to display + /// Maximum depth to display. fn level(&self) -> usize { self.ctx.level.unwrap_or(usize::MAX) } + /// Grab a reference to [Context]. fn context(&self) -> &Context { &self.ctx } - /// Parallel traversal of the root directory and its contents taking `.gitignore` into - /// consideration. Parallel traversal relies on `WalkParallel`. Any filesystem I/O or related - /// system calls are expected to occur during parallel traversal; thus post-processing of all - /// directory entries should be completely CPU-bound. If filesystem I/O or system calls occur - /// outside of the parallel traversal step please report an issue. - fn traverse(ctx: &Context) -> TreeResult { + /// Grabs a reference to `inner`. + fn inner(&self) -> &Arena { + &self.inner + } + + /// Parallel traversal of the root directory and its contents. Parallel traversal relies on + /// `WalkParallel`. Any filesystem I/O or related system calls are expected to occur during + /// parallel traversal; post-processing post-processing of all directory entries should + /// be completely CPU-bound. + fn traverse(ctx: &Context) -> TreeResult<(Arena, NodeId)> { let walker = WalkParallel::try_from(ctx)?; - let (tx, rx) = channel::unbounded::(); + let (tx, rx) = channel::unbounded::(); - // Receives directory entries from the workers used for parallel traversal to construct the - // components needed to assemble a `Tree`. - let tree_components = thread::spawn(move || -> TreeResult { - let mut branches: Branches = HashMap::new(); - let mut inodes = HashSet::new(); - let mut root = None; + thread::scope(|s| { + let mut tree = Arena::new(); - while let Ok(node) = rx.recv() { - if node.is_dir() { - let node_path = node.path(); + let res = s.spawn(|| { + // Key represents path of parent directory and values represent children. + let mut branches: HashMap> = HashMap::new(); - if !branches.contains_key(node_path) { - branches.insert(node_path.to_owned(), vec![]); - } + // Set used to prevent double counting hard-links in the same file-tree hiearchy. + let mut inodes = HashSet::new(); - if node.depth == 0 { - root = Some(node); - continue; - } - } + let mut root_id = None; - if let Some(inode) = node.inode() { - if inode.nlink > 1 { - // If a hard-link is already accounted for skip the subsequent one. - if !inodes.insert(inode.properties()) { + while let Ok(TraversalState::Ongoing(node)) = rx.recv() { + if node.is_dir() { + let node_path = node.path(); + + if !branches.contains_key(node_path) { + branches.insert(node_path.to_owned(), vec![]); + } + + if node.depth == 0 { + root_id = Some(tree.new_node(node)); continue; } } - } - let parent = node.parent_path_buf().ok_or(Error::ExpectedParent)?; + // If a hard-link is already accounted for, skip all subsequent ones. + if let Some(inode) = node.inode() { + if inode.nlink > 1 { + if !inodes.insert(inode.properties()) { + continue; + } + } + } - let update = branches.get_mut(&parent).map(|mut_ref| mut_ref.push(node)); + let parent = node.parent_path().ok_or(Error::ExpectedParent)?.to_owned(); - if update.is_none() { - branches.insert(parent, vec![]); - } - } + let node_id = tree.new_node(node); - let root_node = root.ok_or(Error::MissingRoot)?; + if let None = branches + .get_mut(&parent) + .map(|mut_ref| mut_ref.push(node_id)) + { + branches.insert(parent, vec![]); + } + } - Ok((root_node, branches)) - }); + let root = root_id.ok_or(Error::MissingRoot)?; - // All filesystem I/O and related system-calls should be relegated to this. Directory - // entries that are encountered are sent to the above thread for processing. - walker.run(|| { - Box::new(|entry_res| { - let tx = Sender::clone(&tx); + Self::assemble_tree(&mut tree, root, &mut branches, ctx); - entry_res - .map(|entry| Node::from((&entry, ctx))) - .map(|node| tx.send(node).unwrap()) - .map(|_| WalkState::Continue) - .unwrap_or(WalkState::Skip) - }) - }); + if ctx.prune { + Self::prune_directories(root, &mut tree); + } - drop(tx); + Ok::<(Arena, NodeId), Error>((tree, root)) + }); - let (mut root, mut branches) = tree_components.join().unwrap()?; + let mut visitor_builder = BranchVisitorBuilder::new(ctx, Sender::clone(&tx)); - Self::assemble_tree(&mut root, &mut branches, ctx); + walker.visit(&mut visitor_builder); - if ctx.prune { - root.prune_directories() - } + tx.send(TraversalState::Done).unwrap(); - Ok(root) + res.join().unwrap() + }) } /// Takes the results of the parallel traversal and uses it to construct the [Tree] data /// structure. Sorting occurs if specified. - fn assemble_tree(current_node: &mut Node, branches: &mut Branches, ctx: &Context) { - let children = branches.remove(current_node.path()).unwrap(); - - current_node.set_children(children); + fn assemble_tree( + tree: &mut Arena, + current_node_id: NodeId, + branches: &mut HashMap>, + ctx: &Context, + ) { + let mut children = Node::get(current_node_id, tree) + .map(|n| branches.remove(n.path()).unwrap()) + .unwrap(); let mut dir_size = FileSize::new(0, ctx.disk_usage, ctx.prefix, ctx.scale); - current_node.children_mut().for_each(|node| { - if node.is_dir() { - Self::assemble_tree(node, branches, ctx); - } + for child_id in children.iter() { + let is_dir = { + let inner = Node::get(*child_id, tree).unwrap(); + inner.is_dir() + }; - if let Some(fs) = node.file_size() { - dir_size += fs + if is_dir { + Self::assemble_tree(tree, *child_id, branches, ctx); } - }); + + let inner = Node::get(*child_id, tree).unwrap(); + let file_size = inner.file_size().map(|fs| fs.bytes).unwrap_or(0); + + dir_size += file_size; + } if dir_size.bytes > 0 { - current_node.set_file_size(dir_size) + let current_node = Node::get_mut(current_node_id, tree).unwrap(); + current_node.set_file_size(dir_size); } + // Sort if sorting specified if let Some(func) = Order::from((ctx.sort(), ctx.dirs_first())).comparator() { - current_node.sort_children(func) + children.sort_by(|id_a, id_b| { + let node_a = Node::get(*id_a, tree).unwrap(); + let node_b = Node::get(*id_b, tree).unwrap(); + func(node_a, node_b) + }); + } + + // Append children to current node. + for child_id in children { + current_node_id.append(child_id, tree); + } + } + + /// Function to remove empty directories. + fn prune_directories(root_id: NodeId, tree: &mut Arena) { + let mut to_prune = vec![]; + + for node_id in root_id.descendants(tree) { + let node = Node::get(node_id, tree).unwrap(); + + if node.is_dir() { + if node_id.children(tree).peekable().peek().is_none() { + to_prune.push(node_id); + } + } + } + + for node_id in to_prune { + node_id.remove_subtree(tree) } } } @@ -193,75 +234,79 @@ impl TryFrom<&Context> for WalkParallel { impl Display for Tree { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { - let root = self.root(); + let root = self.root; + let inner = self.inner(); let level = self.level(); - let theme = ui::get_theme(); - let ctx = self.context(); - fn extend_output( - f: &mut Formatter, + + let mut descendants = root.descendants(self.inner()).skip(1).peekable(); + + let root_node = Node::get(root, inner).unwrap(); + + fn display_node( node: &Node, - prefix: &str, + base_prefix: &str, ctx: &Context, + f: &mut Formatter<'_>, ) -> fmt::Result { if ctx.size_left && !ctx.suppress_size { - node.display_size_left(f, prefix, ctx)?; - writeln!(f, "") + node.display_size_left(f, base_prefix, ctx)?; } else { - node.display_size_right(f, prefix, ctx)?; - writeln!(f, "") + node.display_size_right(f, base_prefix, ctx)?; } + + writeln!(f, "") } - fn traverse( - f: &mut Formatter, - children: Iter, - base_prefix: &str, - level: usize, - theme: &ui::ThemesMap, - ctx: &Context, - ) -> fmt::Result { - let mut peekable = children.peekable(); + display_node(&root_node, "", ctx, f)?; - while let Some(child) = peekable.next() { - let last_entry = peekable.peek().is_none(); - let prefix = base_prefix.to_owned() - + if last_entry { - theme.get("uprt").unwrap() - } else { - theme.get("vtrt").unwrap() - }; + let mut prefix_components = vec!["".to_owned()]; - extend_output(f, child, &prefix, ctx)?; + while let Some(current_node_id) = descendants.next() { + let mut current_prefix_components = prefix_components.clone(); - if !child.is_dir() || child.depth + 1 > level { - continue; - } + let current_node = Node::get(current_node_id, inner).unwrap(); + + let theme = if current_node.is_symlink() { + ui::get_link_theme() + } else { + ui::get_theme() + }; - if child.has_children() { - let children = child.children(); + let mut siblings = current_node_id.following_siblings(inner).skip(1).peekable(); - let new_theme = if child.is_symlink() { - ui::get_link_theme() - } else { - theme - }; + let last_sibling = siblings.peek().is_none(); - let new_base = base_prefix.to_owned() - + if last_entry { - ui::SEP - } else { - theme.get("vt").unwrap() - }; + if last_sibling { + current_prefix_components.push(theme.get("uprt").unwrap().to_owned()); + } else { + current_prefix_components.push(theme.get("vtrt").unwrap().to_owned()); + } - traverse(f, children, &new_base, level, new_theme, ctx)?; + let prefix = current_prefix_components.join(""); + + if current_node.depth <= level { + display_node(¤t_node, &prefix, ctx, f)?; + } + + if let Some(next_id) = descendants.peek() { + let next_node = Node::get(*next_id, inner).unwrap(); + + if next_node.depth == current_node.depth + 1 { + if last_sibling { + prefix_components.push(ui::SEP.to_owned()); + } else { + prefix_components.push(theme.get("vt").unwrap().to_owned()); + } + } else if next_node.depth < current_node.depth { + let depth_delta = current_node.depth - next_node.depth; + for _ in 0..depth_delta { + prefix_components.pop(); + } } } - Ok(()) } - extend_output(f, root, "", ctx)?; - traverse(f, root.children(), "", level, theme, ctx)?; Ok(()) } } diff --git a/src/render/tree/node.rs b/src/render/tree/node.rs index b79dccd2..86b71b6e 100644 --- a/src/render/tree/node.rs +++ b/src/render/tree/node.rs @@ -5,12 +5,12 @@ use crate::{ render::{ context::Context, disk_usage::{DiskUsage, FileSize}, - order::NodeComparator, }, }; use ansi_term::Color; use ansi_term::Style; use ignore::DirEntry; +use indextree::{Arena, Node as NodeWrapper, NodeId}; use lscolors::Style as LS_Style; use std::{ borrow::Cow, @@ -19,7 +19,6 @@ use std::{ fmt::{self, Formatter}, fs::{self, FileType}, path::{Path, PathBuf}, - slice::{Iter, IterMut}, }; /// A node of [`Tree`] that can be created from a [DirEntry]. Any filesystem I/O and @@ -32,7 +31,6 @@ use std::{ pub struct Node { pub depth: usize, pub file_size: Option, - children: Vec, file_name: OsString, file_type: Option, inode: Option, @@ -47,7 +45,6 @@ impl Node { pub fn new( depth: usize, file_size: Option, - children: Vec, file_name: OsString, file_type: Option, inode: Option, @@ -57,7 +54,6 @@ impl Node { symlink_target: Option, ) -> Self { Self { - children, depth, file_name, file_size, @@ -70,44 +66,9 @@ impl Node { } } - /// Returns a mutable reference to `children` if any. - pub fn children_mut(&mut self) -> IterMut { - self.children.iter_mut() - } - - /// Returns an iter over a `children` slice if any. - pub fn children(&self) -> Iter { - self.children.iter() - } - - /// Setter for `children`. - pub fn set_children(&mut self, children: Vec) { - self.children = children; - } - - /// Sorts `children` given comparator. - pub fn sort_children(&mut self, comparator: Box>) { - self.children.sort_by(comparator) - } - - /// Whether or not a [Node] has children. - pub fn has_children(&self) -> bool { - !self.children.is_empty() - } - /// Recursively traverse [Node]s, removing any [Node]s that have no children. pub fn prune_directories(&mut self) { - self.children.retain_mut(|node| { - if node.is_dir() { - if node.has_children() { - node.prune_directories(); - return node.has_children(); - } else { - return false; - } - } - true - }); + todo!(); } /// Returns a reference to `file_name`. If file is a symlink then `file_name` is the name of @@ -149,16 +110,9 @@ impl Node { self.file_type.as_ref() } - /// Returns the path to the [Node]'s parent, if any. This is a pretty expensive operation used - /// during parallel traversal. Perhaps an area for optimization. - pub fn parent_path_buf(&self) -> Option { - let mut path_buf = self.path.clone(); - - if path_buf.pop() { - Some(path_buf) - } else { - None - } + /// Returns the path to the [Node]'s parent, if any. + pub fn parent_path(&self) -> Option<&Path> { + self.path.parent() } /// Returns a reference to `path`. @@ -186,6 +140,20 @@ impl Node { self.inode.as_ref() } + /// Get a reference to [Node] from `inner` field of [`Tree`]. + /// + /// [`Tree`]: super::Tree + pub fn get(node_id: NodeId, tree: &Arena) -> Option<&Self> { + tree.get(node_id).map(NodeWrapper::get) + } + + /// Get a mutable reference to [Node] from `inner` field of [`Tree`]. + /// + /// [`Tree`]: super::Tree + pub fn get_mut(node_id: NodeId, tree: &mut Arena) -> Option<&mut Self> { + tree.get_mut(node_id).map(NodeWrapper::get_mut) + } + /// Gets stylized icon for node if enabled. Icons without extensions are styled based on the /// [`LS_COLORS`] foreground configuration of the associated file name. /// @@ -253,8 +221,6 @@ impl From<(&DirEntry, &Context)> for Node { let prefix = *prefix; let icons = *icons; - let children = vec![]; - let depth = dir_entry.depth(); let file_type = dir_entry.file_type(); @@ -300,7 +266,6 @@ impl From<(&DirEntry, &Context)> for Node { Self::new( depth, file_size, - children, file_name, file_type, inode, @@ -312,6 +277,12 @@ impl From<(&DirEntry, &Context)> for Node { } } +impl From<(NodeId, &mut Arena)> for &Node { + fn from((node_id, tree): (NodeId, &mut Arena)) -> Self { + tree.get(node_id).map(NodeWrapper::get).unwrap() + } +} + /// Simple struct to define location to put the `FileSize` while printing a `Node` #[derive(Copy, Clone, Default)] enum SizeLocation { diff --git a/src/render/tree/visitor.rs b/src/render/tree/visitor.rs new file mode 100644 index 00000000..e893dbcb --- /dev/null +++ b/src/render/tree/visitor.rs @@ -0,0 +1,53 @@ +use super::{Context, Node}; +use crossbeam::channel::Sender; +use ignore::{DirEntry, Error as IgnoreError, ParallelVisitor, ParallelVisitorBuilder, WalkState}; + +pub enum TraversalState { + Ongoing(Node), + Done, +} + +pub struct BranchVisitor<'a> { + ctx: &'a Context, + tx: Sender, +} + +pub struct BranchVisitorBuilder<'a> { + ctx: &'a Context, + tx: Sender, +} + +impl<'a> BranchVisitorBuilder<'a> { + pub fn new(ctx: &'a Context, tx: Sender) -> Self { + Self { ctx, tx } + } +} + +impl<'a> BranchVisitor<'a> { + pub fn new(ctx: &'a Context, tx: Sender) -> Self { + Self { ctx, tx } + } +} + +impl From for TraversalState { + fn from(node: Node) -> Self { + TraversalState::Ongoing(node) + } +} + +impl ParallelVisitor for BranchVisitor<'_> { + fn visit(&mut self, entry: Result) -> WalkState { + entry + .map(|e| TraversalState::from(Node::from((&e, self.ctx)))) + .map(|n| self.tx.send(n).unwrap()) + .map(|_| WalkState::Continue) + .unwrap_or(WalkState::Skip) + } +} + +impl<'s> ParallelVisitorBuilder<'s> for BranchVisitorBuilder<'s> { + fn build(&mut self) -> Box { + let visitor = BranchVisitor::new(self.ctx, self.tx.clone()); + Box::new(visitor) + } +}