simplify

4 files changed, 226 insertions, 82 deletions

diff --git a/src/2022/7/mod.rs b/src/2022/7/mod.rs
index bdc5192..9fbaf97 100644
--- a/src/2022/7/mod.rs
+++ b/src/2022/7/mod.rs
@@ -3,74 +3,35 @@
 
 use crate::prelude::*;
 
-enum Node {
-    Dir(HashMap<String, Node>),
+pub enum Node {
+    Dir,
     File(i64),
 }
 
 impl Node {
     fn new_dir() -> Self {
-        Self::Dir(HashMap::new())
+        Self::Dir
     }
-    
+
     fn new_file(size: i64) -> Self {
         Self::File(size)
     }
-    
+
     fn size(&self) -> i64 {
         match self {
+            Self::Dir => 0,
             Self::File(size) => *size,
-            Self::Dir(children) => children
-                .values()
-                .map(|child| child.size())
-                .sum(),
-        }
-    }
-    
-    fn iter(&self) -> NodeIter<'_> {
-        let mut next = std::collections::VecDeque::new();
-        next.push_back(("", self));
-        NodeIter { next }
-    }
-
-    fn resolve_path(&mut self, path: &[String]) -> &mut Self {
-        let mut pwd = self;
-        for name in path {
-            let Node::Dir(children) = pwd 
-            else { panic!("invalid path") };
-
-            pwd = children.get_mut(name).unwrap();
-        }
-        pwd
-    }
-}
-
-struct NodeIter<'a> {
-    next: std::collections::VecDeque<(&'a str, &'a Node)>,
-}
-
-impl<'a> Iterator for NodeIter<'a> {
-    type Item = (&'a str, &'a Node);
-    
-    fn next(&mut self) -> Option<Self::Item> {
-        let Some(next) = self.next.pop_back()
-        else { return None };
-        if let Node::Dir(children) = next.1 {
-            for (name, child) in children.iter() {
-                self.next.push_front((name, child));
-            }
         }
-        Some(next)
     }
 }
 
