Complete day 14

day-14/src/main.rs

@@ -0,0 +1,162 @@
+use std::collections::HashMap;
+use std::fs;
+
+use itertools::Itertools;
+use lazy_static::lazy_static;
+use regex::Regex;
+
+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(&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 Instruction {
+    Mask(u64, u64, u64),
+    Mem(u64, u64),
+}
+
+impl Instruction {
+    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");
+        }
+
+        let caps = INSTRUCTION_RE
+            .captures(instruction)
+            .ok_or("Instruction did not match")?;
+
+        Ok(match &caps[1] {
+            "mask" => {
+                let and = u64::from_str_radix(&caps[3].replace("X", "1"), 2)
+                    .map_err(|e| format!("Could not parse and: {}", e))?;
+                let or = u64::from_str_radix(&caps[3].replace("X", "0"), 2)
+                    .map_err(|e| format!("Could not parse or: {}", e))?;
+
+                Instruction::Mask(and, or, and ^ or)
+            }
+            "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))?;
+
+                Instruction::Mem(location, value)
+            }
+            _ => Err("Invalid instruction")?,
+        })
+    }
+}
+
+fn parse_program(input: &str) -> Result<Vec<Instruction>, String> {
+    input.lines().map(Instruction::new).collect()
+}
+
+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 {
+            Instruction::Mask(and, or, _) => {
+                mask = (*and, *or);
+                None
+            }
+            Instruction::Mem(address, value) => Some((*address as u64, value & mask.0 | mask.1)),
+        })
+        .collect::<Vec<(u64, u64)>>()
+        .into_iter()
+        .rev()
+        .unique_by(|address| address.0)
+        .collect()
+}
+
+fn run_program_chip2(program: &Vec<Instruction>) -> Vec<(u64, u64)> {
+    let mut memory = HashMap::new();
+    let mut mask = (0, 0);
+
+    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;
+                }
+            }
+        }
+    }
+
+    memory.drain().collect()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use indoc::indoc;
+
+    #[test]
+    fn test_simple() -> Result<(), Box<dyn std::error::Error>> {
+        let input = indoc!(
+            "
+             mask = XXXXXXXXXXXXXXXXXXXXXXXXXXXXX1XXXX0X
+             mem[8] = 11
+             mem[7] = 101
+             mem[8] = 0
+            "
+        );
+
+        let program = parse_program(input)?;
+        let memory = run_program_chip1(&program);
+
+        assert_eq!(memory.iter().map(|(_, value)| value).sum::<u64>(), 165);
+        Ok(())
+    }
+
+    #[test]
+    fn test_simple2() -> Result<(), Box<dyn std::error::Error>> {
+        let input = indoc!(
+            "
+             mask = 000000000000000000000000000000X1001X
+             mem[42] = 100
+             mask = 00000000000000000000000000000000X0XX
+             mem[26] = 1
+            "
+        );
+
+        let program = parse_program(input)?;
+        let memory = dbg!(run_program_chip2(&program));
+
+        assert_eq!(memory.iter().map(|(_, value)| value).sum::<u64>(), 208);
+        Ok(())
+    }
+}