refactor: BlockProverInput::prove is always checked, regardless of cf…

…g feature
0xPolygonZero · Jun 21, 2024 · 8e2febd · 8e2febd
1 parent daf294a
commit 8e2febd
Showing 1 changed file with 75 additions and 75 deletions.
diff --git a/zero_bin/prover/src/lib.rs b/zero_bin/prover/src/lib.rs
@@ -2,7 +2,7 @@ use std::future::Future;
 use std::path::PathBuf;
 
 use alloy::primitives::{BlockNumber, U256};
-use anyhow::{Context, Result};
+use anyhow::Context as _;
 use futures::{future::BoxFuture, stream::FuturesOrdered, FutureExt, TryFutureExt, TryStreamExt};
 use num_traits::ToPrimitive as _;
 use ops::TxProof;
@@ -29,84 +29,84 @@ impl BlockProverInput {
         self.other_data.b_data.b_meta.block_number.into()
     }
 
-    #[cfg(not(feature = "test_only"))]
     pub async fn prove(
         self,
         runtime: &Runtime,
-        previous: Option<impl Future<Output = Result<GeneratedBlockProof>>>,
+        previous: Option<impl Future<Output = anyhow::Result<GeneratedBlockProof>>>,
         save_inputs_on_error: bool,
-    ) -> Result<GeneratedBlockProof> {
-        use anyhow::Context as _;
-
-        let block_number = self.get_block_number();
-
-        let txs =
-            trace_decoder::entrypoint(self.block_trace, self.other_data, |_| unimplemented!())?;
-
-        let agg_proof = IndexedStream::from(txs)
-            .map(&TxProof {
-                save_inputs_on_error,
-            })
-            .fold(&ops::AggProof {
-                save_inputs_on_error,
-            })
-            .run(runtime)
-            .await?;
-
-        if let proof_gen::proof_types::AggregatableProof::Agg(proof) = agg_proof {
-            let block_number = block_number
-                .to_u64()
-                .context("block number overflows u64")?;
-            let prev = match previous {
-                Some(it) => Some(it.await?),
-                None => None,
-            };
-
-            let block_proof = paladin::directive::Literal(proof)
-                .map(&ops::BlockProof {
-                    prev,
-                    save_inputs_on_error,
+    ) -> anyhow::Result<GeneratedBlockProof> {
+        match cfg!(feature = "test_only") {
+            true => {
+                let block_number = self.get_block_number();
+                info!("Testing witness generation for block {block_number}.");
+
+                let txs = trace_decoder::entrypoint(
+                    self.block_trace,
+                    self.other_data,
+                    |_| unimplemented!(),
+                )?;
+
+                IndexedStream::from(txs)
+                    .map(&TxProof {
+                        save_inputs_on_error,
+                    })
+                    .run(runtime)
+                    .await?
+                    .try_collect::<Vec<_>>()
+                    .await?;
+
+                // Dummy proof to match expected output type.
+                Ok(GeneratedBlockProof {
+                    b_height: block_number
+                        .to_u64()
+                        .expect("Block number should fit in a u64"),
+                    intern: proof_gen::proof_gen::dummy_proof()?,
                 })
-                .run(runtime)
-                .await?;
-
-            info!("Successfully proved block {block_number}");
-            Ok(block_proof.0)
-        } else {
-            anyhow::bail!("AggProof is is not GeneratedAggProof")
+            }
+            false => {
+                let block_number = self.get_block_number();
+
+                let txs = trace_decoder::entrypoint(
+                    self.block_trace,
+                    self.other_data,
+                    |_| unimplemented!(),
+                )?;
+
+                let agg_proof = IndexedStream::from(txs)
+                    .map(&TxProof {
+                        save_inputs_on_error,
+                    })
+                    .fold(&ops::AggProof {
+                        save_inputs_on_error,
+                    })
+                    .run(runtime)
+                    .await?;
+
+                if let proof_gen::proof_types::AggregatableProof::Agg(proof) = agg_proof {
+                    let block_number = block_number
+                        .to_u64()
+                        .context("block number overflows u64")?;
+                    let prev = match previous {
+                        Some(it) => Some(it.await?),
+                        None => None,
+                    };
+
+                    let block_proof = paladin::directive::Literal(proof)
+                        .map(&ops::BlockProof {
+                            prev,
+                            save_inputs_on_error,
+                        })
+                        .run(runtime)
+                        .await?;
+
+                    info!("Successfully proved block {block_number}");
+                    Ok(block_proof.0)
+                } else {
+                    anyhow::bail!("AggProof is is not GeneratedAggProof")
+                }
+            }
         }
     }
-
-    #[cfg(feature = "test_only")]
-    pub async fn prove(
-        self,
-        runtime: &Runtime,
-        _previous: Option<impl Future<Output = Result<GeneratedBlockProof>>>,
-        save_inputs_on_error: bool,
-    ) -> Result<GeneratedBlockProof> {
-        let block_number = self.get_block_number();
-        info!("Testing witness generation for block {block_number}.");
-
-        let txs =
-            trace_decoder::entrypoint(self.block_trace, self.other_data, |_| unimplemented!())?;
-
-        IndexedStream::from(txs)
-            .map(&TxProof {
-                save_inputs_on_error,
-            })
-            .run(runtime)
-            .await?
-            .try_collect::<Vec<_>>()
-            .await?;
-
-        // Dummy proof to match expected output type.
-        Ok(GeneratedBlockProof {
-            b_height: block_number
-                .to_u64()
-                .expect("Block number should fit in a u64"),
-            intern: proof_gen::proof_gen::dummy_proof()?,
-        })
-    }
 }
 
 #[derive(Debug, Deserialize, Serialize)]
@@ -124,8 +124,8 @@ impl ProverInput {
         previous_proof: Option<GeneratedBlockProof>,
         save_inputs_on_error: bool,
         proof_output_dir: Option<PathBuf>,
-    ) -> Result<Vec<(BlockNumber, Option<GeneratedBlockProof>)>> {
-        let mut prev: Option<BoxFuture<Result<GeneratedBlockProof>>> =
+    ) -> anyhow::Result<Vec<(BlockNumber, Option<GeneratedBlockProof>)>> {
+        let mut prev: Option<BoxFuture<anyhow::Result<GeneratedBlockProof>>> =
             previous_proof.map(|proof| Box::pin(futures::future::ok(proof)) as BoxFuture<_>);
 
         let results: FuturesOrdered<_> = self
@@ -175,7 +175,7 @@ impl ProverInput {
     pub(crate) async fn write_proof(
         output_dir: Option<PathBuf>,
         proof: &GeneratedBlockProof,
-    ) -> Result<()> {
+    ) -> anyhow::Result<()> {
         let proof_serialized = serde_json::to_vec(proof)?;
         let block_proof_file_path =
             output_dir.map(|path| generate_block_proof_file_name(&path.to_str(), proof.b_height));