Add size_hint to GridIterator

Co-authored-by: Christian Duerr <contact@christianduerr.com>

2 files changed, 47 insertions, 2 deletions

diff --git a/alacritty_terminal/src/grid/mod.rs b/alacritty_terminal/src/grid/mod.rs
index fbd2c79e..4e394d9d 100644
--- a/alacritty_terminal/src/grid/mod.rs
+++ b/alacritty_terminal/src/grid/mod.rs
@@ -599,7 +599,7 @@ impl<'a, T> Iterator for GridIterator<'a, T> {
         }
 
         match self.point {
-            Point { column, .. } if column == self.grid.last_column() => {
+            Point { column, .. } if column >= self.grid.last_column() => {
                 self.point.column = Column(0);
                 self.point.line += 1;
             },
@@ -608,6 +608,24 @@ impl<'a, T> Iterator for GridIterator<'a, T> {
 
         Some(Indexed { cell: &self.grid[self.point], point: self.point })
     }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        if self.point >= self.end {
+            return (0, Some(0));
+        }
+
+        let size = if self.point.line == self.end.line {
+            (self.end.column - self.point.column).0
+        } else {
+            let cols_on_first_line = self.grid.columns.saturating_sub(self.point.column.0 + 1);
+            let middle_lines = (self.end.line - self.point.line).0 as usize - 1;
+            let cols_on_last_line = self.end.column + 1;
+
+            cols_on_first_line + middle_lines * self.grid.columns + cols_on_last_line.0
+        };
+
+        (size, Some(size))
+    }
 }
 
 /// Bidirectional iterator.
@@ -621,7 +639,7 @@ impl<T> BidirectionalIterator for GridIterator<'_, T> {
         let last_column = self.grid.last_column();
 
         // Stop once we've reached the end of the grid.
-        if self.point == Point::new(topmost_line, Column(0)) {
+        if self.point <= Point::new(topmost_line, Column(0)) {
             return None;
         }
 
diff --git a/alacritty_terminal/src/grid/tests.rs b/alacritty_terminal/src/grid/tests.rs
index 49cf1a76..f052c75a 100644
--- a/alacritty_terminal/src/grid/tests.rs
+++ b/alacritty_terminal/src/grid/tests.rs
@@ -348,6 +348,33 @@ fn shrink_reflow_disabled() {
     assert_eq!(grid[Line(0)][Column(1)], cell('2'));
 }
 
+#[test]
+fn accurate_size_hint() {
+    let grid = Grid::<Cell>::new(5, 5, 2);
+
+    size_hint_matches_count(grid.iter_from(Point::new(Line(0), Column(0))));
+    size_hint_matches_count(grid.iter_from(Point::new(Line(2), Column(3))));
+    size_hint_matches_count(grid.iter_from(Point::new(Line(4), Column(4))));
+    size_hint_matches_count(grid.iter_from(Point::new(Line(4), Column(2))));
+    size_hint_matches_count(grid.iter_from(Point::new(Line(10), Column(10))));
+    size_hint_matches_count(grid.iter_from(Point::new(Line(2), Column(10))));
+
+    let mut iterator = grid.iter_from(Point::new(Line(3), Column(1)));
+    iterator.next();
+    iterator.next();
+    size_hint_matches_count(iterator);
+
+    size_hint_matches_count(grid.display_iter());
+}
+
+fn size_hint_matches_count<T>(iter: impl Iterator<Item = T>) {
+    let iterator = iter.into_iter();
+    let (lower, upper) = iterator.size_hint();
+    let count = iterator.count();
+    assert_eq!(lower, count);
+    assert_eq!(upper, Some(count));
+}
+
 // https://github.com/rust-lang/rust-clippy/pull/6375
 #[allow(clippy::all)]
 fn cell(c: char) -> Cell {