Refactor filter.rs

src/filter.rs

@@ -5,83 +5,100 @@ use std::collections::HashSet;
 use std::path::Path;
 use std::path::PathBuf;
 
-pub fn get_biggest(
-    top_level_nodes: Vec<Node>,
-    min_size: Option<usize>,
-    only_dir: bool,
-    n: usize,
-    depth: usize,
-    using_a_filter: bool,
-) -> Option<DisplayNode> {
+pub struct AggregateData {
+    pub min_size: Option<usize>,
+    pub only_dir: bool,
+    pub number_of_lines: usize,
+    pub depth: usize,
+    pub using_a_filter: bool,
+}
+
+pub fn get_biggest(top_level_nodes: Vec<Node>, display_data: AggregateData) -> Option<DisplayNode> {
     if top_level_nodes.is_empty() {
         // perhaps change this, bring back Error object?
         return None;
     }
-
     let mut heap = BinaryHeap::new();
     let number_top_level_nodes = top_level_nodes.len();
-
-    let root = get_new_root(top_level_nodes);
-    let mut allowed_nodes = HashSet::new();
-
-    allowed_nodes.insert(root.name.as_path());
+    let root;
 
     if number_top_level_nodes > 1 {
-        heap = add_children(using_a_filter, min_size, only_dir, &root, usize::MAX, heap);
+        let size = top_level_nodes.iter().map(|node| node.size).sum();
+        root = Node {
+            name: PathBuf::from("(total)"),
+            size,
+            children: top_level_nodes,
+            inode_device: None,
+            depth: 0,
+        };
+        // Always include the base nodes if we add a 'parent' (total) node
+        heap = always_add_children(&display_data, &root, heap);
     } else {
-        heap = add_children(using_a_filter, min_size, only_dir, &root, depth, heap);
+        root = top_level_nodes.into_iter().next().unwrap();
+        heap = add_children(&display_data, &root, heap);
     }
 
-    let number_of_lines_in_output = n - number_top_level_nodes;
-    for _ in 0..number_of_lines_in_output {
+    let nol = display_data.number_of_lines;
+    let remaining = nol.saturating_sub(number_top_level_nodes);
+    fill_remaining_lines(heap, &root, remaining, display_data)
+}
+
+pub fn fill_remaining_lines<'a>(
+    mut heap: BinaryHeap<&'a Node>,
+    root: &'a Node,
+    remaining_lines: usize,
+    display_data: AggregateData,
+) -> Option<DisplayNode> {
+    let mut allowed_nodes = HashSet::new();
+    allowed_nodes.insert(root.name.as_path());
+
+    for _ in 0..remaining_lines {
         let line = heap.pop();
         match line {
             Some(line) => {
                 allowed_nodes.insert(line.name.as_path());
-                heap = add_children(using_a_filter, min_size, only_dir, line, depth, heap);
+                heap = add_children(&display_data, line, heap);
             }
             None => break,
         }
     }
-    recursive_rebuilder(&allowed_nodes, &root)
+    recursive_rebuilder(&allowed_nodes, root)
 }
 
 fn add_children<'a>(
-    using_a_filter: bool,
-    min_size: Option<usize>,
-    only_dir: bool,
+    display_data: &AggregateData,
     file_or_folder: &'a Node,
-    depth: usize,
-    mut heap: BinaryHeap<&'a Node>,
+    heap: BinaryHeap<&'a Node>,
 ) -> BinaryHeap<&'a Node> {
-    if depth > file_or_folder.depth {
-        heap.extend(
-            file_or_folder
-                .children
-                .iter()
-                .filter(|c| match min_size {
-                    Some(ms) => c.size > ms as u64,
-                    None => !using_a_filter || c.name.is_file() || c.size > 0,
-                })
-                .filter(|c| if only_dir { c.name.is_dir() } else { true }),
-        )
+    if display_data.depth > file_or_folder.depth {
+        always_add_children(display_data, file_or_folder, heap)
+    } else {
+        heap
     }
-    heap
 }
 
-fn get_new_root(top_level_nodes: Vec<Node>) -> Node {
-    if top_level_nodes.len() != 1 {
-        let size = top_level_nodes.iter().map(|node| node.size).sum();
-        Node {
-            name: PathBuf::from("(total)"),
-            size,
-            children: top_level_nodes,
-            inode_device: None,
-            depth: 0,
-        }
-    } else {
-        top_level_nodes.into_iter().next().unwrap()
-    }
+fn always_add_children<'a>(
+    display_data: &AggregateData,
+    file_or_folder: &'a Node,
+    mut heap: BinaryHeap<&'a Node>,
+) -> BinaryHeap<&'a Node> {
+    heap.extend(
+        file_or_folder
+            .children
+            .iter()
+            .filter(|c| match display_data.min_size {
+                Some(ms) => c.size > ms as u64,
+                None => !display_data.using_a_filter || c.name.is_file() || c.size > 0,
+            })
+            .filter(|c| {
+                if display_data.only_dir {
+                    c.name.is_dir()
+                } else {
+                    true
+                }
+            }),
+    );
+    heap
 }
 
 fn recursive_rebuilder(allowed_nodes: &HashSet<&Path>, current: &Node) -> Option<DisplayNode> {

src/main.rs

@@ -10,6 +10,8 @@ mod platform;
 mod utils;
 
 use crate::cli::build_cli;
+use dir_walker::WalkData;
+use filter::AggregateData;
 use std::collections::HashSet;
 use std::io::BufRead;
 use std::process;
@@ -18,7 +20,7 @@ use sysinfo::{System, SystemExt};
 use self::display::draw_it;
 use clap::Values;
 use config::get_config;
-use dir_walker::{walk_it, WalkData};
+use dir_walker::walk_it;
 use filter::get_biggest;
 use filter_type::get_all_file_types;
 use rayon::ThreadPoolBuildError;
@@ -175,17 +177,19 @@ fn main() {
 
     let iso = config.get_iso(&options);
     let (top_level_nodes, has_errors) = walk_it(simplified_dirs, walk_data);
-
     let tree = match summarize_file_types {
         true => get_all_file_types(&top_level_nodes, number_of_lines),
-        false => get_biggest(
-            top_level_nodes,
-            config.get_min_size(&options, iso),
-            config.get_only_dir(&options),
-            number_of_lines,
-            depth,
-            options.values_of("filter").is_some() || options.value_of("invert_filter").is_some(),
-        ),
+        false => {
+            let agg_data = AggregateData {
+                min_size: config.get_min_size(&options, iso),
+                only_dir: config.get_only_dir(&options),
+                number_of_lines,
+                depth,
+                using_a_filter: options.values_of("filter").is_some()
+                    || options.value_of("invert_filter").is_some(),
+            };
+            get_biggest(top_level_nodes, agg_data)
+        }
     };
 
     if has_errors {