stackthing.rs

use rand::{Rng, SeedableRng};

#[derive(Clone, Debug)]
enum Operation {
    Plus,
    Multiplication,
}

#[derive(Clone, Debug)]
enum StackElement {
    Operation(Operation),
    Literal(usize),
}

pub fn do_calc(s: &str) -> usize {
    let mut stack = Vec::new();
    let mut byte_iter = s.bytes().peekable();
    while let Some(cur_byte) = byte_iter.next() {
        match cur_byte {
            b'+' => stack.push(StackElement::Operation(Operation::Plus)),
            b'*' => stack.push(StackElement::Operation(Operation::Multiplication)),
            b'0'..=b'9' => {
                let mut cur_lit = (cur_byte - b'0') as usize;
                while matches!(byte_iter.peek().unwrap(), b'0'..=b'9') {
                    cur_lit *= 10;
                    cur_lit += (byte_iter.next().unwrap() - b'0') as usize;
                }
                stack.push(StackElement::Literal(cur_lit));
            }
            b'(' => {} // nothing to do
            b',' => {} // nothing to do
            b')' => {
                // go backwards until we find an operand
                let op = stack
                    .iter()
                    .rev()
                    .find_map(|e| match e {
                        StackElement::Operation(op) => Some(op.clone()),
                        _ => None,
                    })
                    .unwrap();
                let result = match op {
                    Operation::Plus => {
                        let mut result = 0;
                        while let Some(StackElement::Literal(lit)) = stack.pop() {
                            result += lit;
                        }
                        result
                    }
                    Operation::Multiplication => {
                        let mut result: usize = 1;
                        while let Some(StackElement::Literal(lit)) = stack.pop() {
                            result = result
                                .checked_mul(lit)
                                .ok_or_else(|| {
                                    format!("overflow multiplying {} and {}", result, lit)
                                })
                                .unwrap();
                        }
                        result
                    }
                };
                stack.push(StackElement::Literal(result));
            }
            _ => panic!("garbage character {}", cur_byte),
        }
    }
    if let Some(StackElement::Literal(final_result)) = stack.pop() {
        assert!(stack.len() == 0);
        return final_result;
    } else {
        panic!("something went wrong");
    }
}

pub fn write_rec_input<R: Rng>(s: &mut String, rng: &mut R, depth: usize) {
    let child_count = rng.gen_range(2..=10);
    let op = if rng.gen_bool(0.1) { '*' } else { '+' };
    s.push(op);
    s.push('(');
    for i in 0..child_count {
        if rng.gen_bool(0.2) && depth > 1 {
            write_rec_input(s, rng, depth - 1);
        } else {
            use std::fmt::Write;
            write!(s, "{}", rng.gen_range(1..=10)).unwrap();
        }
        if i != child_count - 1 {
            s.push(',');
        }
    }
    s.push(')');
}

fn main() {
    let mut out = String::with_capacity(500);
    let mut rng = rand::rngs::StdRng::seed_from_u64(10);
    write_rec_input(&mut out, &mut rng, 20);
    println!("{}", out);
    println!("{}", do_calc(&out));
}