Day 22 part 2

2 files changed, 168 insertions, 63 deletions

diff --git a/Cargo.toml b/Cargo.toml
index 9e2399b..f359858 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -9,3 +9,6 @@ cached = "0.25.0"
 derive_more = "0.99.2"
 nom = "7.0.0"
 thiserror = "1.0.29"
+
+[profile.release]
+# debug = true
diff --git a/src/bin/day_22.rs b/src/bin/day_22.rs
index 5dee0da..85d9ec9 100644
--- a/src/bin/day_22.rs
+++ b/src/bin/day_22.rs
@@ -7,28 +7,29 @@ use nom::{
     sequence::tuple,
     IResult,
 };
-use std::fs;
+use std::{cmp, fs};
 
 fn main() -> Result<(), Box<dyn std::error::Error>> {
     let input = fs::read_to_string("inputs/day_22.txt")?;
     let instructions = parse_instructions(&input).unwrap().1;
     {
-        let mut octtree_50 = OctTree::with_bounds(Block {
+        let bounds_50 = Block {
             min_x: -50,
             max_x: 51,
             min_y: -50,
             max_y: 51,
             min_z: -50,
             max_z: 51,
-        });
+        };
+        let mut octtree_50 = OctTree::default();
         for instruction in &instructions {
-            octtree_50.set_block(&instruction.bounds, instruction.new_state);
+            octtree_50.set_block(&bounds_50, &instruction.bounds, instruction.new_state);
         }
-        dbg!(octtree_50.count_on_blocks());
+        dbg!(octtree_50.count_on_blocks(&bounds_50));
     }
 
     {
-        let mut octtree = OctTree::with_bounds(Block {
+        let problem_boundary = Block {
             min_x: instructions
                 .iter()
                 .map(|i| i.bounds.min_x)
@@ -59,99 +60,179 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
                 .map(|i| i.bounds.max_z)
                 .max()
                 .unwrap_or(0),
-        });
-        for instruction in &instructions {
-            octtree.set_block(&instruction.bounds, instruction.new_state);
         }
-        dbg!(octtree.count_on_blocks());
+        .expand_to_power_cube();
+
+        // This chunking and adding partial solutions is necessary
+        // because I can't fit the whole thing in memory at once :(
+        // It runs really slowly.
+        let mut count = 0;
+        for chunk_index in 0..8 {
+            let problem_boundary = problem_boundary.oct_chunk(chunk_index);
+            for chunk_index in 0..8 {
+                let problem_boundary = problem_boundary.oct_chunk(chunk_index);
+                for chunk_index in 0..8 {
+                    let problem_boundary = problem_boundary.oct_chunk(chunk_index);
+                    let mut octtree = OctTree::default();
+                    for instruction in &instructions {
+                        octtree.set_block(
+                            &problem_boundary,
+                            &instruction.bounds,
+                            instruction.new_state,
+                        );
+                    }
+                    count += dbg!(octtree.count_on_blocks(&problem_boundary));
+                }
+            }
+        }
+        dbg!(count);
     }
 
     Ok(())
 }
 
-#[derive(Clone)]
+#[derive(Default, Clone)]
 struct OctTree {
-    bounds: Block,
     data: OctTreeData,
 }
 
 impl OctTree {
-    fn with_bounds(bounds: Block) -> OctTree {
-        OctTree {
-            bounds,
-            data: OctTreeData::default(),
-        }
-    }
-
-    fn set_block(&mut self, bounds: &Block, new_val: bool) {
-        if bounds.completely_covers(&self.bounds) {
-            self.data = OctTreeData::Same(new_val);
-        } else if bounds.intersects(&self.bounds) {
+    fn set_block(&mut self, self_bounds: &Block, bounds: &Block, new_val: bool) {
+        if bounds.completely_covers(self_bounds) {
+            self.data = new_val.into();
+        } else if bounds.intersects(self_bounds) {
             match &mut self.data {
-                OctTreeData::Same(current_val) => {
-                    if *current_val != new_val {
-                        self.split();
-                        self.set_block(bounds, new_val);
+                OctTreeData::AllOff => {
+                    if new_val {
+                        self.split(self_bounds);
+                        self.set_block(self_bounds, bounds, new_val);
                     }
                 }
-                OctTreeData::Diverse(ref mut subtrees) => {
-                    for sub in subtrees.iter_mut() {
-                        sub.set_block(bounds, new_val);
+                OctTreeData::AllOn => {
+                    if !new_val {
+                        self.split(self_bounds);
+                        self.set_block(self_bounds, bounds, new_val);
                     }
-                    self.try_join();
                 }
-            };
-        }
-    }
+                OctTreeData::BitSet(ref mut bits) => {
+                    let min_x = cmp::max(self_bounds.min_x, bounds.min_x);
+                    let max_x = cmp::min(self_bounds.max_x, bounds.max_x);
+                    let min_x_index = (min_x - self_bounds.min_x) as usize;
+                    let max_x_index = min_x_index + (max_x - min_x) as usize;
+
+                    let min_y = cmp::max(self_bounds.min_y, bounds.min_y);
+                    let max_y = cmp::min(self_bounds.max_y, bounds.max_y);
+                    let min_y_index = (min_y - self_bounds.min_y) as usize;
+                    let max_y_index = min_y_index + (max_y - min_y) as usize;
 
-    fn split(&mut self) {
-        if let OctTreeData::Same(current_value) = &self.data {
-            let subdata = OctTreeData::Same(*current_value);
-            let chunk_labels = [0, 1, 2, 3, 4, 5, 6, 7];
+                    let min_z = cmp::max(self_bounds.min_z, bounds.min_z);
+                    let max_z = cmp::min(self_bounds.max_z, bounds.max_z);
+                    let min_z_index = (min_z - self_bounds.min_z) as usize;
+                    let max_z_index = min_z_index + (max_z - min_z) as usize;
 
-            self.data = OctTreeData::Diverse(Box::new(chunk_labels.map(|label| OctTree {
-                bounds: self.bounds.oct_chunk(label),
-                data: subdata.clone(),
-            })))
+                    for z_index in min_z_index..max_z_index {
+                        let z_bit_index = z_index << 4;
+                        for y_index in min_y_index..max_y_index {
+                            let y_bit_index = y_index << 2;
+                            for x_index in min_x_index..max_x_index {
+                                let x_bit_index = x_index;
+                                let bit_mask = 1u64 << (z_bit_index + y_bit_index + x_bit_index);
+                                if new_val {
+                                    *bits |= bit_mask;
+                                } else {
+                                    *bits &= !bit_mask;
+                                }
+                            }
+                        }
+                    }
+
+                    if *bits == 0 {
+                        self.data = OctTreeData::AllOff;
+                    } else if *bits == !0u64 {
+                        self.data = OctTreeData::AllOn;
+                    }
+                }
+                OctTreeData::Diverse(ref mut subtrees) => {
+                    for (sub_index, sub) in subtrees.iter_mut().enumerate() {
+                        sub.set_block(&self_bounds.oct_chunk(sub_index as u8), bounds, new_val);
+                    }
+                    if subtrees
+                        .iter()
+                        .all(|sub| matches!(sub.data, OctTreeData::AllOn))
+                    {
+                        self.data = OctTreeData::AllOn;
+                    } else if subtrees
+                        .iter()
+                        .all(|sub| matches!(sub.data, OctTreeData::AllOff))
+                    {
+                        self.data = OctTreeData::AllOff;
+                    }
+                }
+            };
         }
     }
 
-    fn try_join(&mut self) {
-        if let OctTreeData::Diverse(subtrees) = &self.data {
-            if subtrees
-                .iter()
-                .all(|sub| matches!(sub.data, OctTreeData::Same(true)))
-            {
-                self.data = OctTreeData::Same(true);
-            } else if subtrees
-                .iter()
-                .all(|sub| matches!(sub.data, OctTreeData::Same(false)))
-            {
-                self.data = OctTreeData::Same(false);
-            }
+    fn split(&mut self, self_bounds: &Block) {
+        assert!(!matches!(self.data, OctTreeData::Diverse(_)));
+        if self_bounds.volume() == 64 {
+            let new_bitset = match self.data {
+                OctTreeData::AllOn => !0u64,
+                OctTreeData::AllOff => 0,
+                _ => panic!("weird split"),
+            };
+            self.data = OctTreeData::BitSet(new_bitset);
+        } else {
+            let template = OctTree {
+                data: self.data.clone(),
+            };
+            self.data = OctTreeData::Diverse(Box::new([
+                template.clone(),
+                template.clone(),
+                template.clone(),
+                template.clone(),
+                template.clone(),
+                template.clone(),
+                template.clone(),
+                template.clone(),
+            ]));
         }
     }
 
-    fn count_on_blocks(&self) -> usize {
+    fn count_on_blocks(&self, self_bounds: &Block) -> usize {
         match &self.data {
-            OctTreeData::Same(false) => 0,
-            OctTreeData::Same(true) => self.bounds.volume(),
-            OctTreeData::Diverse(subtrees) => {
-                subtrees.iter().map(|sub| sub.count_on_blocks()).sum()
-            }
+            OctTreeData::AllOff => 0,
+            OctTreeData::AllOn => self_bounds.volume(),
+            OctTreeData::BitSet(bitset) => bitset.count_ones() as usize,
+            OctTreeData::Diverse(subtrees) => subtrees
+                .iter()
+                .enumerate()
+                .map(|(index, sub)| sub.count_on_blocks(&self_bounds.oct_chunk(index as u8)))
+                .sum(),
         }
     }
 }
 
 #[derive(Clone)]
 enum OctTreeData {
-    Same(bool),
+    AllOff,
+    AllOn,
+    BitSet(u64),
     Diverse(Box<[OctTree; 8]>),
 }
 
 impl Default for OctTreeData {
     fn default() -> OctTreeData {
-        Self::Same(false)
+        Self::AllOff
+    }
+}
+
+impl From<bool> for OctTreeData {
+    fn from(b: bool) -> Self {
+        if b {
+            Self::AllOn
+        } else {
+            Self::AllOff
+        }
     }
 }
 
@@ -220,6 +301,27 @@ impl Block {
             max_z: if lower_z { mid_z } else { self.max_z },
         }
     }
+
+    fn expand_to_power_cube(&self) -> Block {
+        let mag_x = self.max_x - self.min_x;
+        let mag_y = self.max_y - self.min_y;
+        let mag_z = self.max_z - self.min_z;
+        let mag_max = cmp::max(mag_x, cmp::max(mag_y, mag_z));
+        let first_power_of_2 = (0..)
+            .map(|pow| 2_i32.pow(pow))
+            .filter(|pow_size| *pow_size >= mag_max)
+            .next()
+            .unwrap();
+
+        Block {
+            min_x: self.min_x,
+            max_x: self.max_x + first_power_of_2 - mag_x,
+            min_y: self.min_y,
+            max_y: self.max_y + first_power_of_2 - mag_y,
+            min_z: self.min_z,
+            max_z: self.max_z + first_power_of_2 - mag_z,
+        }
+    }
 }
 
 fn parse_instructions(input: &str) -> IResult<&str, Vec<Instruction>> {