2020: day14: part 2

aoc2020/input/day14_provided2.txt

@@ -0,0 +1,4 @@
+mask = X1001X
+mem[42] = 100
+mask = 00X0XX
+mem[26] = 1

aoc2020/src/day14.rs

@@ -9,6 +9,7 @@ pub fn run() -> aoc::Result<String> {
     let mut res = String::with_capacity(128);
 
     writeln!(res, "part 1: {}", part1(INPUT)?)?;
+    writeln!(res, "part 2: {}", part2(INPUT)?)?;
 
     Ok(res)
 }
@@ -21,6 +22,14 @@ fn part1(input: &str) -> aoc::Result<u64> {
     Ok(program.memory_sum())
 }
 
+fn part2(input: &str) -> aoc::Result<u64> {
+    let mut program: Program = input.parse()?;
+
+    program.run_part2()?;
+
+    Ok(program.memory_sum())
+}
+
 #[derive(Debug, Clone, Copy)]
 enum Mask {
     Floating,
@@ -28,6 +37,106 @@ enum Mask {
     Zero,
 }
 
+#[derive(Debug)]
+struct FloatingIterator<'a> {
+    masks: &'a [Mask],
+    current_value: usize,
+    done: bool,
+
+    stack: Vec<(usize, FloatingState)>,
+}
+
+impl<'a> FloatingIterator<'a> {
+    fn new(masks: &'a [Mask], mut n: usize) -> Self {
+        for (offset, mask) in masks.iter().enumerate() {
+            if let Mask::One = mask {
+                n |= 1 << offset;
+            }
+        }
+
+        Self {
+            masks,
+            current_value: n,
+            done: false,
+            stack: Vec::new(),
+        }
+    }
+
+    fn find_next_floating(masks: &[Mask], offset: usize) -> Option<usize> {
+        masks
+            .iter()
+            .enumerate()
+            .skip(offset)
+            .find(|(_, m)| matches!(m, Mask::Floating))
+            .map(|(idx, _)| idx)
+    }
+}
+
+#[derive(Debug)]
+enum FloatingState {
+    Unapplied,
+    AppliedZero,
+    AppliedZeroAndOne,
+}
+
+impl<'a> Iterator for FloatingIterator<'a> {
+    type Item = usize;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.done {
+            return None;
+        }
+
+        if self.stack.is_empty() {
+            // initialize stack with first floating iterator
+            let next = FloatingIterator::find_next_floating(self.masks, 0);
+
+            match next {
+                Some(offset) => self.stack.push((offset, FloatingState::Unapplied)),
+                None => {
+                    // there are no floating bits in this mask, we can return the value directly
+                    self.done = true;
+                    return Some(self.current_value);
+                }
+            }
+        }
+
+        loop {
+            let (offset, state) = self.stack.last_mut().unwrap();
+            match state {
+                FloatingState::Unapplied => {
+                    self.current_value &= !(1 << *offset);
+                    *state = FloatingState::AppliedZero;
+                }
+                FloatingState::AppliedZero => {
+                    self.current_value |= 1 << *offset;
+                    *state = FloatingState::AppliedZeroAndOne;
+                }
+                FloatingState::AppliedZeroAndOne => {
+                    self.stack.pop();
+
+                    if self.stack.is_empty() {
+                        // we've computed all possibilities, we can just stop now
+                        self.done = true;
+                        return None;
+                    }
+
+                    continue;
+                }
+            }
+
+            // we've applied our current mask transform, now we "recurse" and find the next one
+            match FloatingIterator::find_next_floating(self.masks, *offset + 1) {
+                Some(offset) => self.stack.push((offset, FloatingState::Unapplied)),
+                None => {
+                    // we were the last Floating mask, we can return the produced value
+                    return Some(self.current_value);
+                }
+            }
+        }
+    }
+}
+
 #[derive(Debug, Clone)]
 struct BitMask {
     masks: Vec<Mask>,
@@ -46,6 +155,10 @@ impl BitMask {
 
         n
     }
+
+    fn apply(&self, n: usize) -> impl Iterator<Item = usize> + '_ {
+        FloatingIterator::new(&self.masks, n)
+    }
 }
 
 impl std::str::FromStr for BitMask {
@@ -146,6 +259,27 @@ impl Program {
         Ok(())
     }
 
+    fn run_part2(&mut self) -> aoc::Result<()> {
+        for inst in &self.instructions {
+            match inst {
+                Instruction::ChangeMask(bitmask) => self.current_mask = Some(bitmask.clone()),
+
+                Instruction::MemWrite { offset, value } => match &self.current_mask {
+                    Some(bitmask) => {
+                        for offset in bitmask.apply(*offset) {
+                            self.memory.insert(offset, *value);
+                        }
+                    }
+                    None => {
+                        return Err(err!("tried to execute MemWrite but mask isn't initialized"))
+                    }
+                },
+            }
+        }
+
+        Ok(())
+    }
+
     fn memory_sum(&self) -> u64 {
         self.memory.iter().map(|(_, value)| value).sum()
     }
@@ -172,15 +306,26 @@ impl std::str::FromStr for Program {
 mod tests {
     use super::*;
 
-    const PROVIDED: &str = include_str!("../input/day14_provided.txt");
+    const PROVIDED1: &str = include_str!("../input/day14_provided1.txt");
+    const PROVIDED2: &str = include_str!("../input/day14_provided2.txt");
 
     #[test]
     fn part1_provided() {
-        assert_eq!(part1(PROVIDED).unwrap(), 165);
+        assert_eq!(part1(PROVIDED1).unwrap(), 165);
     }
 
     #[test]
     fn part1_real() {
         assert_eq!(part1(INPUT).unwrap(), 4297467072083);
     }
+
+    #[test]
+    fn part2_provided() {
+        assert_eq!(part2(PROVIDED2).unwrap(), 208);
+    }
+
+    #[test]
+    fn part2_real() {
+        assert_eq!(part2(INPUT).unwrap(), 5030603328768);
+    }
 }