Day 7 part 2 implemented with trees

src/main.rs

@@ -1,4 +1,6 @@
 #![feature(test)]
+
+use std::collections::HashMap;
 extern crate test;
 
 const INPUTS: [&str; 2] = [
@@ -12,6 +14,19 @@ fn parse(input: &'static str) -> impl Iterator<Item = &'static str> {
     input
 }
 
+#[derive(Debug)]
+enum Node {
+    File(u32),
+    Directory(Directory),
+}
+
+#[derive(Debug)]
+struct Directory {
+    children: HashMap<String, Node>,
+    parent: *mut Node,
+    size: u32,
+}
+
 fn main() {
     for input in INPUTS.iter() {
         let output = parse(input);
@@ -22,56 +37,111 @@ fn main() {
 }
 
 fn solution(input: impl Iterator<Item = &'static str>) -> u32 {
-    let mut out = Vec::with_capacity(100);
-    let mut stack = Vec::with_capacity(100);
+    let mut tree = Node::Directory(Directory {
+        children: HashMap::new(),
+        parent: std::ptr::null_mut(),
+        size: 0,
+    });
 
-    let mut current_folder_size = 0;
+    let mut current = &mut tree;
 
     for line in input {
         match &line[..4] {
-            "$ ls" | "dir " => continue,
+            "$ ls" => continue,
+
+            "dir " => {
+                let dir = &line[4..];
+                let cptr = current as *mut Node;
+
+                let Node::Directory(Directory { children, ..}) = current else { unreachable!()};
+
+                children.insert(
+                    dir.to_string(),
+                    Node::Directory(Directory {
+                        children: HashMap::new(),
+                        parent: cptr,
+                        size: 0,
+                    }),
+                );
+            }
 
             "$ cd" => match &line[5..6] {
                 // we are supposed to match on .. but this is fine
                 "." => {
-                    let v = stack.pop().unwrap();
-                    out.push(current_folder_size);
-                    current_folder_size += v;
+                    let Node::Directory(Directory { parent , .. }) = *current  else {
+                        unreachable!("current is not a directory");
+                    };
+
+                    current = unsafe { &mut *parent };
                 }
 
-                "/" => continue,
+                "/" => current = &mut tree,
 
                 _ => {
-                    stack.push(current_folder_size);
-                    current_folder_size = 0;
+                    let dir = &line[5..];
+                    let Node::Directory(Directory { children, ..})  = current  else {
+                            unreachable!("current is not a directory");
+                        };
+
+                    current = children.get_mut(dir).unwrap();
                 }
             },
 
             _ => {
-                let size = line
+                let (size, name) = line.split_once(' ').unwrap();
+                let fsize = size
                     .bytes()
                     .take_while(|&c| c != b' ')
                     .fold(0u32, |a, x| a * 10 + (x - b'0') as u32);
-                current_folder_size += size;
+
+                if let Node::Directory(v) = current {
+                    v.children
+                        .entry(name.to_string())
+                        .or_insert(Node::File(fsize));
+                }
             }
         }
     }
 
-    while let Some(v) = stack.pop() {
-        out.push(current_folder_size);
-        current_folder_size += v;
+    compute_dir_size(&mut tree);
+
+    part1(&tree)
+}
+
+fn part1(node: &Node) -> u32 {
+    match node {
+        &Node::File(_) => 0,
+
+        Node::Directory(dir) => {
+            let mut sum = 0;
+            if dir.size <= 100000 {
+                sum += dir.size;
+            }
+
+            for v in dir.children.values() {
+                sum += part1(v);
+            }
+
+            sum
+        }
     }
+}
 
-    out.push(current_folder_size);
+fn compute_dir_size(node: &mut Node) -> u32 {
+    match node {
+        &mut Node::File(v) => v,
 
-    let available_space = 70000000;
-    let required_space = 30000000;
-    let used_space = *out.last().unwrap();
+        Node::Directory(dir) => {
+            let mut sum = 0;
 
-    out.into_iter()
-        .filter(|&c| c >= (required_space + used_space - available_space))
-        .min()
-        .unwrap()
+            for v in dir.children.values_mut() {
+                sum += compute_dir_size(v);
+            }
+
+            dir.size = sum;
+            sum
+        }
+    }
 }
 
 #[bench]