parallel_disk_usage/app/
sub.rs

1use crate::{
2    args::Fraction,
3    data_tree::{DataTree, DataTreeReflection},
4    fs_tree_builder::FsTreeBuilder,
5    get_size::GetSize,
6    json_data::{BinaryVersion, JsonData, SchemaVersion, UnitAndTree},
7    os_string_display::OsStringDisplay,
8    reporter::ParallelReporter,
9    runtime_error::RuntimeError,
10    size,
11    status_board::GLOBAL_STATUS_BOARD,
12    visualizer::{BarAlignment, ColumnWidthDistribution, Direction, Visualizer},
13};
14use serde::Serialize;
15use std::{io::stdout, iter::once, num::NonZeroU64, path::PathBuf};
16
17/// The sub program of the main application.
18pub struct Sub<Size, SizeGetter, Report>
19where
20    Report: ParallelReporter<Size> + Sync,
21    Size: size::Size + Into<u64> + Serialize + Send + Sync,
22    SizeGetter: GetSize<Size = Size> + Copy + Sync,
23    DataTreeReflection<String, Size>: Into<UnitAndTree>,
24{
25    /// List of files and/or directories.
26    pub files: Vec<PathBuf>,
27    /// Print JSON data instead of an ASCII chart.
28    pub json_output: bool,
29    /// Format to be used to [`display`](size::Size::display) the sizes returned by [`size_getter`](Self::size_getter).
30    pub bytes_format: Size::DisplayFormat,
31    /// The direction of the visualization.
32    pub direction: Direction,
33    /// The alignment of the bars.
34    pub bar_alignment: BarAlignment,
35    /// Distribution and number of characters/blocks can be placed in a line.
36    pub column_width_distribution: ColumnWidthDistribution,
37    /// Maximum number of levels that should be visualized.
38    pub max_depth: NonZeroU64,
39    /// [Get the size](GetSize) of files/directories.
40    pub size_getter: SizeGetter,
41    /// Reports measurement progress.
42    pub reporter: Report,
43    /// Minimal size proportion required to appear.
44    pub min_ratio: Fraction,
45    /// Preserve order of entries.
46    pub no_sort: bool,
47}
48
49impl<Size, SizeGetter, Report> Sub<Size, SizeGetter, Report>
50where
51    Size: size::Size + Into<u64> + Serialize + Send + Sync,
52    Report: ParallelReporter<Size> + Sync,
53    SizeGetter: GetSize<Size = Size> + Copy + Sync,
54    DataTreeReflection<String, Size>: Into<UnitAndTree>,
55{
56    /// Run the sub program.
57    pub fn run(self) -> Result<(), RuntimeError> {
58        let Sub {
59            files,
60            json_output,
61            bytes_format,
62            direction,
63            bar_alignment,
64            column_width_distribution,
65            max_depth,
66            size_getter,
67            reporter,
68            min_ratio,
69            no_sort,
70        } = self;
71
72        let max_depth = max_depth.get();
73
74        let mut iter = files
75            .into_iter()
76            .map(|root| -> DataTree<OsStringDisplay, Size> {
77                FsTreeBuilder {
78                    reporter: &reporter,
79                    root,
80                    size_getter,
81                    max_depth,
82                }
83                .into()
84            });
85
86        let data_tree = if let Some(data_tree) = iter.next() {
87            data_tree
88        } else {
89            return Sub {
90                files: vec![".".into()],
91                reporter,
92                ..self
93            }
94            .run();
95        };
96
97        // ExactSizeIterator::is_empty is unstable
98        let data_tree = if iter.len() == 0 {
99            data_tree
100        } else {
101            let children: Vec<_> = once(data_tree).chain(iter).collect();
102            DataTree::dir(
103                OsStringDisplay::os_string_from("(total)"),
104                Size::default(),
105                children,
106            )
107            .into_par_retained(|_, depth| depth + 1 < max_depth)
108        };
109
110        if reporter.destroy().is_err() {
111            eprintln!("[warning] Failed to destroy the thread that reports progress");
112        }
113
114        let min_ratio: f32 = min_ratio.into();
115        let data_tree = {
116            let mut data_tree = data_tree;
117            if min_ratio > 0.0 {
118                data_tree.par_cull_insignificant_data(min_ratio);
119            }
120            if !no_sort {
121                data_tree.par_sort_by(|left, right| left.size().cmp(&right.size()).reverse());
122            }
123            data_tree
124        };
125
126        GLOBAL_STATUS_BOARD.clear_line(0);
127
128        if json_output {
129            let unit_and_tree: UnitAndTree = data_tree
130                .into_reflection() // I really want to use std::mem::transmute here but can't.
131                .par_convert_names_to_utf8() // TODO: allow non-UTF8 somehow.
132                .expect("convert all names from raw string to UTF-8")
133                .into();
134            let json_data = JsonData {
135                schema_version: SchemaVersion,
136                binary_version: Some(BinaryVersion::current()),
137                unit_and_tree,
138            };
139            return serde_json::to_writer(stdout(), &json_data)
140                .map_err(RuntimeError::SerializationFailure);
141        }
142
143        let visualizer = Visualizer {
144            data_tree: &data_tree,
145            bytes_format,
146            direction,
147            bar_alignment,
148            column_width_distribution,
149        };
150
151        print!("{visualizer}"); // visualizer already ends with "\n", println! isn't needed here.
152        Ok(())
153    }
154}