chore: add set_profile to ContinuationProver (#1001)

benchmarks/src/utils.rs

@@ -109,7 +109,7 @@ where
     // generate_app_proof will emit metrics for proof time of each
     let vk = app_pk.app_vm_pk.vm_pk.get_vk();
     let prover = AppProver::new(app_pk.app_vm_pk, committed_exe)
-        .with_profile()
+        .with_profiling()
         .with_program_name(bench_name.to_string());
     let app_proofs = prover.generate_app_proof(input_stream);
     // 6. Verify STARK proofs.

crates/axvm-sdk/src/prover/agg.rs

@@ -82,9 +82,13 @@ impl AggStarkProver {
         self
     }
 
-    pub fn with_profile(mut self) -> Self {
-        self.profile = true;
-        self.leaf_prover.profile = true;
+    pub fn set_profile(&mut self, profile: bool) -> &mut Self {
+        self.profile = profile;
+        self.leaf_prover.profile = profile;
+        self
+    }
+    pub fn with_profiling(mut self) -> Self {
+        self.set_profile(true);
         self
     }
 

crates/axvm-sdk/src/prover/app.rs

@@ -39,17 +39,18 @@ impl<VC> AppProver<VC> {
             ),
         }
     }
-
-    pub fn with_profile(mut self) -> Self {
-        self.profile = true;
+    pub fn set_profile(&mut self, profile: bool) -> &mut Self {
+        self.profile = profile;
+        self
+    }
+    pub fn with_profiling(mut self) -> Self {
+        self.set_profile(true);
         self
     }
-
     pub fn set_program_name(&mut self, program_name: impl AsRef<str>) -> &mut Self {
         self.program_name = Some(program_name.as_ref().to_string());
         self
     }
-
     pub fn with_program_name(mut self, program_name: impl AsRef<str>) -> Self {
         self.set_program_name(program_name);
         self

crates/axvm-sdk/src/prover/mod.rs

@@ -51,6 +51,13 @@ impl<VC> ContinuationProver<VC> {
         }
     }
 
+    /// Flag for enabling/disabling profiling.
+    pub fn set_profile(&mut self, profile: bool) -> &mut Self {
+        self.stark_prover.set_profile(profile);
+        // halo2 profiling is set in CompilerConfig when creating Halo2ProvingKey
+        self
+    }
+
     pub fn set_program_name(&mut self, program_name: impl AsRef<str>) -> &mut Self {
         self.stark_prover.set_program_name(program_name);
         self

crates/axvm-sdk/src/prover/stark.rs

@@ -37,6 +37,11 @@ impl<VC> StarkProver<VC> {
             agg_prover: AggStarkProver::new(agg_stark_pk, app_pk.leaf_committed_exe.clone()),
         }
     }
+    pub fn set_profile(&mut self, profile: bool) -> &mut Self {
+        self.app_prover.set_profile(profile);
+        self.agg_prover.set_profile(profile);
+        self
+    }
     pub fn set_program_name(&mut self, program_name: impl AsRef<str>) -> &mut Self {
         self.app_prover.set_program_name(program_name);
         self