advent-2022-rust/src/days/day11/mod.rs

use crate::common::parser::{
    empty_lines, eol_terminated, extract_result, ignore, trim0, trim_left1, true_false,
};

use super::template::{DayTrait, ResultType};
use nom::{
    branch::alt,
    bytes::complete::tag,
    character::complete::{char, i64, u32},
    combinator::{map, value, verify},
    error::Error,
    multi::{many1, separated_list1},
    sequence::{preceded, tuple},
    Err, IResult, Parser,
};
use std::{iter::zip, str::FromStr};
use thiserror::Error;

const DAY_NUMBER: usize = 11;

pub struct Day;

impl DayTrait for Day {
    fn get_day_number(&self) -> usize {
        DAY_NUMBER
    }

    fn part1(&self, lines: &str) -> anyhow::Result<ResultType> {
        let troop: Troop = lines.parse()?;
        Ok(ResultType::Integer(troop.play(20)))
    }

    fn part2(&self, lines: &str) -> anyhow::Result<ResultType> {
        let troop: Troop = lines.parse()?;
        Ok(ResultType::Integer(troop.play_again(10_000)))
    }
}

#[derive(Debug, Error)]
enum MonkeyError {
    #[error("Error while parsing: {0}")]
    ParsingError(String),
}

impl From<Err<Error<&str>>> for MonkeyError {
    fn from(value: Err<Error<&str>>) -> Self {
        MonkeyError::ParsingError(value.to_string())
    }
}

#[derive(Debug, PartialEq, Clone, Copy)]
enum Operation {
    Squared,
    Times(i64),
    Plus(i64),
}

impl Operation {
    fn parse(input: &str) -> IResult<&str, Self> {
        let (input, _) = tag("new = old ")(input)?;
        alt((
            value(Operation::Squared, tag("* old")),
            preceded(char('*'), trim0(i64.map(|a| Operation::Times(a)))),
            preceded(char('+'), trim0(i64.map(|a| Operation::Plus(a)))),
        ))(input)
    }

    pub fn calc(&self, old: i64) -> i64 {
        match self {
            Operation::Squared => old.pow(2),
            Operation::Times(value) => old * value,
            Operation::Plus(value) => old + value,
        }
    }
}

#[derive(Debug, PartialEq)]
struct Monkey {
    number: usize,
    items: Vec<i64>,
    operation: Operation,
    divisor: i64,
    good_monkey: usize,
    bad_monkey: usize,
}

impl Monkey {
    fn number_parse(input: &str) -> IResult<&str, usize> {
        let input = ignore(tag("Monkey "))(input)?;
        let (input, number) = trim0(u32.map(|val| val as usize))(input)?;
        let input = ignore(char(':'))(input)?;

        Ok((input, number))
    }

    fn starting_items_parse(input: &str) -> IResult<&str, Vec<i64>> {
        let input = ignore(tag("Starting items:"))(input)?;
        separated_list1(char(','), trim0(i64))(input)
    }

    fn operation_parse(input: &str) -> IResult<&str, Operation> {
        let input = ignore(tag("Operation:"))(input)?;
        trim0(Operation::parse)(input)
    }

    fn divisor_parse(input: &str) -> IResult<&str, i64> {
        let input = ignore(tag("Test: divisible by"))(input)?;
        trim0(i64)(input)
    }

    fn target_parse(want_good: bool) -> impl FnMut(&str) -> IResult<&str, usize> {
        move |input: &str| {
            let input = ignore(tuple((
                tag("If"),
                trim_left1(verify(true_false, |is_good| *is_good == want_good)),
                tag(": throw to monkey"),
            )))(input)?;
            trim0(u32.map(|val| val as usize))(input)
        }
    }

