AOC/6 (Rust)

main.rs

  #![feature(test)]

extern crate fxhash;
extern crate test;

#[derive(Clone, Eq, PartialEq)]
struct MemoryBank {
    banks: Vec<i32>,
}

impl MemoryBank {
    fn new<T: Into<Vec<i32>>>(banks: T) -> Self {
        Self { banks: banks.into() }
    }

    fn cycle(&mut self) {
        let max = self.banks.iter().enumerate().fold(
            None,
            |a, (idx, &value)| match a {
                None => Some((idx, value)),
                Some(a) => {
                    if value > a.1 {
                        Some((idx, value))
                    } else {
                        Some(a)
                    }
                }
            },
        );

        if let Some((max_idx, mut distribute)) = max {
            self.banks[max_idx] = 0;
            for entry in self.banks.iter_mut().skip(max_idx + 1) {
                if distribute > 0 {
                    *entry += 1;
                    distribute -= 1;
                } else {
                    break;
                }
            }

            while distribute > 0 {
                for entry in self.banks.iter_mut() {
                    if distribute > 0 {
                        *entry += 1;
                        distribute -= 1;
                    } else {
                        break;
                    }
                }
            }
        }
    }

    /// Spits out a "memory-efficient," hashable representation of the memory bank state.
    ///
    /// ...also known as a string, ok? I think this will be smaller than the actual vector
    /// in the average case. /shrug
    fn state_key(&self) -> StateKey {
        let key = self.banks.iter().fold(self.banks.len(), |a, &b| {
            a.wrapping_mul(17) + b as usize
        });

        StateKey(key)
    }
}

#[derive(Eq, PartialEq, Hash)]
struct StateKey(usize);

fn main() {
    println!("{}", cycles_until_repeat(input()));
    println!("{}", cycle_length(input()));
}

fn cycles_until_repeat<T: Into<Vec<i32>>>(banks: T) -> usize {
    use fxhash::FxHashSet;

    let target_bank = MemoryBank::new(banks);
    let mut bank = target_bank.clone();

    let mut counter = 0;
    let mut states = FxHashSet::default();

    while states.insert(bank.state_key()) {
        bank.cycle();
        counter += 1;
    }

    counter
}

fn cycle_length<T: Into<Vec<i32>>>(banks: T) -> usize {
    use fxhash::FxHashMap;

    let mut bank = MemoryBank::new(banks);
    let mut counter = 0;
    let mut states = FxHashMap::default();

    loop {
        if let Some(x) = states.insert(bank.state_key(), counter) {
            return counter - x;
        }

        bank.cycle();
        counter += 1;
    }
}

fn input() -> Vec<i32> {
    vec![11, 11, 13, 7, 0, 15, 5, 5, 4, 4, 1, 1, 7, 1, 15, 11]
}

#[cfg(test)]
mod tests {
    use test::{self, Bencher};

    #[test]
    fn part_1_works() {
        assert_eq!(5, super::cycles_until_repeat(vec![0, 2, 7, 0]));
    }

    #[test]
    fn part_2_works() {
        assert_eq!(4, super::cycle_length(vec![0, 2, 7, 0]));
    }

    #[bench]
    fn part_1_bench(b: &mut Bencher) {
        b.iter(|| {
            test::black_box(super::cycles_until_repeat(super::input()));
        });
    }

    #[bench]
    fn part_2_bench(b: &mut Bencher) {
        b.iter(|| { test::black_box(super::cycle_length(super::input())); });
    }
}