Refactor code

src/tests.rs

@@ -41,12 +41,21 @@ fn main_output() -> String {
 {}
 {}",
         format_string("src/test_dir", true, " 8.0K", ""),
-        format_string("src/test_dir/many", true, " 8.0K", "└─┬",),
+        format_string("src/test_dir/many", true, " 4.0K", "└─┬",),
         format_string("src/test_dir/many/hello_file", true, " 4.0K", "  ├──",),
         format_string("src/test_dir/many/a_file", false, "   0B", "  └──",),
     )
 }
 
+#[test]
+pub fn test_main_extra_slash() {
+    assert_cli::Assert::main_binary()
+        .with_args(&["src/test_dir/"])
+        .stdout()
+        .is(main_output())
+        .unwrap();
+}
+
 #[test]
 pub fn test_apparent_size() {
     let r = format!(

src/utils/mod.rs

@@ -1,73 +1,80 @@
 use std::collections::HashSet;
 
-use std::fs::{self, ReadDir};
+use std::fs;
-use std::io;
 
 use dust::Node;
+use std::path::Path;
 
 mod platform;
 use self::platform::*;
 
 pub fn get_dir_tree(filenames: &Vec<&str>, apparent_size: bool) -> (bool, Vec<Node>) {
     let mut permissions = true;
+    let mut inodes: HashSet<(u64, u64)> = HashSet::new();
     let mut results = vec![];
     for &b in filenames {
-        let mut new_name = String::from(b);
+        let filename = strip_end_slashes(b);
-        while new_name.chars().last() == Some('/') && new_name.len() != 1 {
+        let (hp, data) = examine_dir(&Path::new(&filename), apparent_size, &mut inodes);
-            new_name.pop();
-        }
-        let (hp, data) = examine_dir_str(&new_name, apparent_size);
         permissions = permissions && hp;
-        results.push(data);
+        match data {
+            Some(d) => results.push(d),
+            None => permissions = false,
+        }
     }
     (permissions, results)
 }
 
-fn examine_dir_str(loc: &str, apparent_size: bool) -> (bool, Node) {
+fn strip_end_slashes(s: &str) -> String {
-    let mut inodes: HashSet<(u64, u64)> = HashSet::new();
+    let mut new_name = String::from(s);
-    let (hp, result) = examine_dir(fs::read_dir(loc), apparent_size, &mut inodes);
+    while new_name.chars().last() == Some('/') && new_name.len() != 1 {
-
+        new_name.pop();
-    // This needs to be folded into the below recursive call somehow
+    }
-    let new_size = result.iter().fold(0, |a, b| a + b.size());
+    new_name
-    (hp, Node::new(loc, new_size, result))
 }
 
 fn examine_dir(
-    a_dir: io::Result<ReadDir>,
+    sdir: &Path,
     apparent_size: bool,
     inodes: &mut HashSet<(u64, u64)>,
-) -> (bool, Vec<Node>) {
+) -> (bool, Option<Node>) {
+    match fs::read_dir(sdir) {
+        Ok(file_iter) => {
             let mut result = vec![];
             let mut have_permission = true;
+            let mut total_size = 0;
 
-    if a_dir.is_ok() {
+            for single_path in file_iter {
-        let paths = a_dir.unwrap();
+                match single_path {
-        for dd in paths {
-            match dd {
                     Ok(d) => {
                         let file_type = d.file_type().ok();
                         let maybe_size_and_inode = get_metadata(&d, apparent_size);
 
                         match (file_type, maybe_size_and_inode) {
-                        (Some(file_type), Some((size, inode))) => {
+                            (Some(file_type), Some((size, maybe_inode))) => {
-                            let s = d.path().to_string_lossy().to_string();
                                 if !apparent_size {
-                                if let Some(inode_dev_pair) = inode {
+                                    if let Some(inode_dev_pair) = maybe_inode {
                                         if inodes.contains(&inode_dev_pair) {
                                             continue;
                                         }
                                         inodes.insert(inode_dev_pair);
                                     }
                                 }
+                                total_size += size;
 
                                 if d.path().is_dir() && !file_type.is_symlink() {
-                                let (hp, recursive) =
+                                    let (hp, child) = examine_dir(&d.path(), apparent_size, inodes);
-                                    examine_dir(fs::read_dir(d.path()), apparent_size, inodes);
                                     have_permission = have_permission && hp;
-                                let new_size = recursive.iter().fold(size, |a, b| a + b.size());
+
-                                result.push(Node::new(s, new_size, recursive))
+                                    match child {
+                                        Some(c) => {
+                                            total_size += c.size();
+                                            result.push(c);
+                                        }
+                                        None => (),
+                                    }
                                 } else {
-                                result.push(Node::new(s, size, vec![]))
+                                    let path_name = d.path().to_string_lossy().to_string();
+                                    result.push(Node::new(path_name, size, vec![]))
                                 }
                             }
                             (_, None) => have_permission = false,
@@ -77,10 +84,11 @@ fn examine_dir(
                     Err(_) => (),
                 }
             }
-    } else {
+            let n = Node::new(sdir.to_string_lossy().to_string(), total_size, result);
-        have_permission = false;
+            (have_permission, Some(n))
+        }
+        Err(_) => (false, None),
     }
-    (have_permission, result)
 }
 
 // We start with a list of root directories - these must be the biggest folders