perf: implement size hints for Rect iterators (#1420)

Co-authored-by: Josh McKinney <joshka@users.noreply.github.com>

benches/main/rect.rs

@@ -1,24 +1,112 @@
-use criterion::{black_box, criterion_group, BenchmarkId, Criterion};
+use criterion::{black_box, criterion_group, BatchSize, Bencher, BenchmarkId, Criterion};
 use ratatui::layout::Rect;
 
-fn rect_rows_benchmark(c: &mut Criterion) {
-    let rect_sizes = vec![
-        Rect::new(0, 0, 1, 16),
-        Rect::new(0, 0, 1, 1024),
-        Rect::new(0, 0, 1, 65535),
-    ];
-    let mut group = c.benchmark_group("rect_rows");
-    for rect in rect_sizes {
-        group.bench_with_input(BenchmarkId::new("rows", rect.height), &rect, |b, rect| {
-            b.iter(|| {
-                for row in rect.rows() {
-                    // Perform any necessary operations on each row
-                    black_box(row);
-                }
-            });
-        });
+fn rect_iters_benchmark(c: &mut Criterion) {
+    let rect_sizes = vec![[16, 16], [64, 64], [255, 255]];
+    let mut group = c.benchmark_group("rect");
+    for rect in rect_sizes.into_iter().map(|[width, height]| Rect {
+        width,
+        height,
+        ..Default::default()
+    }) {
+        group.bench_with_input(
+            BenchmarkId::new("rect_rows_iter", rect.height),
+            &rect,
+            |b, rect| rect_rows_iter(b, *rect),
+        );
+        group.bench_with_input(
+            BenchmarkId::new("rect_rows_collect", rect.height),
+            &rect,
+            |b, rect| rect_rows_collect(b, *rect),
+        );
+        group.bench_with_input(
+            BenchmarkId::new("rect_columns_iter", rect.width),
+            &rect,
+            |b, rect| rect_columns_iter(b, *rect),
+        );
+        group.bench_with_input(
+            BenchmarkId::new("rect_columns_collect", rect.width),
+            &rect,
+            |b, rect| rect_columns_collect(b, *rect),
+        );
+        group.bench_with_input(
+            BenchmarkId::new(
+                "rect_positions_iter",
+                format!("{}x{}", rect.width, rect.height),
+            ),
+            &rect,
+            |b, rect| rect_positions_iter(b, *rect),
+        );
+        group.bench_with_input(
+            BenchmarkId::new(
+                "rect_positions_collect",
+                format!("{}x{}", rect.width, rect.height),
+            ),
+            &rect,
+            |b, rect| rect_positions_collect(b, *rect),
+        );
     }
     group.finish();
 }
 
-criterion_group!(benches, rect_rows_benchmark);
+fn rect_rows_iter(c: &mut Bencher, rect: Rect) {
+    c.iter_batched(
+        || black_box(rect),
+        |rect| {
+            for row in black_box(rect.rows()) {
+                black_box(row);
+            }
+        },
+        BatchSize::LargeInput,
+    );
+}
+
+fn rect_rows_collect(c: &mut Bencher, rect: Rect) {
+    c.iter_batched(
+        || black_box(rect),
+        |rect| black_box(rect.rows()).collect::<Vec<_>>(),
+        BatchSize::LargeInput,
+    );
+}
+
+fn rect_columns_iter(c: &mut Bencher, rect: Rect) {
+    c.iter_batched(
+        || black_box(rect),
+        |rect| {
+            for col in black_box(rect.columns()) {
+                black_box(col);
+            }
+        },
+        BatchSize::LargeInput,
+    );
+}
+
+fn rect_columns_collect(c: &mut Bencher, rect: Rect) {
+    c.iter_batched(
+        || black_box(rect),
+        |rect| black_box(rect.columns()).collect::<Vec<_>>(),
+        BatchSize::LargeInput,
+    );
+}
+
+fn rect_positions_iter(c: &mut Bencher, rect: Rect) {
+    c.iter_batched(
+        || black_box(rect),
+        |rect| {
+            for pos in black_box(rect.positions()) {
+                black_box(pos);
+            }
+        },
+        BatchSize::LargeInput,
+    );
+}
+
+fn rect_positions_collect(b: &mut Bencher, rect: Rect) {
+    b.iter_batched(
+        || black_box(rect),
+        |rect| black_box(rect.positions()).collect::<Vec<_>>(),
+        BatchSize::LargeInput,
+    );
+}
+
+criterion_group!(benches, rect_iters_benchmark);

src/layout/rect/iter.rs

@@ -35,6 +35,17 @@ impl Iterator for Rows {
         self.current_row_fwd += 1;
         Some(row)
     }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let start_count = self.current_row_fwd.saturating_sub(self.rect.top());
+        let end_count = self.rect.bottom().saturating_sub(self.current_row_back);
+        let count = self
+            .rect
+            .height
+            .saturating_sub(start_count)
+            .saturating_sub(end_count) as usize;
+        (count, Some(count))
+    }
 }
 
 impl DoubleEndedIterator for Rows {
@@ -86,6 +97,17 @@ impl Iterator for Columns {
         self.current_column_fwd += 1;
         Some(column)
     }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let start_count = self.current_column_fwd.saturating_sub(self.rect.left());
+        let end_count = self.rect.right().saturating_sub(self.current_column_back);
+        let count = self
+            .rect
+            .width
+            .saturating_sub(start_count)
+            .saturating_sub(end_count) as usize;
+        (count, Some(count))
+    }
 }
 
 impl DoubleEndedIterator for Columns {
@@ -140,6 +162,19 @@ impl Iterator for Positions {
         }
         Some(position)
     }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let row_count = self.rect.bottom().saturating_sub(self.current_position.y);
