Complete day 7 part 1

day7/src/main.rs

@@ -9,11 +9,9 @@ struct File {
 }
 
 impl File {
-    fn new<T: ToString + ?Sized>(raw_line: &T) -> Self {
-        let raw = raw_line.to_string();
-        let parts: Vec<&str> = raw.split(' ').collect();
-        let size = parts[0].parse::<u128>().unwrap();
-        let name = parts[1].to_string();
+    fn new<T: ToString + ?Sized>(size: &T, name: &T) -> Self {
+        let size = size.to_string().parse::<u128>().unwrap();
+        let name = name.to_string();
 
         File { name, size }
     }
@@ -26,7 +24,7 @@ struct Dir {
     parent: String,
     name: String,
     files: Vec<File>,
-    dirs: Vec<String>,
+    subdirs: Vec<String>,
 }
 
 impl Dir {
@@ -35,21 +33,58 @@ impl Dir {
             parent: parent.to_string(),
             name: name.to_string(),
             files: Vec::new(),
-            dirs: Vec::new(),
+            subdirs: Vec::new(),
         }
     }
 
+    fn add_file(&mut self, file: File) {
+        self.files.push(file);
+    }
+
+    fn add_subdir(&mut self, path: String) {
+        self.subdirs.push(path);
+    }
+
     fn get_loose_files_size(&self) -> u128 {
         self.files
             .iter()
             .map(|file| file.size)
             .reduce(|accum, item| accum + item)
-            .unwrap()
+            .unwrap_or(0)
     }
 }
 
 // ----------------------------------------------------------------------------
 
+#[derive(Debug)]
+enum LineType {
+    Cd(String),
+    Ls,
+    Dir(String),
+    FileAndSize(String, String),
+}
+
+use LineType::*;
+
+impl LineType {
+    fn from(line: &str) -> LineType {
+        let parts: Vec<&str> = line.split_ascii_whitespace().collect();
+
+        match parts[0] {
+            "$" => match parts[1] {
+                "cd" => Cd(parts[2].to_string()),
+                "ls" => Ls,
+                _ => panic!("Invalid command"),
+            },
+            "dir" => Dir(parts[1].to_string()),
+            _ => FileAndSize(parts[0].to_string(), parts[1].to_string()),
+        }
+    }
+}
+
+// ----------------------------------------------------------------------------
+
+#[derive(Debug)]
 struct DirMap {
     current_path: String,
     map: HashMap<String, Dir>,
@@ -57,16 +92,135 @@ struct DirMap {
 
 impl DirMap {
     fn new() -> Self {
-        DirMap {
-            current_path: "/".to_string(),
-            map: HashMap::new(),
+        let current_path = "/".to_string();
+        let mut map: HashMap<String, Dir> = HashMap::new();
+        map.insert(current_path.clone(), Dir::new("", ""));
+
+        DirMap { current_path, map }
+    }
+
+    fn cd<T: ToString + ?Sized>(&mut self, new_dir: &T) {
+        let current_path = self.current_path.clone();
+        let new = new_dir.to_string();
+
+        match new.as_str() {
+            "/" => {
+                self.current_path = new.to_string();
+            }
+            ".." => {
+                let mut dir_parts: Vec<&str> = current_path.split('/').collect();
+                let _ = dir_parts.pop();
+                self.current_path = dir_parts.join("/");
+            }
+            _ => {
+                self.current_path.push('/');
+                self.current_path.push_str(&new);
+            }
+        }
+    }
+
+    fn dir<T: ToString + ?Sized>(&mut self, dir: &T) {
+        let parent = self.current_path.clone();
+        let name = dir.to_string();
+
+        let mut full_path = parent.clone();
+        full_path.push('/');
+        full_path.push_str(&name);
+
+        // Add the new Dir to the path map
+        if !self.map.contains_key(&full_path) {
+            self.map.insert(full_path.clone(), Dir::new(&parent, &name));
+        }
+
+        // Add the new Dir to the list of subdirectories to the Dir mapped to the current path
+        self.map
+            .get_mut(&self.current_path)
+            .expect(&format!(
+                "This dir ({}) should already exist",
+                &self.current_path
+            ))
+            .add_subdir(full_path.clone());
+    }
+
+    fn parse(&mut self, item: LineType) {
+        match item {
+            Cd(s) => self.cd(&s),
+            Ls => {}
+            Dir(s) => self.dir(&s),
+            FileAndSize(size, name) => {
+                self.map
+                    .get_mut(&self.current_path)
+                    .expect(&format!(
+                        "This dir ({}) should already exist",
+                        &self.current_path
+                    ))
+                    .add_file(File::new(&size, &name));
+            }
         }
     }
 }
 
 // ----------------------------------------------------------------------------
 
+fn get_path_size_map(dir_map: &DirMap) -> HashMap<String, u128> {
+    let mut size_map: HashMap<String, u128> = HashMap::new();
+
+    // Get the sizes of the leaf node directories
+    dir_map
+        .map
+        .iter()
+        .filter(|(_, v)| v.subdirs.len() == 0)
+        .for_each(|(k, v)| {
+            size_map.insert(k.to_string(), v.get_loose_files_size());
+        });
+
+    // Calculate dir sizes by the length of the path from largest to smallest,
+    // so we can start with the lowest branches of the tree when calculating folder sizes
+    let mut branch_paths: Vec<&String> = dir_map
+        .map
+        .iter()
+        .filter(|(_, v)| v.subdirs.len() > 0)
+        .map(|(k, _)| k)
+        .collect();
+    branch_paths.sort();
+    branch_paths.reverse();
+
+    branch_paths.into_iter().for_each(|path| {
+        let dir = dir_map.map.get(path).unwrap();
+
+        let base_size = dir.get_loose_files_size();
+        let subdir_size: u128 = dir
+            .subdirs
+            .iter()
+            .map(|sub| {
+                *size_map
+                    .get(sub)
+                    .expect("Dir {} should already have had its size calculated")
+            })
+            .reduce(|accum, item| accum + item)
+            .unwrap_or(0);
+
+        size_map.insert(path.to_string(), base_size + subdir_size);
+    });
+
+    size_map
+}
+
 fn main() {
     let file_str = include_str!("input.txt");
-    let path_map = DirMap::new();
+    let mut path_map = DirMap::new();
+
+    file_str
+        .lines()
+        .map(|line| LineType::from(line))
+        .for_each(|cmd| path_map.parse(cmd));
+
+    let size_map = get_path_size_map(&path_map);
+
+    let size_sum = size_map
+        .iter()
+        .filter(|(_, v)| **v < MAX_DIR_SIZE)
+        .fold(0u128, |acc, (_, v)| acc + *v);
+
+    println!("Part 1: Sum of dirs 100K or smaller {:#?}", size_sum);
 }