From 849cc760344b6441818fcceb4808e2537b459712 Mon Sep 17 00:00:00 2001
From: Daniel Seiller <earthnuker@gmail.com>
Date: Tue, 15 Dec 2020 00:35:08 +0100
Subject: [PATCH] Day 14

---
 src/day14.rs | 170 +++++++++++++++++++++++++++++++++++++++++++++++++++
 src/lib.rs   |   1 +
 2 files changed, 171 insertions(+)
 create mode 100644 src/day14.rs

diff --git a/src/day14.rs b/src/day14.rs
new file mode 100644
index 0000000..ee97c72
--- /dev/null
+++ b/src/day14.rs
@@ -0,0 +1,170 @@
+use aoc_runner_derive::{aoc, aoc_generator};
+use std::collections::HashMap;
+
+#[derive(Debug, Clone)]
+pub enum Command {
+    Mask { and: u64, or: u64, raw: Vec<char> },
+    Write { val: u64, addr: u64 },
+}
+
+#[inline(always)]
+fn set_bit(v: u64, bit: usize, val: bool) -> u64 {
+    let mut v = v;
+    let bv = 1 << bit;
+    if val {
+        v |= bv;
+    } else {
+        v &= !bv;
+    }
+    return v;
+}
+
+impl Command {
+    fn apply_mask_v1(&self, val: u64) -> u64 {
+        if let Command::Mask { and, or, .. } = self {
+            return (val & and) | or;
+        };
+        panic!("Can't apply non-mask command");
+    }
+
+    fn apply_mask_v2_inner(&self, addr: &mut u64, s: &Vec<char>, ret: &mut Vec<u64>) {
+        let mut s_c = s.clone();
+        for (idx, c) in s.iter().enumerate() {
+            match *c {
+                '0' => (),
+                '1' => {
+                    *addr = set_bit(*addr, idx, true);
+                }
+                'X' => {
+                    s_c[idx] = '0';
+                    *addr = set_bit(*addr, idx, true);
+                    self.apply_mask_v2_inner(addr, &s_c, ret);
+                    *addr = set_bit(*addr, idx, false);
+                    self.apply_mask_v2_inner(addr, &s_c, ret);
+                }
+                _ => unreachable!(),
+            }
+        }
+        ret.push(*addr);
+    }
+
+    fn apply_mask_v2(&self, addr: u64) -> Vec<u64> {
+        let mut addr = addr;
+        if let Command::Mask { raw, .. } = self {
+            let mut raw = raw.clone();
+            raw.reverse();
+            let mut ret = vec![];
+            self.apply_mask_v2_inner(&mut addr, &raw, &mut ret);
+            ret.sort();
+            ret.dedup();
+            return ret;
+        };
+        panic!("Can't apply non-mask command");
+    }
+}
+
+type Data = Vec<Command>;
+
+#[aoc_generator(day14)]
+pub fn input_generator(input: &str) -> Data {
+    let mut ret = vec![];
+    for line in input.lines() {
+        if line.starts_with("mask") {
+            let mut mask_and = 0;
+            let mut mask_or = 0;
+            let mut char_buf = vec![];
+            for (n, c) in line.split("=").nth(1).unwrap().trim().chars().enumerate() {
+                char_buf.push(c);
+                mask_and <<= 1;
+                mask_or <<= 1;
+                match c {
+                    'X' => {
+                        mask_and |= 1;
+                        mask_or |= 0;
+                    }
+                    '1' => {
+                        mask_and |= 1;
+                        mask_or |= 1;
+                    }
+                    '0' => {
+                        mask_and |= 0;
+                        mask_or |= 0;
+                    }
+                    c => {
+                        panic!("Invalid char in mask: {}", c);
+                    }
+                }
+            }
+            ret.push(Command::Mask {
+                and: mask_and,
+                or: mask_or,
+                raw: char_buf,
+            });
+            continue;
+        }
+        if line.starts_with("mem") {
+            let addr: u64 = line
+                .chars()
+                .skip_while(|&c| c != '[')
+                .skip(1)
+                .take_while(|&c| c != ']')
+                .collect::<String>()
+                .parse()
+                .unwrap();
+            let val: u64 = line
+                .chars()
+                .skip_while(|&c| c != '=')
+                .skip(1)
+                .skip_while(|&c| c.is_whitespace())
+                .collect::<String>()
+                .parse()
+                .unwrap();
+            ret.push(Command::Write { addr, val });
+        }
+    }
+    ret
+}
+
+#[aoc(day14, part1)]
+pub fn solve_part1(input: &Data) -> u64 {
+    let mut mem = HashMap::new();
+    let mut mask = Command::Mask {
+        and: std::u64::MAX,
+        or: 0u64,
+        raw: vec![],
+    };
+    for inst in input {
+        match inst {
+            Command::Write { val, addr } => {
+                mem.insert(*addr, mask.apply_mask_v1(*val));
+            }
+            m => {
+                mask = m.clone();
+            }
+        }
+    }
+    return mem.values().copied().sum();
+}
+
+#[aoc(day14, part2)]
+pub fn solve_part2(input: &Data) -> u64 {
+    let mut mem = HashMap::new();
+    let mut mask = Command::Mask {
+        and: std::u64::MAX,
+        or: 0u64,
+        raw: vec![],
+    };
+    for inst in input {
+        match inst {
+            Command::Write { val, addr } => {
+                for addr in mask.apply_mask_v2(*addr) {
+                    mem.insert(addr, val);
+                }
+            }
+            m => {
+                mask = m.clone();
+            }
+        }
+    }
+    return mem.values().copied().sum();
+}
diff --git a/src/lib.rs b/src/lib.rs
index 072a12a..b5ee9f1 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -12,4 +12,5 @@ pub mod day10;
 pub mod day11;
 pub mod day12;
 pub mod day13;
+pub mod day14;
 aoc_lib! { year = 2020 }