Pre-compute initial explode actions

During part two, we are computing a pairing between every two
lines. Any line that had one or more depth 4 number pairs will cause
the same set of explode actions to be taken for each of the other 99
lines it is paired with, before we get into the main effort of add()
in processing split actions.  Rather than repeat early work that will
produce the same answer, we can instead pre-compute these initial
explode calls during parse time.

On my input, adding a simple counter shows that this reduces the
number of calls to explode() from 180k down to 131k. parse() and
part1() runtimes are about the same, while part2() speeds up from
600us to 560us.

Author: ebblake

src/year2021/day18.rs

@@ -24,9 +24,25 @@
 //!
 //! This means that we can store each snailfish number as an implicit data structure in a fixed-size
 //! array. This is faster, smaller and more convenient than using a traditional struct with pointers.
-//! The root node is stored at index 1 (index 0 is unused). For a node at index `i` its left child
-//! is at index `2i`, right child at index `2i + 1` and parent at index `i / 2`. As leaf nodes are
-//! always greater than or equal to zero, `-1` is used as a special sentinel value for non-leaf nodes.
+//! The root node is stored at index 1 (index 0 is unused by the tree, but see below). For a
+//! node at index `i` its left child is at index `2i`, right child at index `2i + 1` and parent
+//! at index `i / 2`. As leaf nodes are always greater than or equal to zero, `-1` is used as a
+//! special sentinel value for non-leaf nodes.
+//!
+//! Another optimization is realizing that all of the explode actions before the first split can
+//! be pre-computed. Instead of parsing a line `[1,[2,[3,[4,5]]]]` as written, we instead parse it
+//! it as if it had been `[[1,[2,[3,[4,5]]]],0]`, then perform the initial explode actions for that
+//! line up front, such that node 3 contains the value that must be added to the first leaf of a
+//! right-hand value in summation. Node 0 is not used by the implicit tree structure, so we instead
+//! use it as a tri-state value:
+//!
+//! - -2 means this Snailfish number is the result of a sum, for the left side during part one.
+//!   When passed to `add()`, we must shuffle contents one level lower.
+//! - -1 means this Snailfish number was just parsed, but did not explode left. When used on the
+//!   right side of add, any spillover from the left is added to the first leaf on the right.
+//! - Non-negative means this Snailfish number was just parsed, and had an explode that spilled
+//!   left. When used as the right side of an add, any spill from the left is combined with
+//!   this value, then added to the last leaf on the left.
 use crate::util::parse::*;
 use crate::util::thread::*;
 
@@ -49,8 +65,12 @@ pub fn parse(input: &str) -> Vec<Snailfish> {
     input
         .lines()
         .map(|line: &str| {
+            // Treat the line as if it had been `[line,0]`, then perform explode until it is back
+            // at depth 4. This allows later add() operations to do less work. Index 0 and 3
+            // then track the amount spilled left or right from those explodes.
             let mut tree = [-1; 64];
-            let mut i = 1;
+            tree[3] = 0;
+            let mut i = 2;
 
             for b in line.bytes() {
                 match b {
@@ -60,6 +80,11 @@ pub fn parse(input: &str) -> Vec<Snailfish> {
                     b => tree[i] = b.to_decimal() as i32,
                 }
             }
+            for pair in (32..48).step_by(2) {
+                if tree[pair] >= 0 {
+                    explode(&mut tree, pair);
+                }
+            }
 
             tree
         })
@@ -100,29 +125,56 @@ fn worker(iter: ParIter<'_, (&Snailfish, &Snailfish)>) -> Option<i32> {
 /// The initial step creates a new root node then makes the numbers the left and right children
 /// of this new root node, by copying the respective ranges of the implicit trees.
 ///
-/// We can optimize the rules a little. This initial combination is the only time that more than one
-/// pair will be 4 levels deep simultaneously, so we can sweep from left to right on all possible
-/// leaf nodes in one pass.
+/// We can optimize the rules a little. The parse step already ensured that there are no pairs
+/// deeper than 4 levels, and precomputed any explode values to spill between the two halves
+/// of the joined value. All that remains is checking for splits, where each split also takes
+/// care of any additional explodes needed.
 fn add(left: &Snailfish, right: &Snailfish) -> Snailfish {
     let mut tree = [-1; 64];
 
-    tree[4..6].copy_from_slice(&left[2..4]);
-    tree[8..12].copy_from_slice(&left[4..8]);
-    tree[16..24].copy_from_slice(&left[8..16]);
-    tree[32..48].copy_from_slice(&left[16..32]);
-
-    tree[6..8].copy_from_slice(&right[2..4]);
-    tree[12..16].copy_from_slice(&right[4..8]);
-    tree[24..32].copy_from_slice(&right[8..16]);
-    tree[48..64].copy_from_slice(&right[16..32]);
+    if left[0] == -2 {
+        // Left comes from a running sum during part one. We need to increase the depth, which
+        // in turn might cause some depth 5 leaves that need explode.
+        tree[3] = 0;
+        tree[4..6].copy_from_slice(&left[2..4]);
+        tree[8..12].copy_from_slice(&left[4..8]);
+        tree[16..24].copy_from_slice(&left[8..16]);
+        tree[32..48].copy_from_slice(&left[16..32]);
+
+        for pair in (32..48).step_by(2) {
+            if tree[pair] >= 0 {
+                explode(&mut tree, pair);
+            }
+        }
+    } else {
+        // We are adding two just-parsed numbers; the left is already rooted at 2 and has no depth 5
+        // leaves, making it ready to copy into place.
+        tree[3..24].copy_from_slice(&left[3..24]);
+    }
 
-    for pair in (32..64).step_by(2) {
-        if tree[pair] >= 0 {
-            explode(&mut tree, pair);
+    // Copy the right into place. This value is always just-parsed, with no depth 5 leaves.
+    tree[6..8].copy_from_slice(&right[4..6]);
+    tree[12..16].copy_from_slice(&right[8..12]);
+    tree[24..32].copy_from_slice(&right[16..24]);
+
+    // Adjust by the explode spillover between sides. We ensured that tree[3] contains any
+    // value to spill right, but must check right[0] to see if that sum then spills back left.
+    let (mut i, spill) = if right[0] == -1 { (24, tree[3]) } else { (23, tree[3] + right[0]) };
+    loop {
+        if tree[i] >= 0 {
+            tree[i] += spill;
+            break;
         }
+        i /= 2;
     }
+    tree[3] = -1;
 
+    // Now we process all split operations; any further explode actions are done during the split
+    // that creates a temporary depth 5.
     while split(&mut tree) {}
+
+    // Mark this tree as a sum before returning it.
+    tree[0] = -2;
     tree
 }
 
@@ -145,6 +197,9 @@ fn explode(tree: &mut Snailfish, pair: usize) {
             }
             i /= 2;
         }
+    } else {
+        // Store the left spill-out for later use by add().
+        tree[0] = tree[pair];
     }
 
     if pair < 62 {