From 8e7024e7d92935359b7beefc161c92696c72cabf Mon Sep 17 00:00:00 2001
From: MauroFab <maurotoscano2@gmail.com>
Date: Thu, 16 Jul 2026 13:03:42 -0300
Subject: [PATCH] refactor(stark): direct scatter in reconstruct_ood_full

The next-row fill loop scanned next_row_cols per output cell (O(width
x mask_width) per row, degrading toward O(width^2) for AIRs using the
conservative all-columns next-row window). Replace it with a zero-fill
followed by a direct scatter of each masked value into its column.

Preserves the documented never-panic contract (bounds-checked reads via
.get, malformed/short archives yield zero-filled cells) and silently
ignores out-of-range indices in next_row_cols. The scatter is
last-write-wins on duplicate indices where the old scan was
first-match-wins; both agree in practice because split_ood_blocks
never emits duplicate indices with differing values, but the two are
not bit-identical for a pathologically malformed next_row_cols.

Adds unit tests for out-of-range next_row_cols indices and a
short/truncated next_block, both exercising the no-panic path.
---
 crypto/stark/src/ood.rs | 52 ++++++++++++++++++++++++++++++-----------
 1 file changed, 38 insertions(+), 14 deletions(-)

diff --git a/crypto/stark/src/ood.rs b/crypto/stark/src/ood.rs
index a779d7851..56bb84c48 100644
--- a/crypto/stark/src/ood.rs
+++ b/crypto/stark/src/ood.rs
@@ -159,21 +159,23 @@ pub fn reconstruct_ood_full<E: IsField>(
         }
     }
 
-    for r in step_size..num_eval_points {
-        let next_row = r - step_size;
-        for c in 0..width {
-            let mut val = FieldElement::<E>::zero();
-            if mask_width > 0 {
-                for (m, &mc) in next_row_cols.iter().enumerate() {
-                    if mc == c {
-                        if let Some(v) = next_block.get(next_row * mask_width + m) {
-                            val = v.clone();
-                        }
-                        break;
-                    }
-                }
+    // Zero-fill the next-row rows, then scatter the surviving masked values
+    // directly into their columns instead of scanning `next_row_cols` per
+    // cell. `.max` keeps the current-row block intact even if
+    // `num_eval_points < step_size` (defensive only: for a well-formed AIR
+    // `num_eval_points` is always a positive multiple of `step_size`).
+    data.resize(
+        data.len().max(num_eval_points * width),
+        FieldElement::<E>::zero(),
+    );
+    for next_row in 0..num_eval_points.saturating_sub(step_size) {
+        let row_base = (step_size + next_row) * width;
+        for (m, &mc) in next_row_cols.iter().enumerate() {
+            if mc < width
+                && let Some(v) = next_block.get(next_row * mask_width + m)
+            {
+                data[row_base + mc] = v.clone();
             }
-            data.push(val);
         }
     }
 
@@ -245,6 +247,28 @@ mod tests {
         assert_eq!(recon.get_row(1), &[Fe::zero(), Fe::zero()]);
     }
 
+    #[test]
+    fn out_of_range_next_row_col_is_ignored_not_panicking() {
+        // width = 3, but next_row_cols advertises column 5 -- out of range.
+        let current_block = vec![fe(1), fe(2), fe(3)];
+        let next_block = vec![fe(99)]; // would-be value for the bogus column
+        let recon = reconstruct_ood_full(&current_block, 3, &next_block, 2, 1, &[5]);
+        assert_eq!(recon.get_row(0), &[fe(1), fe(2), fe(3)]);
+        assert_eq!(recon.get_row(1), &[Fe::zero(), Fe::zero(), Fe::zero()]);
+    }
+
+    #[test]
+    fn short_next_block_leaves_missing_cells_zero_not_panicking() {
+        // width = 3, 3 eval points (step_size 1) => 2 next rows, mask = {0, 2}
+        // so the mask implies 4 next-row values, but next_block only has 1.
+        let current_block = vec![fe(1), fe(2), fe(3)];
+        let next_block = vec![fe(99)];
+        let recon = reconstruct_ood_full(&current_block, 3, &next_block, 3, 1, &[0, 2]);
+        assert_eq!(recon.get_row(0), &[fe(1), fe(2), fe(3)]);
+        assert_eq!(recon.get_row(1), &[fe(99), Fe::zero(), Fe::zero()]); // only present value scattered
+        assert_eq!(recon.get_row(2), &[Fe::zero(), Fe::zero(), Fe::zero()]); // fully missing -> zero
+    }
+
     #[test]
     fn pruned_coeffs_are_zero_off_the_window() {
         // 4 surviving powers for W=3, num_eval_points=2, mask={1}.