adventofcode2022/day11/common.rs

// Copyright 2022 Christian Ulrich
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

use std::io;

#[derive(Debug, Clone)]
enum Operation {
    Add { arg: u64 },
    Multiply { arg: u64 },
    Square,
}

#[derive(Debug, Clone)]
pub struct Monkey {
    items: Vec<u64>,
    operation: Operation,
    modulus: u64,
    target_devisible: usize,
    target_indevisible: usize,
    item_count: usize,
}

pub struct Game {
    monkeys: Vec<Monkey>,
    modulus_lcm: u64,
    divisor: u64,
}

pub struct Round {
    monkeys: Vec<Monkey>,
}

impl Monkey {
    fn from_stdin() -> Result<Option<Self>, &'static str> {
        const ERROR: &str = "cannot parse monkey";
        let mut line = String::new();
        let _ = io::stdin().read_line(&mut line).map_err(|_| ERROR)?;
        if line.len() == 0 {
            return Ok(None);
        }
        if !line.starts_with("Monkey ") {
            return Err(ERROR);
        }
        line.clear();

        // parse starting items
        let _ = io::stdin().read_line(&mut line).map_err(|_| ERROR)?;
        if !line.starts_with("  Starting items: ") {
            return Err(ERROR);
        }
        let mut items = vec![];
        for part in line[18..].split(',') {
            items.push(part.trim().parse::<u64>().map_err(|_| ERROR)?);
        }
        line.clear();

        // parse operation
        let _ = io::stdin().read_line(&mut line).map_err(|_| ERROR)?;
        if !line.starts_with("  Operation: new = old ") {
            return Err(ERROR);
        }
        let operation = match line[23..].trim() {
            "* old" => Operation::Square,
            arg if { arg.starts_with("+ ") } => Operation::Add {
                arg: arg[2..].parse().map_err(|_| ERROR)?,
            },
            arg if { arg.starts_with("* ") } => Operation::Multiply {
                arg: arg[2..].parse().map_err(|_| ERROR)?,
            },
            _ => return Err(ERROR),
        };
        line.clear();

        // parse modulus
        let _ = io::stdin().read_line(&mut line).map_err(|_| ERROR)?;
        if !line.starts_with("  Test: divisible by ") {
            return Err(ERROR);
        }
        let modulus = line[21..].trim().parse().map_err(|_| ERROR)?;
        line.clear();

        // parse action if devisible
        let _ = io::stdin().read_line(&mut line).map_err(|_| ERROR)?;
        if !line.starts_with("    If true: throw to monkey ") {
            return Err(ERROR);
        }
        let target_devisible = line[29..].trim().parse().map_err(|_| ERROR)?;
        line.clear();

        // parse action if indevisible
        let _ = io::stdin().read_line(&mut line).map_err(|_| ERROR)?;
        if !line.starts_with("    If false: throw to monkey ") {
            return Err(ERROR);
        }
        let target_indevisible = line[30..].trim().parse().map_err(|_| ERROR)?;
        line.clear();
        let _ = io::stdin().read_line(&mut line).map_err(|_| ERROR)?;
        if line.trim().len() > 0 {
            return Err(ERROR);
        }
        Ok(Some(Self {
            items: items,
            operation: operation,
            modulus: modulus,
            target_devisible: target_devisible,
            target_indevisible: target_indevisible,
            item_count: 0,
        }))
    }

    fn throw_item(&mut self, divisor: u64) -> Result<Option<(u64, usize)>, &'static str> {
        const ERROR: &str = "integer overflow";
        if self.items.len() == 0 {
            return Ok(None);
        }
        let mut item = self.items.remove(0);
        match self.operation {
            Operation::Add { arg } => item = item.checked_add(arg).ok_or(ERROR)?,
            Operation::Multiply { arg } => item = item.checked_mul(arg).ok_or(ERROR)?,
            Operation::Square => item = item.checked_mul(item).ok_or(ERROR)?,
        };
        item /= divisor;
        let target = match item % self.modulus {
            0 => self.target_devisible,
            _ => self.target_indevisible,
        };
        self.item_count += 1;
        Ok(Some((item, target)))
    }

    pub fn add_item(&mut self, item: u64) {
        self.items.push(item);
    }

    pub fn items(&self) -> &[u64] {
        &self.items[..]
    }

    pub fn item_count(&self) -> usize {
        self.item_count
    }
}

impl Game {
    pub fn from_stdin(divisor: u64) -> Result<Self, &'static str> {
        let mut monkeys = vec![];
        while let Some(monkey) = Monkey::from_stdin()? {
            monkeys.push(monkey);
        }
        let modulus_lcm = monkeys.iter().map(|m| m.modulus).product();
        Ok(Self {
            monkeys: monkeys,
            modulus_lcm: modulus_lcm,
            divisor: divisor,
        })
    }

    pub fn monkeys(&self) -> &[Monkey] {
        &self.monkeys[..]
    }
}

impl Iterator for Game {
    type Item = Result<Round, &'static str>;

    fn next<'a>(&'a mut self) -> Option<Self::Item> {
        for i in 0..self.monkeys.len() {
            loop {
                match self.monkeys[i].throw_item(self.divisor) {
                    Ok(None) => break,
                    Ok(Some((item, target))) => {
                        self.monkeys[target].add_item(item % self.modulus_lcm);
                    }
                    Err(e) => return Some(Err(e)),
                }
            }
        }
        Some(Ok(Round {
            monkeys: self.monkeys.clone(),
        }))
    }
}

impl Round {
    pub fn monkeys(&self) -> &[Monkey] {
        &self.monkeys[..]
    }
}