parallel_disk_usage/
tree_builder.rs

1pub mod info;
2
3pub use info::Info;
4
5use super::{data_tree::DataTree, size};
6use rayon::prelude::*;
7
8/// Collection of functions and starting points in order to build a [`DataTree`] with [`From`] or [`Into`].
9#[derive(Debug)]
10pub struct TreeBuilder<Path, Name, Size, GetInfo, JoinPath>
11where
12    Path: Send + Sync,
13    Name: Send + Sync,
14    GetInfo: Fn(&Path) -> Info<Name, Size> + Copy + Send + Sync,
15    JoinPath: Fn(&Path, &Name) -> Path + Copy + Send + Sync,
16    Size: size::Size + Send,
17{
18    /// Path to the root.
19    pub path: Path,
20    /// Name of the root.
21    pub name: Name,
22    /// Function to extract necessary information from `path` (`size` and `children`).
23    pub get_info: GetInfo,
24    /// Function to join parent's `path` with a child's name to make the child's `name`.
25    pub join_path: JoinPath,
26    /// Deepest level of descendent to store as arrays. The sizes beyond the max depth still count toward total.
27    pub max_depth: u64,
28}
29
30impl<Path, Name, Size, GetInfo, JoinPath> From<TreeBuilder<Path, Name, Size, GetInfo, JoinPath>>
31    for DataTree<Name, Size>
32where
33    Path: Send + Sync,
34    Name: Send + Sync,
35    GetInfo: Fn(&Path) -> Info<Name, Size> + Copy + Send + Sync,
36    JoinPath: Fn(&Path, &Name) -> Path + Copy + Send + Sync,
37    Size: size::Size + Send,
38{
39    /// Create a [`DataTree`] from a [`TreeBuilder`].
40    fn from(builder: TreeBuilder<Path, Name, Size, GetInfo, JoinPath>) -> Self {
41        let TreeBuilder {
42            path,
43            name,
44            get_info,
45            join_path,
46            max_depth,
47        } = builder;
48
49        let Info { size, children } = get_info(&path);
50        let max_depth = max_depth.saturating_sub(1);
51
52        let children = children
53            .into_par_iter()
54            .map(|name| TreeBuilder {
55                path: join_path(&path, &name),
56                name,
57                get_info,
58                join_path,
59                max_depth,
60            })
61            .map(Self::from);
62
63        if max_depth > 0 {
64            DataTree::dir(name, size, children.collect())
65        } else {
66            let size = size + children.map(|child| child.size()).sum();
67            DataTree::dir(name, size, Vec::new())
68        }
69    }
70}