path: root/src/tree/mod.rs

use crate::{
    context::{column, Context},
    disk_usage::file_size::FileSize,
    fs::inode::Inode,
    progress::{IndicatorHandle, Message},
    utils,
};
use count::FileCount;
use error::Error;
use ignore::{WalkBuilder, WalkParallel};
use indextree::{Arena, NodeId};
use node::{cmp::NodeComparator, Node};
use std::{
    collections::{HashMap, HashSet},
    convert::TryFrom,
    fs,
    path::PathBuf,
    result::Result as StdResult,
    sync::mpsc::{self, Sender},
    thread,
};
use visitor::{BranchVisitorBuilder, TraversalState};

/// Operations to handle and display aggregate file counts based on their type.
pub mod count;

/// Errors related to traversal, [Tree] construction, and the like.
pub mod error;

/// Contains components of the [`Tree`] data structure that derive from [`ignore::DirEntry`].
pub mod node;

/// Custom visitor that operates on each thread during filesystem traversal.
mod visitor;

/// Virtual data structure that represents local file-system hierarchy.
pub struct Tree {
    arena: Arena<Node>,
    root_id: NodeId,
}

pub type Result<T> = StdResult<T, Error>;

impl Tree {
    /// Constructor for [Tree].
    pub const fn new(arena: Arena<Node>, root_id: NodeId) -> Self {
        Self { arena, root_id }
    }

    /// Initiates file-system traversal and [Tree] as well as updates the [Context] object with
    /// various properties necessary to render output.
    pub fn try_init(
        mut ctx: Context,
        indicator: Option<&IndicatorHandle>,
    ) -> Result<(Self, Context)> {
        let mut column_properties = column::Properties::from(&ctx);

        let (arena, root_id) = Self::traverse(&ctx, &mut column_properties, indicator)?;

        ctx.update_column_properties(&column_properties);

        if ctx.truncate {
            ctx.set_window_width();
        }

        let tree = Self::new(arena, root_id);

        if tree.is_stump() {
            return Err(Error::NoMatches);
        }

        Ok((tree, ctx))
    }

    /// Returns `true` if there are no entries to show excluding the `root_id`.
    pub fn is_stump(&self) -> bool {
        self.root_id
            .descendants(self.arena())
            .skip(1)
            .peekable()
            .next()
            .is_none()
    }

    /// Grab a reference to `root_id`.
    pub const fn root_id(&self) -> NodeId {
        self.root_id
    }

    /// Grabs a reference to `arena`.
    pub const fn arena(&self) -> &Arena<Node> {
        &self.arena
    }

    /// Parallel traversal of the `root_id` 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,
        column_properties: &mut column::Properties,
        indicator: Option<&IndicatorHandle>,
    ) -> Result<(Arena<Node>, NodeId)> {
        let walker = WalkParallel::try_from(ctx)?;
        let (tx, rx) = mpsc::channel();

        let progress_indicator_mailbox = indicator.map(IndicatorHandle::mailbox);

        thread::scope(|s| {
            let res = s.spawn(move || {
                let mut tree = Arena::new();
                let mut branches: HashMap<PathBuf, Vec<NodeId>> = HashMap::new();
                let mut root_id = None;

                while let Ok(TraversalState::Ongoing(node)) = rx.recv() {
                    if let Some(ref mailbox) = progress_indicator_mailbox {
                        if mailbox.send(Message::Index).is_err() {
                            return Err(Error::Terminated);
                        }
                    }

                    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().ok_or(Error::ExpectedParent)?.to_owned();

                    let node_id = tree.new_node(node);

                    if branches
                        .get_mut(&parent)
                        .map(|mut_ref| mut_ref.push(node_id))
                        .is_none()
                    {
                        branches.insert(parent, vec![]);
                    }
                }

                if let Some(ref mailbox) = progress_indicator_mailbox {
                    if mailbox.send(Message::DoneIndexing).is_err() {
                        return Err(Error::Terminated);
                    }
                }

                let root_id = root_id.ok_or(Error::MissingRoot)?;
                let node_comparator = node::cmp::comparator(ctx);
                let mut inodes = HashSet::new();

                Self::assemble_tree(
                    &mut tree,
                    root_id,
                    &mut branches,
                    &node_comparator,
                    &mut inodes,
                    column_properties,
                    ctx,
                );

                if ctx.prune || ctx.pattern.is_some() {
                    Self::prune_directories(root_id, &mut tree);
                }

                if ctx.dirs_only {
                    Self::filter_directories(root_id, &mut tree);
                }

                Ok((tree, root_id))
            });

            let mut visitor_builder = BranchVisitorBuilder::new(ctx, Sender::clone(&tx));

            walker.visit(&mut visitor_builder);

            let _ = tx.send(TraversalState::Done);

            res.join().unwrap()
        })
    }

    /// Takes the results of the parallel traversal and uses it to construct the [Tree] data
    /// structure. Sorting occurs if specified. The amount of columns needed to fit all of the disk
    /// usages is also computed here.
    fn assemble_tree(
        tree: &mut Arena<Node>,
        current_node_id: NodeId,
        branches: &mut HashMap<PathBuf, Vec<NodeId>>,
        node_comparator: &NodeComparator,
        inode_set: &mut HashSet<Inode>,
        column_properties: &mut column::Properties,
        ctx: &Context,
    ) {
        let current_node = tree[current_node_id].get_mut();

        let mut children = branches.remove(current_node.path()).unwrap();

        let mut dir_size = FileSize::from(ctx);

        for child_id in &children {
            let index = *child_id;

            if tree[index].get().is_dir() {
                Self::assemble_tree(
                    tree,
                    index,
                    branches,
                    node_comparator,
                    inode_set,
                    column_properties,
                    ctx,
                );
            }

            let node = tree[index].get();

            #[cfg(unix)]
            Self::update_column_properties(column_properties, node, ctx);

            #[cfg(not(unix))]
            Self::update_column_properties(column_properties, node, ctx);

            // If a hard-link is already accounted for then don't increment parent dir size.
            if let Some(inode) = node.inode() {
                if inode.nlink > 1 && !inode_set.insert(inode) {
                    continue;
                }
            }

            if let Some(file_size) = node.file_size() {
                dir_size += file_size;
            }
        }

        if dir_size.value() > 0 {
            let dir = tree[current_node_id].get_mut();

            dir.set_file_size(dir_size);
        }

        let dir = tree[current_node_id].get();

        Self::update_column_properties(column_properties, dir, ctx);

        children.sort_by(|&id_a, &id_b| {
            let node_a = tree[id_a].get