From ec5b46d793d1e9861e89186369b0958c8d351137 Mon Sep 17 00:00:00 2001
From: jazzpi <jasper@mezzo.de>
Date: Wed, 21 Dec 2022 17:51:41 +0100
Subject: [PATCH] Use trees crate

---
 Cargo.lock       |  7 +++++++
 Cargo.toml       |  1 +
 src/bin/d21p1.rs |  2 +-
 src/day21.rs     | 47 +++++++++++++++++++++++++++++++++--------------
 4 files changed, 42 insertions(+), 15 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index be7672f..409f3ba 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -22,6 +22,7 @@ dependencies = [
  "num",
  "once_cell",
  "regex",
+ "trees",
 ]
 
 [[package]]
@@ -340,6 +341,12 @@ dependencies = [
  "winapi-util",
 ]
 
+[[package]]
+name = "trees"
+version = "0.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0de5f738ceab88e2491a94ddc33c3feeadfa95fedc60363ef110845df12f3878"
+
 [[package]]
 name = "winapi"
 version = "0.3.9"
diff --git a/Cargo.toml b/Cargo.toml
index 43b7e0a..9b0b0c4 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -13,6 +13,7 @@ env_logger = "^0.10"
 num = "^0.1"
 lazy_static = "^1.4"
 once_cell = "^1.16"
+trees = "^0.4"
 
 [[bin]]
 name = "d1p1"
diff --git a/src/bin/d21p1.rs b/src/bin/d21p1.rs
index 7364953..bd9ff0f 100644
--- a/src/bin/d21p1.rs
+++ b/src/bin/d21p1.rs
@@ -3,5 +3,5 @@ use aoc22::{day21, util};
 pub fn main() {
     let monkeys = day21::parse_monkeys(&util::parse_input());
 
-    println!("root yells: {}", day21::yell(&monkeys, "root"));
+    println!("root yells: {}", day21::yell(monkeys.root()));
 }
diff --git a/src/day21.rs b/src/day21.rs
index 537823c..e2a4004 100644
--- a/src/day21.rs
+++ b/src/day21.rs
@@ -1,8 +1,9 @@
 use std::collections::HashMap;
 
 use regex::Regex;
+use trees::{tr, Node, Tree};
 
-#[derive(Debug)]
+#[derive(Debug, Clone, Copy)]
 pub enum Operation {
     Add,
     Sub,
@@ -22,14 +23,14 @@ impl Operation {
     }
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub enum Monkey {
     Number(isize),
     Calculate(String, String, Operation),
 }
 
-pub fn parse_monkeys(input: &String) -> HashMap<String, Monkey> {
-    let mut result = HashMap::new();
+pub fn parse_monkeys(input: &String) -> Tree<Monkey> {
+    let mut monkeys = HashMap::new();
 
     let re = Regex::new(r"^(\w+): (?:(\d+)|(\w+) ([+\-*/]) (\w+))$").unwrap();
     for line in input.lines() {
@@ -44,21 +45,39 @@ pub fn parse_monkeys(input: &String) -> HashMap<String, Monkey> {
                 Operation::from(captures.get(4).unwrap().as_str()),
             )
         };
-        result.insert(name.to_owned(), monkey);
+        monkeys.insert(name.to_owned(), monkey);
     }
 
-    result
+    let tree = to_tree(&monkeys, "root");
+
+    tree
 }
 
-pub fn yell(monkeys: &HashMap<String, Monkey>, monkey: &str) -> isize {
+fn to_tree(monkeys: &HashMap<String, Monkey>, monkey: &str) -> Tree<Monkey> {
     let monkey = &monkeys[monkey];
-    match monkey {
+    let mut tree = tr(monkey.clone());
+
+    if let Monkey::Calculate(a, b, _) = monkey {
+        let forest = -to_tree(monkeys, a) - to_tree(monkeys, b);
+        tree.append(forest);
+    }
+
+    tree
+}
+
+pub fn yell(monkey: &Node<Monkey>) -> isize {
+    match monkey.data() {
         Monkey::Number(n) => *n,
-        Monkey::Calculate(a, b, op) => match op {
-            Operation::Add => yell(monkeys, a) + yell(monkeys, b),
-            Operation::Sub => yell(monkeys, a) - yell(monkeys, b),
-            Operation::Mul => yell(monkeys, a) * yell(monkeys, b),
-            Operation::Div => yell(monkeys, a) / yell(monkeys, b),
-        },
+        Monkey::Calculate(_, _, op) => {
+            let mut children = monkey.iter();
+            let a = children.next().unwrap();
+            let b = children.next().unwrap();
+            match op {
+                Operation::Add => yell(a) + yell(b),
+                Operation::Sub => yell(a) - yell(b),
+                Operation::Mul => yell(a) * yell(b),
+                Operation::Div => yell(a) / yell(b),
+            }
+        }
     }
 }