fix program poseidon panic blocking

AbdelStark · Aug 5, 2024 · 04c1c2a · 04c1c2a
1 parent 11a3250
commit 04c1c2a
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 15 deletions.
diff --git a/crates/core/src/prover_service.rs b/crates/core/src/prover_service.rs
@@ -189,8 +189,13 @@ impl ProverService {
                         poseidon_req.log_n_instances
                     );
 
-                    if log_n_rows <= MIN_FFT_LOG_SIZE as u32
-                        || poseidon_req.log_n_instances <= MIN_FFT_LOG_SIZE as u32
+                    println!(
+                        "poseidon_req.log_n_instances < MIN_FFT_LOG_SIZE{}",
+                        poseidon_req.log_n_instances < MIN_FFT_LOG_SIZE
+                    );
+
+                    if poseidon_req.log_n_instances < MIN_FFT_LOG_SIZE as u32
+                        || log_n_rows < MIN_FFT_LOG_SIZE as u32
                     {
                         println!(
                             "log_n_elements >= MIN_FFT_LOG_SIZE as usize {}",
@@ -206,14 +211,14 @@ impl ProverService {
                     let poseidon = PoseidonStruct::new(poseidon_req.log_n_instances);
 
                     // TODO fix prove poseidon with inputs_requirements
-                    // match poseidon {
-                    //     Ok(poseidon) => match poseidon.prove() {
-                    //         Ok(proof) => Ok(GenericProvingResponse::new(request.clone(), proof)),
-                    //         Err(e) => Err(e.to_string()),
-                    //     },
-                    //     Err(e) => Err(ProvingError::ConstraintsNotSatisfied.to_string()),
-                    // }
-                    Err(ProvingError::ConstraintsNotSatisfied.to_string())
+                    match poseidon {
+                        Ok(poseidon) => match poseidon.prove() {
+                            Ok(proof) => Ok(GenericProvingResponse::new(request.clone(), proof)),
+                            Err(e) => Err(e.to_string()),
+                        },
+                        Err(e) => Err(ProvingError::ConstraintsNotSatisfied.to_string()),
+                    }
+                    // Err(ProvingError::ConstraintsNotSatisfied.to_string())
                 }
                 ProgramInternalContractName::WideFibonnaciProvingRequest => {
                     // Err(ProvingError::ConstraintsNotSatisfied.to_string())

diff --git a/crates/stwo_wasm/src/poseidon.rs b/crates/stwo_wasm/src/poseidon.rs
@@ -3,6 +3,7 @@ use stwo_prover::constraint_framework::constant_columns::gen_is_first;
 use stwo_prover::constraint_framework::logup::LookupElements;
 use stwo_prover::core::backend::simd::fft::MIN_FFT_LOG_SIZE;
 use stwo_prover::core::backend::simd::SimdBackend;
+use stwo_prover::core::backend::CpuBackend;
 use stwo_prover::core::channel::{Blake2sChannel, Channel};
 use stwo_prover::core::fields::m31::BaseField;
 use stwo_prover::core::fields::IntoSlice;
@@ -21,6 +22,7 @@ use stwo_prover::examples::poseidon::{
     PoseidonAir,
     PoseidonComponent, //  PoseidonComponent,
 };
+use stwo_prover::trace_generation::commit_and_prove;
 use wasm_bindgen::prelude::*;
 
 use crate::StwoResult;
@@ -107,6 +109,7 @@ impl PoseidonStruct {
     }
     pub fn prove(&self) -> Result<StarkProof, ProvingError> {
         // let (trace, lookup_data) = gen_trace(self.air.component.log_n_rows);
+        // let res = PoseidonStruct::prove_poseidon(self.air.component.log_n_rows);
         let res = PoseidonStruct::prove_poseidon(self.air.component.log_n_rows);
 
         match res {
@@ -122,7 +125,10 @@ impl PoseidonStruct {
     }
 
     // @TODO handle correctly error to not crash
-    pub fn prove_poseidon(log_n_instances: u32) -> Result<StarkProof, String> {
+    pub fn prove_poseidon(
+        // air:PoseidonAir,
+        log_n_instances: u32,
+    ) -> Result<StarkProof, String> {
         if log_n_instances < N_LOG_INSTANCES_PER_ROW as u32 {
             return Err("log_n_rows < LOG_N_LANES".to_string());
         }
@@ -141,6 +147,23 @@ impl PoseidonStruct {
             return Err("llog_n_elements >= MIN_FFT_LOG_SIZE as usize".to_string());
         }
 
+        if log_n_rows < MIN_FFT_LOG_SIZE || log_n_instances < MIN_FFT_LOG_SIZE {
+            println!(
+                " log_n_rows < MIN_FFT_LOG_SIZE || log_n_instances < MIN_FFT_LOG_SIZE  {}",
+                log_n_rows >= LOG_N_LANES
+            );
+            return Err(
+                " log_n_rows < MIN_FFT_LOG_SIZE || log_n_instances < MIN_FFT_LOG_SIZE ".to_string(),
+            );
+        }
+        if log_n_rows < LOG_N_LANES + 2 {
+            println!(
+                "log_n_rows < LOG_N_LANES + 2  {}",
+                log_n_rows < LOG_N_LANES + 2
+            );
+            return Err("log_n_rows < LOG_N_LANES + 2 ".to_string());
+        }
+
         // Precompute twiddles.
         // let span = span!(Level::INFO, "Precompute twiddles").entered();
         let twiddles = SimdBackend::precompute_twiddles(
@@ -161,7 +184,6 @@ impl PoseidonStruct {
 
         // Draw lookup element.
         let lookup_elements = LookupElements::draw(channel);
-
         // Interaction trace.
         let (trace, claimed_sum) = gen_interaction_trace(log_n_rows, lookup_data, lookup_elements);
         let mut tree_builder = commitment_scheme.tree_builder();
@@ -180,15 +202,21 @@ impl PoseidonStruct {
             claimed_sum,
         };
         let air = PoseidonAir { component };
-        let proof = prove::<SimdBackend>(
+        let proof = prove(
             &air,
             channel,
             &InteractionElements::default(),
             commitment_scheme,
         )
-        .unwrap();
+        .map_err(|op| Err::<StarkProof, ProvingError>(op));
 
-        Ok(proof)
+        match proof {
+            Ok(p) => Ok(p),
+            Err(e) => Err("PROVING_ERROR".to_string()),
+        }
+        // } else {
+        //     return Err("llog_n_elements >= MIN_FFT_LOG_SIZE as usize".to_string());
+        // }
     }
 
     pub fn verify(&self, proof: StarkProof) -> Result<(), VerificationError> {