Advent of Code 2021 - Day 18: Snailfish

snailfish.rs

fn magnitude(numbers: &[&str]) -> u64 {
    let total = SNumber::from(numbers[0]);

    numbers[1..]
        .iter()
        .fold(total, |total, &number| total + SNumber::from(number))
        .magnitude()
}

fn max_magnitude(numbers: &[&str]) -> u64 {
    (0..numbers.len())
        .flat_map(|a| {
            (0..numbers.len()).map(move |b| {
                if a == b {
                    u64::MIN
                } else {
                    (SNumber::from(numbers[a]) + SNumber::from(numbers[b])).magnitude()
                }
            })
        })
        .max()
        .unwrap()
}

enum SNumber {
    Number(u64),
    Pair(Box<(SNumber, SNumber)>),
}

impl Add for SNumber {
    type Output = SNumber;

    fn add(self, rhs: Self) -> Self::Output {
        let mut sum = SNumber::Pair(Box::new((self, rhs)));
        sum.reduce();
        sum
    }
}

impl From<&str> for SNumber {
    fn from(value: &str) -> Self {
        let mut chars = value.chars().collect::<VecDeque<_>>();

        Self::parse_pair(&mut chars)
    }
}

impl SNumber {
    fn parse_pair(chars: &mut VecDeque<char>) -> SNumber {
        chars.pop_front().unwrap();

        let left = match chars[0] {
            '[' => Self::parse_pair(chars),
            '0'..='9' => Self::parse_number(chars),
            _ => panic!("Unrecognised symbol!"),
        };

        chars.pop_front().unwrap();

        let right = match chars[0] {
            '[' => Self::parse_pair(chars),
            '0'..='9' => Self::parse_number(chars),
            _ => panic!("Unrecognised symbol!"),
        };

        chars.pop_front().unwrap();

        SNumber::Pair(Box::new((left, right)))
    }

    fn parse_number(chars: &mut VecDeque<char>) -> SNumber {
        SNumber::Number(chars.pop_front().unwrap().to_digit(10).unwrap() as u64)
    }

    fn magnitude(&self) -> u64 {
        match self {
            SNumber::Number(value) => *value,
            SNumber::Pair(pair) => pair.0.magnitude() * 3 + pair.1.magnitude() * 2,
        }
    }

    fn reduce(&mut self) {
        while self.explode() || self.split() {}
    }

    fn explode(&mut self) -> bool {
        if let Some((pos, left, right)) = self.find_explode(0, &mut 0) {
            if pos > 0 {
                self.add_to(pos - 1, left, &mut 0);
            }
            self.add_to(pos + 1, right, &mut 0);
            return true;
        }
        false
    }

    fn find_explode(&mut self, depth: usize, current_pos: &mut usize) -> Option<(usize, u64, u64)> {
        match self {
            SNumber::Number(_) => {
                *current_pos += 1;
                None
            }
            SNumber::Pair(pair) => match (&pair.0, &pair.1) {
                (&SNumber::Number(left), &SNumber::Number(right)) if depth == 4 => {
                    *self = SNumber::Number(0);
                    Some((*current_pos, left, right))
                }
                _ => {
                    if let Some(result) = pair.0.find_explode(depth + 1, current_pos) {
                        return Some(result);
                    }
                    pair.1.find_explode(depth + 1, current_pos)
                }
            },
        }
    }

    fn add_to(&mut self, pos: usize, value: u64, current_pos: &mut usize) {
        match self {
            SNumber::Number(n) => {
                if pos == *current_pos {
                    *self = SNumber::Number(*n + value);
                }
                *current_pos += 1;
            }
            SNumber::Pair(pair) => {
                pair.0.add_to(pos, value, current_pos);
                pair.1.add_to(pos, value, current_pos);
            }
        }
    }

    fn split(&mut self) -> bool {
        match self {
            SNumber::Number(value) => {
                if *value >= 10 {
                    let left = *value / 2;
                    let right = *value - left;
                    *self =
                        SNumber::Pair(Box::new((SNumber::Number(left), SNumber::Number(right))));
                    return true;
                }
                false
            }
            SNumber::Pair(pair) => {
                if pair.0.split() {
                    return true;
                }

                pair.1.split()
            }
        }
    }
}