Optimize 2024 day 4

Rewriting to allow vectorization reduces runtime from ~180µs to 60µs, or ~40µs with target_cpu=native.

Author: ictrobot

crates/year2024/src/day04.rs

@@ -4,80 +4,60 @@ use utils::prelude::*;
 /// Counting matches in a word search.
 #[derive(Clone, Debug)]
 pub struct Day04 {
-    rows: usize,
     cols: usize,
     grid: Vec<u8>,
 }
 
 impl Day04 {
     pub fn new(input: &str, _: InputType) -> Result<Self, InputError> {
-        let (rows, cols, grid) = grid::from_str_padded(input, 3, b'\0', |c| match c {
+        let (_, cols, grid) = grid::from_str_padded(input, 3, b'\0', |c| match c {
             b'X' | b'M' | b'A' | b'S' => Some(c),
             _ => None,
         })?;
-        Ok(Self { rows, cols, grid })
+        Ok(Self { cols, grid })
     }
 
     #[must_use]
     pub fn part1(&self) -> u32 {
-        let offsets = [
-            self.cols as isize,
-            -(self.cols as isize),
-            1,
-            -1,
-            (self.cols as isize) + 1,
-            (self.cols as isize) - 1,
-            -(self.cols as isize) + 1,
-            -(self.cols as isize) - 1,
-        ];
+        self.check_offset(self.cols as isize)
+            + self.check_offset(-(self.cols as isize))
+            + self.check_offset(1)
+            + self.check_offset(-1)
+            + self.check_offset((self.cols as isize) + 1)
+            + self.check_offset((self.cols as isize) - 1)
+            + self.check_offset(-(self.cols as isize) + 1)
+            + self.check_offset(-(self.cols as isize) - 1)
+    }
 
+    fn check_offset(&self, offset: isize) -> u32 {
+        let start = 3 * self.cols + 3;
         let mut count = 0;
-        for r in 3..self.rows - 3 {
-            for c in 3..self.cols - 3 {
-                let i = r * self.cols + c;
-                if self.grid[i] != b'X' {
-                    continue;
-                }
-
-                for o in offsets {
-                    if self.grid[i.wrapping_add_signed(o)] == b'M'
-                        && self.grid[i.wrapping_add_signed(o * 2)] == b'A'
-                        && self.grid[i.wrapping_add_signed(o * 3)] == b'S'
-                    {
-                        count += 1;
-                    }
-                }
-            }
+        for (((first, second), third), fourth) in self.grid[start..]
+            .iter()
+            .zip(&self.grid[start.wrapping_add_signed(offset)..])
+            .zip(&self.grid[start.wrapping_add_signed(offset * 2)..])
+            .zip(&self.grid[start.wrapping_add_signed(offset * 3)..])
+        {
+            count += u32::from(
+                (*first == b'X') & (*second == b'M') & (*third == b'A') & (*fourth == b'S'),
+            );
         }
         count
     }
 
     #[must_use]
     pub fn part2(&self) -> u32 {
         let mut count = 0;
-        for r in 4..self.rows - 4 {
-            for c in 4..self.cols - 4 {
-                let i = r * self.cols + c;
-                if self.grid[i] != b'A' {
-                    continue;
-                }
-
-                let (nw, ne, sw, se) = (
-                    self.grid[i.wrapping_add_signed(-(self.cols as isize) - 1)],
-                    self.grid[i.wrapping_add_signed(-(self.cols as isize) + 1)],
-                    self.grid[i.wrapping_add_signed((self.cols as isize) - 1)],
-                    self.grid[i.wrapping_add_signed((self.cols as isize) + 1)],
-                );
-
-                // Given each variable is one of (b'\0', b'X', b'M', b'A', b'S') this is
-                // equivalent to and slightly faster than
-                //  ((nw == b'M' && se == b'S') || (nw == b'S' && se == b'M'))
-                //      && ((ne == b'M' && sw == b'S') || (ne == b'S' && sw == b'M'))
-                // As no other pair XORed equals b'M' ^ b'S'
-                if (nw ^ se) == (b'M' ^ b'S') && (ne ^ sw) == (b'M' ^ b'S') {
-                    count += 1;
-                }
-            }
+        for ((((middle, nw), ne), sw), se) in self.grid[self.cols * 4 + 4..]
+            .iter()
+            .zip(&self.grid[self.cols * 3 + 3..])
+            .zip(&self.grid[self.cols * 3 + 5..])
+            .zip(&self.grid[self.cols * 5 + 3..])
+            .zip(&self.grid[self.cols * 5 + 5..])
+        {
+            count += u32::from(
+                (*middle == b'A') & ((*nw ^ *se) == (b'M' ^ b'S')) & ((*ne ^ *sw) == (b'M' ^ b'S')),
+            );
         }
         count
     }