Optimized the Grid entries function quite a lot

Author: marcdejonge

2022/src/bin/day23.rs

@@ -3,7 +3,6 @@
 use advent_lib::direction::Direction;
 use advent_lib::geometry::{point2, vector2, Vector};
 use advent_lib::{
-    direction::CardinalDirection,
     grid::{Grid, Location},
     *,
 };
@@ -39,10 +38,10 @@ const DIRECTIONS: [(Direction, [usize; 3]); 4] = [
     (Direction::East, [1, 2, 3]),
 ];
 
-fn find_possible_move(round: usize, has_elfs: [bool; 8]) -> Option<Direction> {
+fn find_possible_move(round: usize, has_elfs: [&Ground; 8]) -> Option<Direction> {
     for ix in 0..4 {
         let (dir, ixs) = DIRECTIONS[(round + ix) % 4];
-        if ixs.iter().all(|&ix| !has_elfs[ix]) {
+        if ixs.iter().all(|&ix| has_elfs[ix] == &Ground::Empty) {
             return Some(dir);
         }
     }
@@ -52,16 +51,16 @@ fn find_possible_move(round: usize, has_elfs: [bool; 8]) -> Option<Direction> {
 fn step(elfs: &mut Elfs, round: usize) -> bool {
     let mut proposals = Vec::<(Location, Location)>::with_capacity(elfs.len());
     let mut proposal_count = FxHashMap::with_capacity_and_hasher(elfs.len(), Default::default());
-    elfs.entries().filter(|&(_, g)| g == &Ground::Elf).for_each(|(loc, _)| {
-        let has_elfs = CardinalDirection::ALL.map(|dir| elfs.get(loc + dir) == Some(&Ground::Elf));
-        if has_elfs.iter().any(|&has_elf| has_elf) {
-            if let Some(dir) = find_possible_move(round, has_elfs) {
+    for (loc, _) in elfs.entries().filter(|&(_, g)| g == &Ground::Elf) {
+        let neighbours = elfs.cardinal_neighbours(loc).unwrap();
+        if neighbours.contains(&&Ground::Elf) {
+            if let Some(dir) = find_possible_move(round, neighbours) {
                 let target = loc + dir;
                 proposals.push((loc, target));
                 *proposal_count.entry(target).or_default() += 1;
             }
         }
-    });
+    }
 
     let mut some_change = false;
     for &(current, target) in proposals.iter() {
@@ -126,4 +125,4 @@ fn calculate_part2(elfs: &Elfs) -> usize {
 
 day_main!(prepare => calculate_part1, calculate_part2);
 day_test!(23, example => 110, 20 ; crate::prepare );
-day_test!(23 => 3762 ; crate::prepare ); // Part 2 is 997 but takes 5 seconds in test mode
+day_test!(23 => 3762, 997 ; crate::prepare );

shared/src/grid.rs

@@ -132,6 +132,37 @@ impl<T> Grid<T> {
         }
     }
 
+    pub fn cardinal_neighbours(&self, location: Location) -> Option<[&T; 8]>
+    where
+        T: Default + Copy,
+    {
+        // Check boundaries, we don't support fetching the neighbours at the edge
+        if location.x() <= 0
+            || location.x() >= self.width()
+            || location.y() <= 0
+            || location.y() >= self.height()
+        {
+            return None;
+        }
+
+        let ix = (location.x() + location.y() * self.width()) as usize;
+        let y_step = self.width() as usize;
+
+        // We've checked the bounds
+        unsafe {
+            Some([
+                self.items.get_unchecked(ix - y_step),
+                self.items.get_unchecked(ix - y_step + 1),
+                self.items.get_unchecked(ix + 1),
+                self.items.get_unchecked(ix + y_step + 1),
+                self.items.get_unchecked(ix + y_step),
+                self.items.get_unchecked(ix + y_step - 1),
+                self.items.get_unchecked(ix - 1),
+                self.items.get_unchecked(ix - y_step - 1),
+            ])
+        }
+    }
+
     pub fn get(&self, location: Location) -> Option<&T> {
         let ix = self.index_from_location(location)?;
         self.items.get(ix)
@@ -172,17 +203,12 @@ impl<T> Grid<T> {
     }
 
     pub fn entries(&self) -> impl Iterator<Item = (Location, &T)> {
-        self.items
-            .iter()
-            .enumerate()
-            .map(|(index, value)| ((self.size, index).into(), value))
+        Indexed::new(self.items.iter(), self.width())
     }
 
     pub fn entries_mut(&mut self) -> impl Iterator<Item = (Location, &mut T)> {
-        self.items
-            .iter_mut()
-            .enumerate()
-            .map(|(index, value)| ((self.size, index).into(), value))
+        let width = self.width();
+        Indexed::new(self.items.iter_mut(), width)
     }
 
     pub fn values(&self) -> impl Iterator<Item = &T> { self.items.iter() }
@@ -658,3 +684,37 @@ mod tests {
         );
     }
 }
+
+#[derive(Clone, Debug)]
+pub struct Indexed<I> {
+    iter: I,
+    x: i32,
+    y: i32,
+    width: i32,
+}
+
+impl<I> Indexed<I> {
+    fn new(iter: I, width: i32) -> Indexed<I> { Indexed { iter, x: 0, y: 0, width } }
+}
+
+impl<I> Iterator for Indexed<I>
+where
+    I: Iterator,
+{
+    type Item = (Location, <I as Iterator>::Item);
+
+    #[inline]
+    fn next(&mut self) -> Option<Self::Item> {
+        let item = self.iter.next()?;
+        let location = point2(self.x, self.y);
+        self.x += 1;
+        if self.x >= self.width {
+            self.x = 0;
+            self.y += 1;
+        }
+        Some((location, item))
+    }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) { self.iter.size_hint() }
+}