+        if row_count == 0 {
+            return (0, Some(0));
+        }
+        let column_count = self.rect.right().saturating_sub(self.current_position.x);
+        // subtract 1 from the row count to account for the current row
+        let count = (row_count - 1)
+            .saturating_mul(self.rect.width)
+            .saturating_add(column_count) as usize;
+        (count, Some(count))
+    }
 }
 
 #[cfg(test)]
@@ -150,68 +185,106 @@ mod tests {
     fn rows() {
         let rect = Rect::new(0, 0, 2, 3);
         let mut rows = Rows::new(rect);
+        assert_eq!(rows.size_hint(), (3, Some(3)));
         assert_eq!(rows.next(), Some(Rect::new(0, 0, 2, 1)));
+        assert_eq!(rows.size_hint(), (2, Some(2)));
         assert_eq!(rows.next(), Some(Rect::new(0, 1, 2, 1)));
+        assert_eq!(rows.size_hint(), (1, Some(1)));
         assert_eq!(rows.next(), Some(Rect::new(0, 2, 2, 1)));
+        assert_eq!(rows.size_hint(), (0, Some(0)));
         assert_eq!(rows.next(), None);
+        assert_eq!(rows.size_hint(), (0, Some(0)));
         assert_eq!(rows.next_back(), None);
+        assert_eq!(rows.size_hint(), (0, Some(0)));
     }
 
     #[test]
     fn rows_back() {
         let rect = Rect::new(0, 0, 2, 3);
         let mut rows = Rows::new(rect);
+        assert_eq!(rows.size_hint(), (3, Some(3)));
         assert_eq!(rows.next_back(), Some(Rect::new(0, 2, 2, 1)));
+        assert_eq!(rows.size_hint(), (2, Some(2)));
         assert_eq!(rows.next_back(), Some(Rect::new(0, 1, 2, 1)));
+        assert_eq!(rows.size_hint(), (1, Some(1)));
         assert_eq!(rows.next_back(), Some(Rect::new(0, 0, 2, 1)));
+        assert_eq!(rows.size_hint(), (0, Some(0)));
         assert_eq!(rows.next_back(), None);
+        assert_eq!(rows.size_hint(), (0, Some(0)));
         assert_eq!(rows.next(), None);
+        assert_eq!(rows.size_hint(), (0, Some(0)));
     }
 
     #[test]
     fn rows_meet_in_the_middle() {
         let rect = Rect::new(0, 0, 2, 4);
         let mut rows = Rows::new(rect);
+        assert_eq!(rows.size_hint(), (4, Some(4)));
         assert_eq!(rows.next(), Some(Rect::new(0, 0, 2, 1)));
+        assert_eq!(rows.size_hint(), (3, Some(3)));
         assert_eq!(rows.next_back(), Some(Rect::new(0, 3, 2, 1)));
+        assert_eq!(rows.size_hint(), (2, Some(2)));
         assert_eq!(rows.next(), Some(Rect::new(0, 1, 2, 1)));
+        assert_eq!(rows.size_hint(), (1, Some(1)));
         assert_eq!(rows.next_back(), Some(Rect::new(0, 2, 2, 1)));
+        assert_eq!(rows.size_hint(), (0, Some(0)));
         assert_eq!(rows.next(), None);
+        assert_eq!(rows.size_hint(), (0, Some(0)));
         assert_eq!(rows.next_back(), None);
+        assert_eq!(rows.size_hint(), (0, Some(0)));
     }
 
     #[test]
     fn columns() {
         let rect = Rect::new(0, 0, 3, 2);
         let mut columns = Columns::new(rect);
+        assert_eq!(columns.size_hint(), (3, Some(3)));
         assert_eq!(columns.next(), Some(Rect::new(0, 0, 1, 2)));
+        assert_eq!(columns.size_hint(), (2, Some(2)));
         assert_eq!(columns.next(), Some(Rect::new(1, 0, 1, 2)));
+        assert_eq!(columns.size_hint(), (1, Some(1)));
         assert_eq!(columns.next(), Some(Rect::new(2, 0, 1, 2)));
+        assert_eq!(columns.size_hint(), (0, Some(0)));
         assert_eq!(columns.next(), None);
+        assert_eq!(columns.size_hint(), (0, Some(0)));
         assert_eq!(columns.next_back(), None);
+        assert_eq!(columns.size_hint(), (0, Some(0)));
     }
 
     #[test]
     fn columns_back() {
         let rect = Rect::new(0, 0, 3, 2);
         let mut columns = Columns::new(rect);
+        assert_eq!(columns.size_hint(), (3, Some(3)));
         assert_eq!(columns.next_back(), Some(Rect::new(2, 0, 1, 2)));
+        assert_eq!(columns.size_hint(), (2, Some(2)));
         assert_eq!(columns.next_back(), Some(Rect::new(1, 0, 1, 2)));
+        assert_eq!(columns.size_hint(), (1, Some(1)));
         assert_eq!(columns.next_back(), Some(Rect::new(0, 0, 1, 2)));
+        assert_eq!(columns.size_hint(), (0, Some(0)));
         assert_eq!(columns.next_back(), None);
+        assert_eq!(columns.size_hint(), (0, Some(0)));
         assert_eq!(columns.next(), None);
+        assert_eq!(columns.size_hint(), (0, Some(0)));
     }
 
     #[test]
     fn columns_meet_in_the_middle() {
         let rect = Rect::new(0, 0, 4, 2);
         let mut columns = Columns::new(rect);
+        assert_eq!(columns.size_hint(), (4, Some(4)));
         assert_eq!(columns.next(), Some(Rect::new(0, 0, 1, 2)));
+        assert_eq!(columns.size_hint(), (3, Some(3)));
         assert_eq!(columns.next_back(), Some(Rect::new(3, 0, 1, 2)));
+        assert_eq!(columns.size_hint(), (2, Some(2)));
         assert_eq!(columns.next(), Some(Rect::new(1, 0, 1, 2)));
+        assert_eq!(columns.size_hint(), (1, Some(1)));
         assert_eq!(columns.next_back(), Some(Rect::new(2, 0, 1, 2)));
+        assert_eq!(columns.size_hint(), (0, Some(0)));
         assert_eq!(columns.next(), None);
+        assert_eq!(columns.size_hint(), (0, Some(0)));
         assert_eq!(columns.next_back(), None);
+        assert_eq!(columns.size_hint(), (0, Some(0)));
     }
 
     /// We allow a total of `65536` columns in the range `(0..=65535)`.  In this test we iterate
@@ -241,10 +314,16 @@ mod tests {
     fn positions() {
         let rect = Rect::new(0, 0, 2, 2);
         let mut positions = Positions::new(rect);
+        assert_eq!(positions.size_hint(), (4, Some(4)));
         assert_eq!(positions.next(), Some(Position::new(0, 0)));
+        assert_eq!(positions.size_hint(), (3, Some(3)));
         assert_eq!(positions.next(), Some(Position::new(1, 0)));
+        assert_eq!(positions.size_hint(), (2, Some(2)));
         assert_eq!(positions.next(), Some(Position::new(0, 1)));
+        assert_eq!(positions.size_hint(), (1, Some(1)));
         assert_eq!(positions.next(), Some(Position::new(1, 1)));
+        assert_eq!(positions.size_hint(), (0, Some(0)));
         assert_eq!(positions.next(), None);
+        assert_eq!(positions.size_hint(), (0, Some(0)));
     }
 }