We got the first star!

src/bin/p07.rs

@@ -1,9 +1,12 @@
-#![feature(iterator_try_collect, absolute_path)]
+#![feature(iterator_try_collect, absolute_path, hash_set_entry)]
 
 use std::collections::HashMap;
+use std::error::Error;
+use std::fmt::Display;
 use std::fs;
+use std::hash::{Hash, Hasher};
 use std::io::{self, BufRead, Read};
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
 
 mod parsing {
     use nom::{
@@ -42,15 +45,6 @@ mod parsing {
         File(String, usize),
     }
 
-    impl Listing {
-        pub fn path(&self) -> &String {
-            match self {
-                Self::Dir(s) => s,
-                Self::File(s, _) => s,
-            }
-        }
-    }
-
     fn dir(input: &str) -> IResult<&str, Listing> {
         let (input, f_name) = preceded(pair(tag("dir"), multispace0), filename)(input)?;
         Ok((input, Listing::Dir(f_name)))
@@ -143,22 +137,129 @@ use path_clean::clean;
 #[derive(Debug)]
 struct State {
     cur_path: PathBuf,
-    filesystem: HashMap<String, usize>,
+    filesystem: FileNode,
+}
+
+#[derive(Debug)]
+struct FileNode {
+    name: String,
+    children: HashMap<String, FileNode>,
+    size: Option<usize>,
+}
+
+impl Display for FileNode {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{} ", if self.is_dir() { "dir" } else { "file" })?;
+        if let Some(size) = self.size {
+            write!(f, "{} ", size)?;
+        }
+        write!(f, "{}", self.name)
+    }
+}
+
+impl PartialEq for FileNode {
+    fn eq(&self, other: &Self) -> bool {
+        self.name.eq(&other.name)
+    }
+}
+impl Eq for FileNode {}
+
+impl Hash for FileNode {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.name.hash(state);
+    }
+}
+
+impl From<String> for FileNode {
+    fn from(value: String) -> Self {
+        FileNode::new(value)
+    }
+}
+
+impl<'a> IntoIterator for &'a FileNode {
+    type Item = &'a FileNode;
+
+    type IntoIter = std::vec::IntoIter<Self::Item>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        fn append<'a>(tree: &'a FileNode, v: &mut Vec<&'a FileNode>) {
+            v.push(tree);
+            for child in tree.children.values() {
+                append(child, v);
+            }
+        }
+
+        let mut result = vec![];
+        append(self, &mut result);
+        result.into_iter()
+    }
+}
+
+impl FileNode {
+    fn is_dir(&self) -> bool {
+        !self.children.is_empty()
+    }
+
+    fn new(name: String) -> Self {
+        Self {
+            name,
+            children: HashMap::new(),
+            size: None,
+        }
+    }
+
+    fn calculate_size(&mut self) -> usize {
+        *self.size.get_or_insert_with(|| {
+            self.children
+                .iter_mut()
+                .fold(0, |acc, (_, v)| acc + v.calculate_size())
+        })
+    }
+
+    fn add(&mut self, path: String, size: Option<usize>) -> Result<(), Box<dyn Error>> {
+        let path = Path::new(&path);
+        if path.components().count() == 1 {
+            self.size = self.size.or(size);
+            return Ok(());
+        }
+
+        let path = path.strip_prefix(&self.name)?;
+        let first_comp = path
+            .components()
+            .next()
+            .ok_or("Expected at least one component")?
+            .as_os_str()
+            .to_string_lossy()
+            .to_string();
+
+        let path_str = path.as_os_str().to_string_lossy().to_string();
+
+        if let Some(child) = self.children.get_mut(&first_comp) {
+            child.add(path_str, size)?;
+        } else {
+            let mut child = FileNode::new(first_comp.clone());
+            child.add(path_str, size)?;
+            self.children.insert(first_comp, child);
+        }
+
+        Ok(())
+    }
 }
 
 impl State {
     fn new() -> Self {
         State {
             cur_path: PathBuf::from("/"),
-            filesystem: HashMap::new(),
+            filesystem: FileNode::new("/".to_owned()),
         }
     }
 
-    fn process(&mut self, line: Line) {
+    fn process(&mut self, line: Line) -> Result<(), Box<dyn Error>> {
         match line {
             Line::Command(cmd) => self.run(cmd),
-            Line::Result(listing) => self.process_result(listing),
+            Line::Result(listing) => self.process_result(listing)?,
         }
+        Ok(())
     }
 
     fn run(&mut self, cmd: Command) {
@@ -167,11 +268,15 @@ impl State {
         }
     }
 
-    fn process_result(&mut self, res: Listing) {
-        if let Listing::File(path, size) = res {
-            let filename = clean(&self.cur_path.join(path).to_string_lossy());
-            self.filesystem.insert(filename, size);
-        }
+    fn process_result(&mut self, res: Listing) -> Result<(), Box<dyn Error>> {
+        let (path, maybe_size) = match res {
+            Listing::File(path, size) => (path, Some(size)),
+            Listing::Dir(path) => (path, None),
+        };
+
+        let filename = clean(&self.cur_path.join(path).to_string_lossy());
+        self.filesystem.add(filename, maybe_size)?;
+        Ok(())
     }
 }
 
@@ -181,9 +286,19 @@ fn part1(lines: &[String]) {
     lines
         .iter()
         .flat_map(|l| line(l))
-        .for_each(|(_, l)| state.process(l));
+        .for_each(|(_, l)| state.process(l).unwrap());
 
-    dbg!(state.filesystem);
+    let full_size = state.filesystem.calculate_size();
+    dbg!(full_size);
+
+    const MAX_DIR_SIZE: usize = 100_000;
+    dbg!(state
+        .filesystem
+        .into_iter()
+        .filter(|f| f.is_dir())
+        .flat_map(|f| f.size)
+        .filter(|&s| s < MAX_DIR_SIZE)
+        .sum::<usize>());
 }
 
-fn part2(lines: &[String]) {}
+fn part2(_lines: &[String]) {}