    fn parse(input: &str) -> IResult<&str, Monkey> {
        let input = empty_lines(input)?;
        let (input, number) = eol_terminated(trim0(Monkey::number_parse))(input)?;
        let (input, items) = eol_terminated(trim0(Monkey::starting_items_parse))(input)?;
        let (input, operation) = eol_terminated(trim0(Monkey::operation_parse))(input)?;
        let (input, divisor) = eol_terminated(trim0(Monkey::divisor_parse))(input)?;
        let (input, good_monkey) = eol_terminated(trim0(Monkey::target_parse(true)))(input)?;
        let (input, bad_monkey) = eol_terminated(trim0(Monkey::target_parse(false)))(input)?;

        Ok((
            input,
            Monkey {
                number,
                items,
                operation,
                divisor,
                good_monkey,
                bad_monkey,
            },
        ))
    }

    #[inline]
    pub fn process(&self, value: i64) -> i64 {
        self.operation.calc(value)
    }

    #[inline]
    pub fn check(&self, value: i64) -> bool {
        value % self.divisor == 0
    }

    #[allow(dead_code)]
    fn parse_one(lines: &str) -> Result<Monkey, MonkeyError> {
        Ok(extract_result(Monkey::parse)(lines)?)
    }
}

struct Troop {
    monkeys: Vec<Monkey>,
}

impl FromStr for Troop {
    type Err = MonkeyError;

    fn from_str(lines: &str) -> Result<Self, Self::Err> {
        Ok(extract_result(Troop::parse)(lines)?)
    }
}

impl Troop {
    fn parse(input: &str) -> IResult<&str, Troop> {
        map(many1(Monkey::parse), |monkeys| Troop { monkeys })(input)
    }

    fn do_play<F>(&self, rounds: usize, alter: F) -> i64
    where
        F: Fn(i64) -> i64,
    {
        let mut holding = vec![Vec::new(); self.monkeys.len()];
        let mut inspected = vec![0i64; self.monkeys.len()];
        for (holding, monkey) in zip(&mut holding, &self.monkeys) {
            holding.extend(&monkey.items);
        }
        for _ in 0..rounds {
            for (pos, monkey) in self.monkeys.iter().enumerate() {
                let (good, bad): (Vec<_>, Vec<_>) = holding[pos]
                    .iter()
                    .map(|value| alter(monkey.process(*value)))
                    .partition(|value| monkey.check(*value));
                inspected[pos] += holding[pos].len() as i64;
                holding.get_mut(pos).unwrap().clear();
                holding.get_mut(monkey.good_monkey).unwrap().extend(&good);
                holding.get_mut(monkey.bad_monkey).unwrap().extend(&bad);
            }
        }

        inspected.sort_by(|a, b| b.cmp(a));
        inspected.iter().take(2).product()
    }

    pub fn play(&self, rounds: usize) -> i64 {
        self.do_play(rounds, |v| v / 3)
    }

    pub fn play_again(&self, rounds: usize) -> i64 {
        let divisor: i64 = self.monkeys.iter().map(|monkey| monkey.divisor).product();
        self.do_play(rounds, |v| v % divisor)
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::common::file::read_string;
    use anyhow::Result;

    #[test]
    fn test_parse() -> Result<()> {
        let day = Day {};
        let lines = read_string(day.get_day_number(), "example01.txt")?;
        let expected = Monkey {
            number: 0,
            items: vec![79, 98],
            operation: Operation::Times(19),
            divisor: 23,
            good_monkey: 2,
            bad_monkey: 3,
        };
        let result = Monkey::parse_one(&lines)?;
        assert_eq!(result, expected);

        Ok(())
    }

    #[test]
    fn test_parse_all() -> Result<()> {
        let day = Day {};
        let lines = read_string(day.get_day_number(), "example01.txt")?;
        let expected = 4;
        let result: Troop = lines.parse()?;
        assert_eq!(result.monkeys.len(), expected);

        Ok(())
    }

    #[test]
    fn test_part1() -> Result<()> {
        let day = Day {};
        let lines = read_string(day.get_day_number(), "example01.txt")?;
        let expected = ResultType::Integer(10605);
        let result = day.part1(&lines)?;
        assert_eq!(result, expected);

        Ok(())
    }

    #[test]
    fn test_part2() -> Result<()> {
        let day = Day {};
        let lines = read_string(day.get_day_number(), "example01.txt")?;
        let expected = ResultType::Integer(2713310158);
        let result = day.part2(&lines)?;
        assert_eq!(result, expected);

        Ok(())
    }
}