Complete day 14

day-14/src/main.rs

@@ -1,6 +1,5 @@
-use std::convert::TryFrom;
+use std::collections::HashMap;
 use std::fs;
-use std::iter;
 
 use itertools::Itertools;
 use lazy_static::lazy_static;
@@ -10,22 +9,26 @@ const MASK_36_BITS: u64 = 0xFFFFFFFFF;
 
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     let input = fs::read_to_string("input")?;
-    let program = parse_program_chip1(&input)?;
+    let program = parse_program(&input)?;
 
     // Part 1
     let memory = run_program_chip1(&program);
     println!("{}", memory.iter().map(|(_, value)| value).sum::<u64>());
 
+    // Part 2
+    let memory = run_program_chip2(&program);
+    println!("{}", memory.iter().map(|(_, value)| value).sum::<u64>());
+
     Ok(())
 }
 
 #[derive(Copy, Clone, Debug)]
-enum Chip1Instruction {
-    Mask(u64, u64),
-    Mem(usize, u64),
+enum Instruction {
+    Mask(u64, u64, u64),
+    Mem(u64, u64),
 }
 
-impl Chip1Instruction {
+impl Instruction {
     fn new(instruction: &str) -> Result<Self, String> {
         lazy_static! {
             static ref INSTRUCTION_RE: Regex =
@@ -44,7 +47,7 @@ impl Chip1Instruction {
                 let or = u64::from_str_radix(&caps[3].replace("X", "0"), 2)
                     .map_err(|e| format!("Could not parse or: {}", e))?;
 
-                Chip1Instruction::Mask(and, or)
+                Instruction::Mask(and, or, and ^ or)
             }
             "mem" => {
                 let location = caps[2]
@@ -54,30 +57,28 @@ impl Chip1Instruction {
                     .parse()
                     .map_err(|e| format!("Could not parse value: {}", e))?;
 
-                Chip1Instruction::Mem(location, value)
+                Instruction::Mem(location, value)
             }
             _ => Err("Invalid instruction")?,
         })
     }
 }
 
-fn parse_program_chip1(input: &str) -> Result<Vec<Chip1Instruction>, String> {
-    input.lines().map(Chip1Instruction::new).collect()
+fn parse_program(input: &str) -> Result<Vec<Instruction>, String> {
+    input.lines().map(Instruction::new).collect()
 }
 
-fn run_program_chip1(program: &Vec<Chip1Instruction>) -> Vec<(u64, u64)> {
+fn run_program_chip1(program: &Vec<Instruction>) -> Vec<(u64, u64)> {
     let mut mask = (u64::max_value() & MASK_36_BITS, 0);
 
     program
         .iter()
         .filter_map(|instruction| match instruction {
-            Chip1Instruction::Mask(and, or) => {
+            Instruction::Mask(and, or, _) => {
                 mask = (*and, *or);
                 None
             }
-            Chip1Instruction::Mem(address, value) => {
-                Some((*address as u64, value & mask.0 | mask.1))
-            }
+            Instruction::Mem(address, value) => Some((*address as u64, value & mask.0 | mask.1)),
         })
         .collect::<Vec<(u64, u64)>>()
         .into_iter()
@@ -86,86 +87,36 @@ fn run_program_chip1(program: &Vec<Chip1Instruction>) -> Vec<(u64, u64)> {
         .collect()
 }
 
-#[derive(Clone, Debug)]
-enum Chip2Instruction {
-    Mask(u64, u64),
-    Mem(usize, u64),
-}
+fn run_program_chip2(program: &Vec<Instruction>) -> Vec<(u64, u64)> {
+    let mut memory = HashMap::new();
+    let mut mask = (0, 0);
 
-impl Chip2Instruction {
-    fn new(instruction: &str) -> Result<Self, String> {
-        lazy_static! {
-            static ref INSTRUCTION_RE: Regex =
-                Regex::new(r"([[:alpha:]]+?)\[?(\d+?)?\]? = ([0-9X]+)")
-                    .expect("Regex should compile");
+    for instruction in program {
+        match instruction {
+            Instruction::Mask(_, or, floating) => {
+                mask = (*or, *floating);
+            }
+            Instruction::Mem(address, value) => {
+                // We set floating bits to 0 here, so that they can be
+                // toggled with addition
+                let base = (address | mask.0) & !mask.1 & MASK_36_BITS;
+
+                let mut floating = !mask.1 & MASK_36_BITS;
+                while floating <= MASK_36_BITS {
+                    // For each unset bit in the floating mask, we
+                    // toggle the bit in the address
+                    memory.insert(base | (!floating & MASK_36_BITS), *value);
+                    floating += 1;
+
+                    // Since add does overflows, we reset 1s to 0s if
+                    // they're unset in the original mask
+                    floating |= !mask.1 & MASK_36_BITS;
+                }
+            }
         }
-
-        let caps = INSTRUCTION_RE
-            .captures(instruction)
-            .ok_or("Instruction did not match")?;
-
-        Ok(match &caps[1] {
-            "mask" => {
-                let or = u64::from_str_radix(&caps[3].replace("X", "0"), 2)
-                    .map_err(|e| format!("Could not parse or: {}", e))?;
-
-                let floating = u64::from_str_radix(&caps[3].replace("X", "0"), 2)
-                    .map_err(|e| format!("Could not parse floating: {}", e))?;
-
-                Chip2Instruction::Mask(or, floating)
-            }
-            "mem" => {
-                let location = caps[2]
-                    .parse()
-                    .map_err(|e| format!("Could not parse location: {}", e))?;
-                let value = caps[3]
-                    .parse()
-                    .map_err(|e| format!("Could not parse value: {}", e))?;
-
-                Chip2Instruction::Mem(location, value)
-            }
-            _ => Err("Invalid instruction")?,
-        })
     }
-}
 
-fn parse_program_chip2(input: &str) -> Result<Vec<Chip2Instruction>, String> {
-    input.lines().map(Chip2Instruction::new).collect()
-}
-
-fn run_program_chip2(program: &Vec<Chip2Instruction>) -> Vec<(u64, u64)> {
-    let mut mask = (
-        u64::max_value() & MASK_36_BITS,
-        u64::max_value() & MASK_36_BITS,
-    );
-
-    program
-        .iter()
-        .filter_map(|instruction| match instruction {
-            Chip2Instruction::Mask(and, or) => {
-                mask = (*and, *or);
-                None
-            }
-            Chip2Instruction::Mem(address, value) => Some(
-                mask.1
-                    .iter()
-                    .flat_map(|index| {
-                        let address = u64::try_from(*address).expect("<=64 bit usize required");
-
-                        iter::once((address | 1 << index | mask.0) & MASK_36_BITS).chain(
-                            iter::once((address & !(1 << index) | mask.0) & MASK_36_BITS),
-                        )
-                    })
-                    .map(|address| (address, *value))
-                    .collect::<Vec<(u64, u64)>>(),
-            ),
-        })
-        .flatten()
-        .collect::<Vec<(u64, u64)>>()
-        .into_iter()
-        .rev()
-        .unique_by(|address| address.0)
-        .collect()
+    memory.drain().collect()
 }
 
 #[cfg(test)]
@@ -184,7 +135,7 @@ mod tests {
             "
         );
 
-        let program = parse_program_chip1(input)?;
+        let program = parse_program(input)?;
         let memory = run_program_chip1(&program);
 
         assert_eq!(memory.iter().map(|(_, value)| value).sum::<u64>(), 165);
@@ -202,7 +153,7 @@ mod tests {
             "
         );
 
-        let program = parse_program_chip2(input)?;
+        let program = parse_program(input)?;
         let memory = dbg!(run_program_chip2(&program));
 
         assert_eq!(memory.iter().map(|(_, value)| value).sum::<u64>(), 208);