From 8725fc174a5e524f5c55527ccd1ec30f6d26d4a5 Mon Sep 17 00:00:00 2001
From: ginavalent <ginavalent@outlook.com>
Date: Sun, 13 Sep 2026 17:41:11 +0800
Subject: [PATCH] fix exact-size iterator contracts for rows and columns

Signed-off-by: ginavalent <ginavalent@outlook.com>
---
 air/src/proof/table.rs          | 40 +++++++++++++++++++++++++++++----
 prover/src/matrix/col_matrix.rs | 20 ++++++++++-------
 prover/src/matrix/tests.rs      | 37 ++++++++++++++++++++++++++++++
 3 files changed, 85 insertions(+), 12 deletions(-)

diff --git a/air/src/proof/table.rs b/air/src/proof/table.rs
index 9805cc031..879fc3a77 100644
--- a/air/src/proof/table.rs
+++ b/air/src/proof/table.rs
@@ -136,12 +136,44 @@ impl<'a, E: FieldElement> Iterator for RowIterator<'a, E> {
             },
         }
     }
-}
 
-impl<E: FieldElement> ExactSizeIterator for RowIterator<'_, E> {
-    fn len(&self) -> usize {
-        self.table.num_rows()
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let remaining = self.table.num_rows() - self.cursor;
+        (remaining, Some(remaining))
     }
 }
 
+impl<E: FieldElement> ExactSizeIterator for RowIterator<'_, E> {}
+
 impl<E: FieldElement> FusedIterator for RowIterator<'_, E> {}
+
+#[cfg(test)]
+mod tests {
+    use math::fields::f64::BaseElement;
+    use utils::Serializable;
+
+    use super::Table;
+
+    #[test]
+    fn row_iterator_tracks_remaining_length() {
+        let values = vec![
+            BaseElement::new(1),
+            BaseElement::new(2),
+            BaseElement::new(3),
+            BaseElement::new(4),
+        ];
+        let table: Table<BaseElement> = Table::from_bytes(&values.to_bytes(), 4, 1).unwrap();
+        let mut iter = table.rows();
+
+        assert_eq!(iter.len(), 4);
+        assert_eq!(iter.size_hint(), (4, Some(4)));
+
+        iter.next();
+        assert_eq!(iter.len(), 3);
+        assert_eq!(iter.size_hint(), (3, Some(3)));
+
+        iter.by_ref().for_each(drop);
+        assert_eq!(iter.len(), 0);
+        assert_eq!(iter.size_hint(), (0, Some(0)));
+    }
+}
diff --git a/prover/src/matrix/col_matrix.rs b/prover/src/matrix/col_matrix.rs
index 6b15f8088..9ab4dab5f 100644
--- a/prover/src/matrix/col_matrix.rs
+++ b/prover/src/matrix/col_matrix.rs
@@ -331,14 +331,17 @@ impl<'a, E: FieldElement> Iterator for ColumnIter<'a, E> {
             None => None,
         }
     }
-}
 
-impl<E: FieldElement> ExactSizeIterator for ColumnIter<'_, E> {
-    fn len(&self) -> usize {
-        self.matrix.map(|matrix| matrix.num_cols()).unwrap_or_default()
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let remaining =
+            self.matrix.map(|matrix| matrix.num_cols() - self.cursor).unwrap_or_default();
+
+        (remaining, Some(remaining))
     }
 }
 
+impl<E: FieldElement> ExactSizeIterator for ColumnIter<'_, E> {}
+
 impl<E: FieldElement> FusedIterator for ColumnIter<'_, E> {}
 
 impl<E: FieldElement> Default for ColumnIter<'_, E> {
@@ -380,12 +383,13 @@ impl<'a, E: FieldElement> Iterator for ColumnIterMut<'a, E> {
             },
         }
     }
-}
 
-impl<E: FieldElement> ExactSizeIterator for ColumnIterMut<'_, E> {
-    fn len(&self) -> usize {
-        self.matrix.num_cols()
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let remaining = self.matrix.num_cols() - self.cursor;
+        (remaining, Some(remaining))
     }
 }
 
+impl<E: FieldElement> ExactSizeIterator for ColumnIterMut<'_, E> {}
+
 impl<E: FieldElement> FusedIterator for ColumnIterMut<'_, E> {}
diff --git a/prover/src/matrix/tests.rs b/prover/src/matrix/tests.rs
index d37911b33..95a60d5cf 100644
--- a/prover/src/matrix/tests.rs
+++ b/prover/src/matrix/tests.rs
@@ -41,6 +41,43 @@ fn test_eval_poly_with_offset_matrix() {
     }
 }
 
+#[test]
+fn column_iterators_track_remaining_length() {
+    let columns = vec![
+        vec![BaseElement::new(1), BaseElement::new(2)],
+        vec![BaseElement::new(3), BaseElement::new(4)],
+        vec![BaseElement::new(5), BaseElement::new(6)],
+        vec![BaseElement::new(7), BaseElement::new(8)],
+    ];
+    let matrix = ColMatrix::new(columns.clone());
+    let mut iter = matrix.columns();
+
+    assert_eq!(iter.len(), 4);
+    assert_eq!(iter.size_hint(), (4, Some(4)));
+
+    iter.next();
+    assert_eq!(iter.len(), 3);
+    assert_eq!(iter.size_hint(), (3, Some(3)));
+
+    iter.by_ref().for_each(drop);
+    assert_eq!(iter.len(), 0);
+    assert_eq!(iter.size_hint(), (0, Some(0)));
+
+    let mut matrix = ColMatrix::new(columns);
+    let mut iter = matrix.columns_mut();
+
+    assert_eq!(iter.len(), 4);
+    assert_eq!(iter.size_hint(), (4, Some(4)));
+
+    iter.next();
+    assert_eq!(iter.len(), 3);
+    assert_eq!(iter.size_hint(), (3, Some(3)));
+
+    iter.by_ref().for_each(drop);
+    assert_eq!(iter.len(), 0);
+    assert_eq!(iter.size_hint(), (0, Some(0)));
+}
+
 // HELPER FUNCTIONS
 // ================================================================================================