update witness util

plonky2/src/fri/witness_util.rs

@@ -1,9 +1,10 @@
-use anyhow::Result;
+use anyhow::{anyhow, Result};
 use itertools::Itertools;
+use plonky2_field::types::Field;
 
 use crate::field::extension::Extendable;
 use crate::fri::proof::{FriProof, FriProofTarget};
-use crate::hash::hash_types::RichField;
+use crate::hash::hash_types::{HashOut, RichField};
 use crate::iop::witness::WitnessWrite;
 use crate::plonk::config::AlgebraicHasher;
 
@@ -20,13 +21,26 @@ where
 {
     witness.set_target(fri_proof_target.pow_witness, fri_proof.pow_witness)?;
 
-    for (&t, &x) in fri_proof_target
-        .final_poly
-        .0
-        .iter()
-        .zip_eq(&fri_proof.final_poly.coeffs)
-    {
-        witness.set_extension_target(t, x)?;
+    let target_len = fri_proof_target.final_poly.0.len();
+    let coeffs_len = fri_proof.final_poly.coeffs.len();
+
+    if target_len < coeffs_len {
+        return Err(anyhow!(
+            "fri_proof->final_poly's target length is less than the proof length"
+        ));
+    }
+
+    // Set overlapping elements
+    for i in 0..coeffs_len {
+        witness.set_extension_target(
+            fri_proof_target.final_poly.0[i],
+            fri_proof.final_poly.coeffs[i],
+        )?;
+    }
+
+    // Set remaining elements in target to ZERO if target is longer
+    for i in coeffs_len..target_len {
+        witness.set_extension_target(fri_proof_target.final_poly.0[i], F::Extension::ZERO)?;
     }
 
     for (t, x) in fri_proof_target
@@ -51,22 +65,44 @@ where
             for (&t, &x) in at.0.iter().zip_eq(&a.0) {
                 witness.set_target(t, x)?;
             }
-            for (&t, &x) in at.1.siblings.iter().zip_eq(&a.1.siblings) {
-                witness.set_hash_target(t, x)?;
+            let target_len = at.1.siblings.len();
+            let siblings_len = a.1.siblings.len();
+
+            if target_len < siblings_len {
+                return Err(anyhow!("fri_proof->query_round_proofs->initial_trees_proof->evals_proofs->siblings' target length is less than the proof length"));
+            }
+
+            // Set overlapping elements
+            for i in 0..siblings_len {
+                witness.set_hash_target(at.1.siblings[i], a.1.siblings[i])?;
+            }
+
+            // Set remaining elements in target to ZERO if target is longer
+            for i in siblings_len..target_len {
+                witness.set_hash_target(at.1.siblings[i], HashOut::ZERO)?;
             }
         }
 
         for (st, s) in qt.steps.iter().zip_eq(&q.steps) {
             for (&t, &x) in st.evals.iter().zip_eq(&s.evals) {
                 witness.set_extension_target(t, x)?;
             }
-            for (&t, &x) in st
-                .merkle_proof
-                .siblings
-                .iter()
-                .zip_eq(&s.merkle_proof.siblings)
-            {
-                witness.set_hash_target(t, x)?;
+
+            let target_len = st.merkle_proof.siblings.len();
+            let siblings_len = s.merkle_proof.siblings.len();
+
+            if target_len < siblings_len {
+                return Err(anyhow!("fri_proof->query_round_proofs->steps->merkle_proof->siblings' target length is less than the proof length"));
+            }
+
+            // Set overlapping elements
+            for i in 0..siblings_len {
+                witness.set_hash_target(st.merkle_proof.siblings[i], s.merkle_proof.siblings[i])?;
+            }
+
+            // Set remaining elements in target to ZERO if target is longer
+            for i in siblings_len..target_len {
+                witness.set_hash_target(st.merkle_proof.siblings[i], HashOut::ZERO)?;
             }
         }
     }