From 17d41ef6c66e9ca7bedd38bb72446a811931aeaf Mon Sep 17 00:00:00 2001
From: Vinzenz Schroeter <vinzenz.f.s@gmail.com>
Date: Sun, 2 Jun 2024 14:34:22 +0200
Subject: [PATCH] add more iterators

---
 crates/servicepoint/src/bit_vec.rs    |  8 +++--
 crates/servicepoint/src/byte_grid.rs  | 47 ++++++++++++++++++++++-----
 crates/servicepoint/src/grid.rs       |  7 ++++
 crates/servicepoint/src/pixel_grid.rs | 33 +++++++++++++++++++
 4 files changed, 84 insertions(+), 11 deletions(-)

diff --git a/crates/servicepoint/src/bit_vec.rs b/crates/servicepoint/src/bit_vec.rs
index 5af9da8..23719a7 100644
--- a/crates/servicepoint/src/bit_vec.rs
+++ b/crates/servicepoint/src/bit_vec.rs
@@ -106,6 +106,7 @@ impl BitVec {
         Iter {
             bit_vec: self,
             index: 0,
+            end: self.size,
         }
     }
 
@@ -152,15 +153,16 @@ impl From<&[u8]> for BitVec {
 }
 
 pub struct Iter<'t> {
-    bit_vec: &'t BitVec,
-    index: usize,
+    pub(crate) bit_vec: &'t BitVec,
+    pub(crate) index: usize,
+    pub(crate) end: usize,
 }
 
 impl<'t> Iterator for Iter<'t> {
     type Item = bool;
 
     fn next(&mut self) -> Option<Self::Item> {
-        if self.index >= self.bit_vec.size {
+        if self.index >= self.end {
             return None;
         }
 
diff --git a/crates/servicepoint/src/byte_grid.rs b/crates/servicepoint/src/byte_grid.rs
index 7b3f78b..5dbc654 100644
--- a/crates/servicepoint/src/byte_grid.rs
+++ b/crates/servicepoint/src/byte_grid.rs
@@ -43,13 +43,6 @@ impl ByteGrid {
         }
     }
 
-    pub fn iter(&self) -> Iter {
-        Iter {
-            byte_grid: self,
-            index: 0,
-        }
-    }
-
     fn check_indexes(&self, x: usize, y: usize) {
         assert!(
             x < self.width,
@@ -99,6 +92,21 @@ impl Grid<u8> for ByteGrid {
         self.data[x + y * self.width]
     }
 
+    fn iter(&self) -> impl Iterator<Item = u8> {
+        Iter {
+            byte_grid: self,
+            index: 0,
+            end: self.data.len(),
+        }
+    }
+
+    fn iter_rows(&self) -> impl Iterator<Item = impl Iterator<Item = u8>> {
+        IterRows {
+            byte_grid: self,
+            row: 0,
+        }
+    }
+
     fn fill(&mut self, value: u8) {
         self.data.fill(value);
     }
@@ -162,13 +170,14 @@ impl RefGrid<u8> for ByteGrid {
 pub struct Iter<'t> {
     byte_grid: &'t ByteGrid,
     index: usize,
+    end: usize,
 }
 
 impl<'t> Iterator for Iter<'t> {
     type Item = u8;
 
     fn next(&mut self) -> Option<Self::Item> {
-        if self.index >= self.byte_grid.data.len() {
+        if self.index >= self.end {
             return None;
         }
 
@@ -178,6 +187,28 @@ impl<'t> Iterator for Iter<'t> {
     }
 }
 
+pub struct IterRows<'t> {
+    byte_grid: &'t ByteGrid,
+    row: usize,
+}
+
+impl<'t> Iterator for IterRows<'t> {
+    type Item = Iter<'t>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.row >= self.byte_grid.height {
+            return None;
+        }
+        let result = Some(Iter {
+            byte_grid: self.byte_grid,
+            index: self.row * self.byte_grid.width,
+            end: (self.row + 1) * self.byte_grid.width,
+        });
+        self.row += 1;
+        result
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::{ByteGrid, DataRef, Grid};
diff --git a/crates/servicepoint/src/grid.rs b/crates/servicepoint/src/grid.rs
index 76c80db..c9d93e0 100644
--- a/crates/servicepoint/src/grid.rs
+++ b/crates/servicepoint/src/grid.rs
@@ -54,6 +54,13 @@ pub trait Grid<T> {
         }
     }
 
+    /// Get an iterator over every cell in the grid.
+    ///
+    /// Iteration is done in memory order, rows first.
+    fn iter(&self) -> impl Iterator<Item = T>;
+
+    fn iter_rows(&self) -> impl Iterator<Item = impl Iterator<Item = T>>;
+
     /// Sets all cells in the grid to the specified value
     fn fill(&mut self, value: T);
 
diff --git a/crates/servicepoint/src/pixel_grid.rs b/crates/servicepoint/src/pixel_grid.rs
index c6d02dd..3f55a70 100644
--- a/crates/servicepoint/src/pixel_grid.rs
+++ b/crates/servicepoint/src/pixel_grid.rs
@@ -96,6 +96,17 @@ impl Grid<bool> for PixelGrid {
         self.bit_vec.get(x + y * self.width)
     }
 
+    fn iter(&self) -> impl Iterator<Item = bool> {
+        self.bit_vec.iter()
+    }
+
+    fn iter_rows(&self) -> impl Iterator<Item = impl Iterator<Item = bool>> {
+        IterRows {
+            pixel_grid: self,
+            row: 0,
+        }
+    }
+
     /// Sets the state of all pixels in the `PixelGrid`.
     ///
     /// # Arguments
@@ -132,6 +143,28 @@ impl From<PixelGrid> for Vec<u8> {
     }
 }
 
+pub struct IterRows<'t> {
+    pixel_grid: &'t PixelGrid,
+    row: usize,
+}
+
+impl<'t> Iterator for IterRows<'t> {
+    type Item = crate::bit_vec::Iter<'t>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.row >= self.pixel_grid.height {
+            return None;
+        }
+        let result = Some(crate::bit_vec::Iter {
+            bit_vec: &self.pixel_grid.bit_vec,
+            index: self.row * self.pixel_grid.width,
+            end: (self.row + 1) * self.pixel_grid.width,
+        });
+        self.row += 1;
+        result
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::{DataRef, Grid, PixelGrid};