Advent of Code 2022 solution for day 21

part1.rs

use std::collections::HashMap;
use num_rational::Rational64;

pub type Number = Rational64;

#[derive(Debug)]
pub enum Expr<'a> {
    Unknown,
    Number(Number),
    Add(&'a str, &'a str),
    Sub(&'a str, &'a str),
    Mul(&'a str, &'a str),
    Div(&'a str, &'a str),
}
use Expr::*;

pub fn parse_expr(s: &str) -> Expr {
    if let Ok(n) = s.parse::<Number>() {
        return Number(n);
    }
    let (a, rest) = s.split_once(' ').unwrap();
    let (op, b) = rest.split_once(' ').unwrap();
    match op {
        "+" => Add(a, b),
        "-" => Sub(a, b),
        "*" => Mul(a, b),
        "/" => Div(a, b),
        _ => panic!("bad operation"),
    }
}

pub type Monkeys<'a> = HashMap<&'a str, Expr<'a>>;

pub fn parse_monkeys(input: &str) -> Monkeys {
    input.lines().map(|line| {
        let (name, expr) = line.split_once(": ").unwrap();
        (name, parse_expr(expr))
    }).collect()
}

pub fn eval<'a>(monkeys: &Monkeys<'a>, which: &'a str, cache: &mut HashMap<&'a str, Number>) -> Number {
    match cache.get(which) {
        Some(n) => *n,
        None => {
            let n = match monkeys.get(which) {
                Some(Number(n)) => *n,
                Some(Add(a, b)) => eval(monkeys, a, cache) + eval(monkeys, b, cache),
                Some(Sub(a, b)) => eval(monkeys, a, cache) - eval(monkeys, b, cache),
                Some(Mul(a, b)) => eval(monkeys, a, cache) * eval(monkeys, b, cache),
                Some(Div(a, b)) => eval(monkeys, a, cache) / eval(monkeys, b, cache),
                Some(Unknown) => panic!("this is for part 2"),
                None => panic!("bad monkey reference"),
            };
            cache.insert(which, n);
            n
        }
    }
}

pub fn part1(input: &str) -> Number {
    let monkeys = parse_monkeys(input);
    eval(&monkeys, "root", &mut HashMap::new())
}

part2.rs

use crate::part1::{parse_monkeys, Expr, Monkeys, Number};
use std::collections::HashMap;
use std::ops::{Add, Div, Mul, Sub};
use num_traits::sign::Signed;
use Expr::*;

#[derive(Copy, Clone)]
pub struct Formula {
    // let's try a * x + b
    a: Number,
    b: Number,
}

impl std::fmt::Debug for Formula {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let zero = 0.into();
        let one: Number = 1.into();
        match (self.a, self.b) {
            (a, b) if a == zero => write!(f, "{}", b),
            (a, b) if a == one && b == zero => write!(f, "x"),
            (a, b) if a == one && b > zero => write!(f, "x + {}", b),
            (a, b) if a == one && b < zero => write!(f, "x - {}", b.abs()),
            (a, b) if b == zero => write!(f, "{} * x", a),
            (a, b) if b > zero => write!(f, "{} * x + {}", a, b),
            (a, b) if b < zero => write!(f, "{} * x - {}", a, b.abs()),
            _ => unreachable!(),
        }
    }
}

impl Add<Formula> for Formula {
    type Output = Formula;

    fn add(self, other: Formula) -> Self::Output {
        Formula {
            a: self.a + other.a,
            b: self.b + other.b,
        }
    }
}

impl Sub<Formula> for Formula {
    type Output = Formula;

    fn sub(self, other: Formula) -> Self::Output {
        Formula {
            a: self.a - other.a,
            b: self.b - other.b,
        }
    }
}

impl Mul<Formula> for Formula {
    type Output = Formula;

    fn mul(self, other: Formula) -> Self::Output {
        // (ax + b) * (cx + d) = acx**2 + (ad + bc) x + bd
        let (a, b) = (self.a, self.b);
        let (c, d) = (other.a, other.b);
        if a * c != 0.into() {
            panic!("unexpected quadratic");
        }
        Formula {
            a: (a * d + b * c),
            b: b * d,
        }
    }
}

impl Div<Formula> for Formula {
    type Output = Formula;

    fn div(self, other: Formula) -> Self::Output {
        let (a, b) = (self.a, self.b);
        let (c, d) = (other.a, other.b);
        if c != 0.into() {
            panic!("unexpected division by x");
        }
        Formula { a: a / d, b: b / d }
    }
}

pub fn eval<'a>(
    monkeys: &Monkeys<'a>,
    which: &'a str,
    cache: &mut HashMap<&'a str, Formula>,
    debug: bool,
) -> Formula {
    match cache.get(which) {
        Some(n) => *n,
        None => {
            let n = match monkeys.get(which) {
                Some(Number(n)) => Formula { a: 0.into(), b: *n },
                Some(Add(a, b)) => eval(monkeys, a, cache, debug) + eval(monkeys, b, cache, debug),
                Some(Sub(a, b)) => eval(monkeys, a, cache, debug) - eval(monkeys, b, cache, debug),
                Some(Mul(a, b)) => eval(monkeys, a, cache, debug) * eval(monkeys, b, cache, debug),
                Some(Div(a, b)) => eval(monkeys, a, cache, debug) / eval(monkeys, b, cache, debug),
                Some(Unknown) => Formula {
                    a: 1.into(),
                    b: 0.into(),
                },
                None => panic!("bad monkey reference"),
            };
            if debug {
                println!("Evaluating {}, which is {:?}", which, monkeys.get(which));
                println!("got {:?}", n);
            }
            cache.insert(which, n);
            n
        }
    }
}

pub fn solve(input: &str, debug: bool) -> Number {
    let mut monkeys = parse_monkeys(input);
    let (a, b) = match monkeys.get("root") {
        Some(Add(a, b)) => (a, b),
        Some(Sub(a, b)) => (a, b),
        Some(Mul(a, b)) => (a, b),
        Some(Div(a, b)) => (a, b),
        _ => panic!("root monkey is not operating on two numbers"),
    };
    monkeys.insert("root", Sub(a, b));
    monkeys.insert("humn", Unknown);
    let equation = eval(&monkeys, "root", &mut HashMap::new(), debug);
    dbg!(equation);
    // at this point we know that ax + b = 0, so x = -b/a
    -equation.b / equation.a
}

#[allow(dead_code)]
pub fn part2(input: &str) -> Number {
    solve(input, false)
}