Day 3 Part 2

src/bin/day03.rs

@@ -2,41 +2,80 @@ use color_eyre::{
     eyre::{bail, ensure},
     Result,
 };
-use std::io::stdin;
+use std::{cell::Cell, collections::HashMap, io::stdin, rc::Rc};
 
-fn part1(
-    matrix: &[u8],
+#[derive(Clone, Copy)]
+struct SliceCols<'a, T> {
     cols: usize,
-) -> Result<impl Iterator<Item = Result<u64, std::num::ParseIntError>> + '_> {
-    let cols_sig = isize::try_from(cols)?;
-
-    Ok(matrix
-        .chunks(cols)
-        .flat_map(|chunk| chunk.split(|pred| !(*pred as char).is_ascii_digit()))
-        .filter(|s| !s.is_empty())
-        .map(move |slice| {
-            let ix = slice.as_ptr() as usize - matrix.as_ptr() as usize;
-            (ix..slice.len() + ix)
-                .flat_map(|ix| {
-                    (ix % cols != 0)
-                        .then_some([-1, cols_sig - 1, -cols_sig - 1])
-                        .into_iter()
-                        .chain((ix % cols != cols - 1).then_some([1, cols_sig + 1, -cols_sig + 1]))
-                        .flatten()
-                        .chain([-cols_sig, cols_sig])
-                        .flat_map(move |offset| ix.checked_add_signed(offset))
-                })
-                .flat_map(|ix| matrix.get(ix))
-                .any(|&b| (b as char).is_ascii_punctuation() && b != b'.')
-                .then(|| String::from_utf8_lossy(slice).parse::<u64>())
-                .transpose()
-                .map(|x| x.unwrap_or_default())
-        }))
+    data: &'a [T],
 }
 
-fn part2(matrix: &[u8], cols: usize) -> Result<impl Iterator<Item = Result<u64>> + '_> {
+fn separate_numbers(matrix: SliceCols<u8>) -> impl Iterator<Item = &[u8]> {
+    matrix
+        .data
+        .chunks(matrix.cols)
+        .flat_map(|chunk| chunk.split(|&b| !char::from(b).is_ascii_digit()))
+        .filter(|s| !s.is_empty())
+}
 
-    Ok(std::iter::once(Ok(0)))
+fn around_indices(ix: usize, cols: usize) -> impl Iterator<Item = usize> {
+    let cols_sig = isize::try_from(cols).expect("Length can't be larger than isize::MAX");
+    (ix % cols != 0)
+        .then_some([-1, cols_sig - 1, -cols_sig - 1])
+        .into_iter()
+        .chain((ix % cols != cols - 1).then_some([1, cols_sig + 1, -cols_sig + 1]))
+        .flatten()
+        .chain([-cols_sig, cols_sig])
+        .flat_map(move |offset| ix.checked_add_signed(offset))
+}
+
+fn part1(matrix: SliceCols<u8>) -> u64 {
+    separate_numbers(matrix)
+        .map(move |slice| {
+            let ix = slice.as_ptr() as usize - matrix.data.as_ptr() as usize;
+            (ix..slice.len() + ix)
+                .flat_map(|ix| around_indices(ix, matrix.cols))
+                .flat_map(|ix| matrix.data.get(ix))
+                .any(|&b| (b as char).is_ascii_punctuation() && b != b'.')
+                .then(|| {
+                    String::from_utf8_lossy(slice)
+                        .parse::<u64>()
+                        .expect("Filter already isolated numbers")
+                })
+                .unwrap_or_default()
+        })
+        .sum()
+}
+
+fn part2(matrix: SliceCols<u8>) -> u64 {
+    let numbers: HashMap<usize, Rc<Cell<Option<u64>>>> = separate_numbers(matrix)
+        .flat_map(|slice| {
+            let num = Rc::new(Cell::new(Some(
+                String::from_utf8_lossy(slice)
+                    .parse::<u64>()
+                    .expect("Filter already isolated numbers"),
+            )));
+            slice
+                .iter()
+                .map(|i| i as *const u8 as usize - matrix.data.as_ptr() as usize)
+                .zip(std::iter::repeat(num))
+        })
+        .collect();
+
+    matrix
+        .data
+        .iter()
+        .filter(|&&b| b == b'*')
+        .map(|i| i as *const u8 as usize - matrix.data.as_ptr() as usize)
+        .map(|ix| {
+            let mut it = around_indices(ix, matrix.cols).flat_map(|ix| {
+                let i = numbers.get(&ix)?;
+                i.take()
+            });
+
+            it.next().unwrap_or_default() * it.next().unwrap_or_default() * it.product::<u64>()
+        })
+        .sum()
 }
 
 fn main() -> Result<()> {
@@ -60,10 +99,12 @@ fn main() -> Result<()> {
         Ok(())
     })?;
 
-    println!(
-        "{}",
-        part1(&matrix, cols)?.try_fold::<_, _, Result<_>>(0, |acc, x| Ok(acc + x?))?
-    );
+    let matrix = SliceCols {
+        cols,
+        data: &matrix,
+    };
+
+    println!("{}", part2(matrix));
 
     Ok(())
 }