-pub struct DirTree {
-    root: Node,
+fn tree_size(tree: &Tree<String, Node>) -> i64 {
+    tree.bfs().map(|(_, tree)| tree.data().size()).sum()
 }
 
-pub fn parse(fh: File) -> Result<DirTree> {
+pub fn parse(fh: File) -> Result<Tree<String, Node>> {
     let mut lines = parse::lines(fh).peekable();
-    let mut root = Node::new_dir();
+    let mut root = Tree::new(Node::Dir);
     let mut path = vec![];
 
     while let Some(cmdline) = lines.by_ref().next() {
@@ -82,7 +43,7 @@ pub fn parse(fh: File) -> Result<DirTree> {
         let Some(cmd) = captures.get(1)
         else { bail!("no cmd found in cmdline: '{}'", cmdline) };
 
-        let pwd = root.resolve_path(&path);
+        let pwd = root.at_mut(&path).unwrap();
 
         match cmd.as_str() {
             "cd" => {
@@ -94,11 +55,8 @@ pub fn parse(fh: File) -> Result<DirTree> {
                     }
                     "/" => path = vec![],
                     dirname => {
-                        let Node::Dir(children) = pwd 
-                        else { panic!("invalid path") };
-
-                        if !children.contains_key(dirname) {
-                            children.insert(
+                        if pwd.at(&[dirname.to_string()]).is_none() {
+                            pwd.add_child(
                                 dirname.to_string(),
                                 Node::new_dir(),
                             );
@@ -107,45 +65,40 @@ pub fn parse(fh: File) -> Result<DirTree> {
                     }
                 }
             }
-            "ls" => {
-                loop {
-                    let Node::Dir(children) = pwd 
-                    else { panic!("invalid path") };
-
-                    let Some(lsline) = lines.peek()
+            "ls" => loop {
+                let Some(lsline) = lines.peek()
                     else { break };
-                    let Some(captures) = &regex_captures!(
+                let Some(captures) = &regex_captures!(
                         r"^(dir|[0-9]+) (.*)$",
                         lsline
                     )
                     else { break };
-                    let Some(data) = captures.get(1)
+                let Some(data) = captures.get(1)
                     else { break };
-                    let Some(name) = captures.get(2)
+                let Some(name) = captures.get(2)
                     else { break };
 
-                    let node = match data.as_str() {
-                        "dir" => Node::new_dir(),
-                        size => Node::new_file(size.parse().unwrap()),
-                    };
-                    children.insert(name.as_str().to_string(), node);
-                    lines.next();
-                }
-            }
+                let node = match data.as_str() {
+                    "dir" => Node::new_dir(),
+                    size => Node::new_file(size.parse().unwrap()),
+                };
+                pwd.add_child(name.as_str().to_string(), node);
+                lines.next();
+            },
             _ => unreachable!(),
         }
     }
 
-    Ok(DirTree { root })
+    Ok(root)
 }
 
-pub fn part1(tree: DirTree) -> Result<i64> {
+pub fn part1(tree: Tree<String, Node>) -> Result<i64> {
     let mut total = 0;
-    for (name, node) in tree.root.iter() {
-        if matches!(node, Node::File(_)) {
+    for (_, tree) in tree.bfs() {
+        if matches!(tree.data(), Node::File(_)) {
             continue;
         }
-        let size = node.size();
+        let size = tree_size(tree);
         if size <= 100_000 {
             total += size;
         }
@@ -153,16 +106,16 @@ pub fn part1(tree: DirTree) -> Result<i64> {
     Ok(total)
 }
 
-pub fn part2(tree: DirTree) -> Result<i64> {
-    let total = tree.root.size();
+pub fn part2(tree: Tree<String, Node>) -> Result<i64> {
+    let total = tree_size(&tree);
     let free = 70_000_000 - total;
     let needed = 30_000_000 - free;
     let mut possible = vec![];
-    for (name, node) in tree.root.iter() {
-        if matches!(node, Node::File(_)) {
+    for (_, tree) in tree.bfs() {
+        if matches!(tree.data(), Node::File(_)) {
             continue;
         }
-        let size = node.size();
+        let size = tree_size(tree);
         if size > needed {
             possible.push(size);
         }
diff --git a/src/main.rs b/src/main.rs
index e04fedb..490fa02 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -16,6 +16,7 @@ pub mod graph;
 pub mod grid;
 pub mod parse;
 pub mod prelude;
+pub mod tree;
 
 #[path = "2020/mod.rs"]
 mod year2020;
diff --git a/src/prelude.rs b/src/prelude.rs
index 6c53874..046b4fb 100644
--- a/src/prelude.rs
+++ b/src/prelude.rs
@@ -1,6 +1,7 @@
 pub use crate::graph::Graph as _;
 pub use crate::grid::{Col, Grid, Row};
 pub use crate::parse;
+pub use crate::tree::Tree;
 
 pub use std::cmp::Ordering;
 pub use std::collections::VecDeque;
diff --git a/src/tree.rs b/src/tree.rs
new file mode 100644
index 0000000..3c727dd
--- /dev/null
+++ b/src/tree.rs
@@ -0,0 +1,189 @@
+use crate::prelude::*;
+
+#[derive(Debug)]
+pub struct Tree<K: std::hash::Hash + std::cmp::Eq, V> {
+    data: V,
+    children: HashMap<K, Tree<K, V>>,
+}
+
+impl<K, V> Tree<K, V>
+where
+    K: std::hash::Hash + std::cmp::Eq + std::clone::Clone,
+{
+    pub fn new(data: V) -> Self {
+        Self {
+            data,
+            children: HashMap::new(),
+        }
+    }
+
+    pub fn data(&self) -> &V {
+        &self.data
+    }
+
+    pub fn data_mut(&mut self) -> &mut V {
+        &mut self.data
+    }
+
+    pub fn add_child(&mut self, name: K, data: V) {
+        self.children.insert(name, Self::new(data));
+    }
+
+    pub fn at<'a, I>(&'a self, path: I) -> Option<&'a Self>
+    where
+        I: IntoIterator<Item = &'a K>,
+    {
+        let mut ret = self;
+        for entry in path {
+            let Some(child) = ret.children.get(entry)
+            else { return None };
+            ret = child;
+        }
+        Some(ret)
+    }
+
+    pub fn at_mut<'a, I>(&'a mut self, path: I) -> Option<&'a mut Self>
+    where
+        I: IntoIterator<Item = &'a K>,
+    {
+        let mut ret = self;
+        for entry in path {
+            let Some(child) = ret.children.get_mut(entry)
+            else { return None };
+            ret = child;
+        }
+        Some(ret)
+    }
+
+    pub fn bfs(&self) -> BFS<'_, K, V> {
+        let mut next = std::collections::VecDeque::new();
+        next.push_front((vec![], self));
+        BFS { next }
+    }
+
+    pub fn dfs(&self) -> DFS<'_, K, V> {
+        DFS {
+            next: vec![(vec![], self)],
+        }
+    }
+}
+
+pub struct BFS<'a, K, V>
+where
+    K: std::cmp::Eq + std::hash::Hash + std::clone::Clone,
+{
+    next: std::collections::VecDeque<(Vec<K>, &'a Tree<K, V>)>,
+}
+
+impl<'a, K, V> Iterator for BFS<'a, K, V>
+where
+    K: std::cmp::Eq + std::hash::Hash + std::clone::Clone,
+{
+    type Item = (Vec<K>, &'a Tree<K, V>);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let Some(next) = self.next.pop_front()
+        else { return None };
+
+        self.next.extend(next.1.children.iter().map(|(k, v)| {
+            (
+                next.0
+                    .clone()
+                    .into_iter()
+                    .chain(std::iter::once(k.clone()))
+                    .collect(),
+                v,
+            )
+        }));
+        Some(next)
+    }
+}
+
+pub struct DFS<'a, K, V>
+where
+    K: std::cmp::Eq + std::hash::Hash + std::clone::Clone,
+{
+    next: Vec<(Vec<K>, &'a Tree<K, V>)>,
+}
+
+impl<'a, K, V> Iterator for DFS<'a, K, V>
+where
+    K: std::cmp::Eq + std::hash::Hash + std::clone::Clone,
+{
+    type Item = (Vec<K>, &'a Tree<K, V>);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let Some(next) = self.next.pop()
+        else { return None };
+
+        self.next.extend(next.1.children.iter().map(|(k, v)| {
+            (
+                next.0
+                    .clone()
+                    .into_iter()
+                    .chain(std::iter::once(k.clone()))
+                    .collect(),
+                v,
+            )
+        }));
+        Some(next)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_basic() {
+        enum Entry {
+            Node,
+            Leaf(u32),
+        }
+
+        impl Entry {
+            fn value(&self) -> u32 {
+                match self {
+                    Self::Node => 0,
+                    Self::Leaf(value) => *value,
+                }
+            }
+        }
+
+        let mut tree = Tree::new(Entry::Node);
+        tree.add_child("foo", Entry::Node);
+        tree.add_child("bar", Entry::Leaf(23));
+        tree.at_mut(&["foo"])
+            .unwrap()
+            .add_child("baz", Entry::Leaf(42));
+
+        let bfs = tree.bfs().map(|(path, _)| path).collect::<Vec<_>>();
+        assert_eq!(bfs[0], Vec::<&str>::new());
+        assert!(bfs[1] == vec!["foo"] || bfs[1] == vec!["bar"]);
+        assert!(bfs[2] == vec!["foo"] || bfs[2] == vec!["bar"]);
+        assert!(bfs[1] != bfs[2]);
+        assert_eq!(bfs[3], vec!["foo", "baz"]);
+
+        let dfs = tree.dfs().map(|(path, _)| path).collect::<Vec<_>>();
+        assert_eq!(dfs[0], Vec::<&str>::new());
+        assert!(dfs[1] == vec!["foo"] || dfs[1] == vec!["bar"]);
+        if dfs[1] == vec!["foo"] {
+            assert_eq!(dfs[2], vec!["foo", "baz"]);
+            assert_eq!(dfs[3], vec!["bar"]);
+        } else if dfs[1] == vec!["bar"] {
+            assert_eq!(dfs[2], vec!["foo"]);
+            assert_eq!(dfs[3], vec!["foo", "baz"]);
+        } else {
+            unreachable!()
+        }
+
+        assert_eq!(
+            tree.bfs().map(|(_, tree)| tree.data().value()).sum::<u32>(),
+            65
+        );
+        assert_eq!(
+            tree.dfs().map(|(_, tree)| tree.data().value()).sum::<u32>(),
+            65
+        );
+    }
+}