day18.rs

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum Token {
    Number(i64),
    Add,
    Mul,
    OParen,
    CParen,
}

use Token::*;

#[aoc_generator(day18)]
fn input_generator(input: &str) -> Vec<Vec<Token>> {
    input.lines().map(parse_expr).collect()
}

fn parse_expr(expr: &str) -> Vec<Token> {
    let mut tokens = vec![];
    let mut index = 0;
    while index < expr.len() {
        match expr.as_bytes()[index] {
            b' ' => (),
            b'+' => tokens.push(Add),
            b'*' => tokens.push(Mul),
            b'(' => tokens.push(OParen),
            b')' => tokens.push(CParen),
            _ => tokens.push(Number(parse_number(expr, &mut index))),
        }
        index += 1;
    }
    tokens
}

fn parse_number(input: &str, index: &mut usize) -> i64 {
    let first = *index;
    while *index + 1 < input.len() && (b'0'..=b'9').contains(&input.as_bytes()[*index + 1]) {
        *index += 1;
    }
    unsafe {
        std::str::from_utf8_unchecked(&input.as_bytes()[first..=*index])
            .parse()
            .unwrap()
    }
}

#[aoc(day18, part1)]
fn part1(exprs: &[Vec<Token>]) -> i64 {
    exprs.iter().map(|t| eval(&*t)).sum()
}

#[aoc(day18, part2)]
fn part2(exprs: &[Vec<Token>]) -> i64 {
    exprs.iter().map(|t| eval2(&*t)).sum()
}

fn eval(tokens: &[Token]) -> i64 {
    let mut stack = vec![];
    for &t in tokens {
        if t == CParen {
            let v = stack.pop().unwrap();
            stack.pop();
            stack.push(v);
        } else {
            stack.push(t);
        }
        if stack.len() >= 3 {
            match stack[stack.len() - 3..stack.len()] {
                [Number(a), Add, Number(b)] => {
                    stack.truncate(stack.len() - 3);
                    stack.push(Number(a + b))
                }
                [Number(a), Mul, Number(b)] => {
                    stack.truncate(stack.len() - 3);
                    stack.push(Number(a * b))
                }
                _ => (),
            }
        }
    }
    if let Number(a) = stack.pop().unwrap() {
        a
    } else {
        unreachable!()
    }
}

fn eval2(tokens: &[Token]) -> i64 {
    let mut stack = vec![];
    for &t in tokens {
        if t == CParen {
            while stack[stack.len() - 2] != OParen {
                match stack[stack.len() - 3..stack.len()] {
                    [Number(a), Mul, Number(b)] => {
                        stack.truncate(stack.len() - 3);
                        stack.push(Number(a * b));
                    }
                    _ => unreachable!(),
                }
            }
            let v = stack.pop().unwrap();
            stack.pop();
            stack.push(v);
        } else {
            stack.push(t);
        }
        if stack.len() >= 3 {
            match stack[stack.len() - 3..stack.len()] {
                [Number(a), Add, Number(b)] => {
                    stack.truncate(stack.len() - 3);
                    stack.push(Number(a + b))
                }
                _ => (),
            }
        }
    }
    stack
        .into_iter()
        .step_by(2)
        .map(|t| match t {
            Number(a) => a,
            _ => unreachable!(),
        })
        .product()
}