aztec_pxe/kernel/
prover.rs

1//! Private kernel prover trait and BB binary integration.
2//!
3//! Ports the TS `PrivateKernelProver` interface and `BBPrivateKernelProver`
4//! implementation that shells out to the `bb` binary for proof generation.
5
6use std::path::{Path, PathBuf};
7use std::process::Stdio;
8
9use async_trait::async_trait;
10use aztec_core::error::Error;
11use serde::{Deserialize, Serialize};
12use tokio::process::Command;
13
14// ---------------------------------------------------------------------------
15// Types
16// ---------------------------------------------------------------------------
17
18/// Output from a kernel circuit simulation or witness generation.
19#[derive(Debug, Clone, Serialize, Deserialize)]
20pub struct PrivateKernelSimulateOutput {
21    /// The circuit public inputs (opaque JSON matching TS types).
22    pub public_inputs: serde_json::Value,
23    /// Verification key data (raw bytes or base64).
24    #[serde(default)]
25    pub verification_key: Vec<u8>,
26    /// Output witness map (serialized).
27    #[serde(default)]
28    pub output_witness: Vec<u8>,
29    /// Circuit bytecode (gzipped ACIR).
30    #[serde(default)]
31    pub bytecode: Vec<u8>,
32}
33
34/// A single step in private execution, used to build a ChonkProof.
35#[derive(Debug, Clone, Serialize, Deserialize)]
36pub struct PrivateExecutionStep {
37    /// Function name for tracing.
38    pub function_name: String,
39    /// Gzipped circuit bytecode.
40    pub bytecode: Vec<u8>,
41    /// Serialized witness map.
42    pub witness: Vec<u8>,
43    /// Verification key bytes.
44    pub vk: Vec<u8>,
45    /// Timings metadata.
46    #[serde(default)]
47    pub timings: StepTimings,
48}
49
50/// Timing metadata for an execution step.
51#[derive(Debug, Clone, Default, Serialize, Deserialize)]
52pub struct StepTimings {
53    /// Witness generation time in ms.
54    pub witgen_ms: u64,
55    /// Gate count (if profiled).
56    pub gate_count: Option<u64>,
57    /// Oracle resolution time in ms.
58    pub oracles_ms: u64,
59}
60
61/// ChonkProof with public inputs — the final aggregated proof.
62#[derive(Debug, Clone, Serialize, Deserialize)]
63pub struct ChonkProofWithPublicInputs {
64    /// The aggregated proof bytes.
65    pub proof: Vec<u8>,
66    /// The public inputs bytes.
67    pub public_inputs: Vec<u8>,
68}
69
70/// Result from BB binary execution.
71#[derive(Debug)]
72pub enum BbResult {
73    Success {
74        duration_ms: u64,
75        proof_path: Option<PathBuf>,
76        vk_directory_path: Option<PathBuf>,
77    },
78    Failure {
79        reason: String,
80        retry: bool,
81    },
82}
83
84/// Proving timings for the entire transaction.
85#[derive(Debug, Clone, Default, Serialize, Deserialize)]
86pub struct ProvingTimings {
87    /// Total proving time in ms.
88    pub total_ms: u64,
89    /// Per-circuit timings.
90    pub circuits: Vec<(String, u64)>,
91}
92
93// ---------------------------------------------------------------------------
94// Trait
95// ---------------------------------------------------------------------------
96
97/// Private kernel prover interface.
98///
99/// Matches the TS `PrivateKernelProver` interface with methods for each
100/// kernel circuit in the proving sequence.
101#[async_trait]
102pub trait PrivateKernelProver: Send + Sync {
103    /// Generate witness + output for the init kernel circuit.
104    async fn generate_init_output(
105        &self,
106        inputs: &serde_json::Value,
107    ) -> Result<PrivateKernelSimulateOutput, Error>;
108
109    /// Simulate the init circuit (no witness generation).
110    async fn simulate_init(
111        &self,
112        inputs: &serde_json::Value,
113    ) -> Result<PrivateKernelSimulateOutput, Error>;
114
115    /// Generate witness + output for an inner kernel circuit.
116    async fn generate_inner_output(
117        &self,
118        inputs: &serde_json::Value,
119    ) -> Result<PrivateKernelSimulateOutput, Error>;
120
121    /// Simulate an inner circuit.
122    async fn simulate_inner(
123        &self,
124        inputs: &serde_json::Value,
125    ) -> Result<PrivateKernelSimulateOutput, Error>;
126
127    /// Generate witness + output for a reset kernel circuit.
128    async fn generate_reset_output(
129        &self,
130        inputs: &serde_json::Value,
131    ) -> Result<PrivateKernelSimulateOutput, Error>;
132
133    /// Simulate a reset circuit.
134    async fn simulate_reset(
135        &self,
136        inputs: &serde_json::Value,
137    ) -> Result<PrivateKernelSimulateOutput, Error>;
138
139    /// Generate witness + output for the tail kernel circuit.
140    async fn generate_tail_output(
141        &self,
142        inputs: &serde_json::Value,
143    ) -> Result<PrivateKernelSimulateOutput, Error>;
144
145    /// Simulate the tail circuit.
146    async fn simulate_tail(
147        &self,
148        inputs: &serde_json::Value,
149    ) -> Result<PrivateKernelSimulateOutput, Error>;
150
151    /// Generate hiding kernel output for rollup path.
152    async fn generate_hiding_to_rollup_output(
153        &self,
154        inputs: &serde_json::Value,
155    ) -> Result<PrivateKernelSimulateOutput, Error>;
156
157    /// Generate hiding kernel output for public path.
158    async fn generate_hiding_to_public_output(
159        &self,
160        inputs: &serde_json::Value,
161    ) -> Result<PrivateKernelSimulateOutput, Error>;
162
163    /// Create the aggregated ChonkProof from execution steps.
164    async fn create_chonk_proof(
165        &self,
166        execution_steps: &[PrivateExecutionStep],
167    ) -> Result<ChonkProofWithPublicInputs, Error>;
168
169    /// Compute gate count for a circuit (profiling).
170    async fn compute_gate_count_for_circuit(
171        &self,
172        bytecode: &[u8],
173        circuit_name: &str,
174    ) -> Result<u64, Error>;
175}
176
177// ---------------------------------------------------------------------------
178// BB Binary Prover Implementation
179// ---------------------------------------------------------------------------
180
181/// Configuration for the BB prover.
182#[derive(Debug, Clone)]
183pub struct BbProverConfig {
184    /// Path to the `bb` binary.
185    pub bb_binary_path: PathBuf,
186    /// Working directory for temporary proof artifacts.
187    pub working_directory: PathBuf,
188    /// Number of threads for hardware concurrency.
189    pub hardware_concurrency: Option<u32>,
190    /// Whether to skip cleanup of temporary files.
191    pub skip_cleanup: bool,
192}
193
194impl Default for BbProverConfig {
195    fn default() -> Self {
196        // Default bb binary path for macOS ARM64
197        let home = std::env::var("HOME").unwrap_or_default();
198        let bb_path = std::env::var("BB_BINARY_PATH").unwrap_or_else(|_| {
199            format!(
200                "{home}/.aztec/versions/4.2.0-aztecnr-rc.2/node_modules/@aztec/bb.js/build/arm64-macos/bb"
201            )
202        });
203        let work_dir = std::env::var("BB_WORKING_DIRECTORY")
204            .unwrap_or_else(|_| format!("{home}/.aztec/bb-working"));
205
206        Self {
207            bb_binary_path: PathBuf::from(bb_path),
208            working_directory: PathBuf::from(work_dir),
209            hardware_concurrency: std::env::var("HARDWARE_CONCURRENCY")
210                .ok()
211                .and_then(|v| v.parse().ok()),
212            skip_cleanup: std::env::var("BB_SKIP_CLEANUP")
213                .map(|v| v == "true" || v == "1")
214                .unwrap_or(false),
215        }
216    }
217}
218
219/// Private kernel prover backed by the Barretenberg (`bb`) binary.
220///
221/// Shells out to the `bb` binary for proof generation, matching the TS
222/// `BBBundlePrivateKernelProver` implementation.
223pub struct BbPrivateKernelProver {
224    config: BbProverConfig,
225}
226
227impl BbPrivateKernelProver {
228    pub fn new(config: BbProverConfig) -> Self {
229        Self { config }
230    }
231
232    /// Create with default configuration (auto-detects bb path).
233    pub fn with_defaults() -> Self {
234        Self::new(BbProverConfig::default())
235    }
236
237    /// Execute the bb binary with given command and arguments.
238    async fn execute_bb(&self, command: &str, args: &[&str]) -> Result<BbResult, Error> {
239        let bb_path = &self.config.bb_binary_path;
240
241        if !bb_path.exists() {
242            return Ok(BbResult::Failure {
243                reason: format!("bb binary not found at {}", bb_path.display()),
244                retry: false,
245            });
246        }
247
248        let start = std::time::Instant::now();
249
250        let mut cmd = Command::new(bb_path);
251        cmd.arg(command);
252        cmd.args(args);
253        cmd.stdout(Stdio::piped());
254        cmd.stderr(Stdio::piped());
255
256        // Set hardware concurrency
257        if let Some(concurrency) = self.config.hardware_concurrency {
258            cmd.env("HARDWARE_CONCURRENCY", concurrency.to_string());
259        }
260
261        tracing::debug!(
262            "Executing BB: {} {} {}",
263            bb_path.display(),
264            command,
265            args.join(" ")
266        );
267
268        let output = cmd
269            .output()
270            .await
271            .map_err(|e| Error::InvalidData(format!("failed to execute bb binary: {e}")))?;
272
273        let duration_ms = start.elapsed().as_millis() as u64;
274
275        if output.status.success() {
276            Ok(BbResult::Success {
277                duration_ms,
278                proof_path: None,
279                vk_directory_path: None,
280            })
281        } else {
282            let stderr = String::from_utf8_lossy(&output.stderr);
283            Ok(BbResult::Failure {
284                reason: format!(
285                    "bb {} failed with exit code {:?}: {}",
286                    command,
287                    output.status.code(),
288                    stderr
289                ),
290                retry: false,
291            })
292        }
293    }
294
295    /// Generate a ChonkProof by writing inputs and invoking bb prove --scheme chonk.
296    async fn execute_chonk_proof(
297        &self,
298        inputs_path: &Path,
299        output_path: &Path,
300    ) -> Result<BbResult, Error> {
301        let args = vec![
302            "-o",
303            output_path.to_str().unwrap_or(""),
304            "--ivc_inputs_path",
305            inputs_path.to_str().unwrap_or(""),
306            "-v",
307            "--scheme",
308            "chonk",
309        ];
310
311        self.execute_bb("prove", &args).await
312    }
313
314    /// Ensure working directory exists.
315    async fn ensure_working_dir(&self) -> Result<PathBuf, Error> {
316        let dir = &self.config.working_directory;
317        tokio::fs::create_dir_all(dir)
318            .await
319            .map_err(|e| Error::InvalidData(format!("failed to create working directory: {e}")))?;
320        Ok(dir.clone())
321    }
322
323    /// Simulate a protocol circuit by executing it with the bb binary.
324    ///
325    /// In the TS implementation, this uses CircuitSimulator (ACVM-based).
326    /// Here we delegate to bb for now, returning opaque JSON public inputs.
327    async fn simulate_circuit(
328        &self,
329        _inputs: &serde_json::Value,
330        circuit_type: &str,
331    ) -> Result<PrivateKernelSimulateOutput, Error> {
332        Err(Error::InvalidData(format!(
333            "kernel circuit simulation for {circuit_type} is not wired to real artifacts yet"
334        )))
335    }
336
337    /// Generate circuit output with witness and VK.
338    async fn generate_circuit_output(
339        &self,
340        _inputs: &serde_json::Value,
341        circuit_type: &str,
342    ) -> Result<PrivateKernelSimulateOutput, Error> {
343        let _ = self.ensure_working_dir().await?;
344        Err(Error::InvalidData(format!(
345            "kernel witness generation for {circuit_type} is not wired to real artifacts yet"
346        )))
347    }
348}
349
350#[async_trait]
351impl PrivateKernelProver for BbPrivateKernelProver {
352    async fn generate_init_output(
353        &self,
354        inputs: &serde_json::Value,
355    ) -> Result<PrivateKernelSimulateOutput, Error> {
356        self.generate_circuit_output(inputs, "PrivateKernelInitArtifact")
357            .await
358    }
359
360    async fn simulate_init(
361        &self,
362        inputs: &serde_json::Value,
363    ) -> Result<PrivateKernelSimulateOutput, Error> {
364        self.simulate_circuit(inputs, "PrivateKernelInitArtifact")
365            .await
366    }
367
368    async fn generate_inner_output(
369        &self,
370        inputs: &serde_json::Value,
371    ) -> Result<PrivateKernelSimulateOutput, Error> {
372        self.generate_circuit_output(inputs, "PrivateKernelInnerArtifact")
373            .await
374    }
375
376    async fn simulate_inner(
377        &self,
378        inputs: &serde_json::Value,
379    ) -> Result<PrivateKernelSimulateOutput, Error> {
380        self.simulate_circuit(inputs, "PrivateKernelInnerArtifact")
381            .await
382    }
383
384    async fn generate_reset_output(
385        &self,
386        inputs: &serde_json::Value,
387    ) -> Result<PrivateKernelSimulateOutput, Error> {
388        self.generate_circuit_output(inputs, "PrivateKernelResetArtifact")
389            .await
390    }
391
392    async fn simulate_reset(
393        &self,
394        inputs: &serde_json::Value,
395    ) -> Result<PrivateKernelSimulateOutput, Error> {
396        self.simulate_circuit(inputs, "PrivateKernelResetArtifact")
397            .await
398    }
399
400    async fn generate_tail_output(
401        &self,
402        inputs: &serde_json::Value,
403    ) -> Result<PrivateKernelSimulateOutput, Error> {
404        self.generate_circuit_output(inputs, "PrivateKernelTailArtifact")
405            .await
406    }
407
408    async fn simulate_tail(
409        &self,
410        inputs: &serde_json::Value,
411    ) -> Result<PrivateKernelSimulateOutput, Error> {
412        self.simulate_circuit(inputs, "PrivateKernelTailArtifact")
413            .await
414    }
415
416    async fn generate_hiding_to_rollup_output(
417        &self,
418        inputs: &serde_json::Value,
419    ) -> Result<PrivateKernelSimulateOutput, Error> {
420        self.generate_circuit_output(inputs, "HidingKernelToRollup")
421            .await
422    }
423
424    async fn generate_hiding_to_public_output(
425        &self,
426        inputs: &serde_json::Value,
427    ) -> Result<PrivateKernelSimulateOutput, Error> {
428        self.generate_circuit_output(inputs, "HidingKernelToPublic")
429            .await
430    }
431
432    async fn create_chonk_proof(
433        &self,
434        execution_steps: &[PrivateExecutionStep],
435    ) -> Result<ChonkProofWithPublicInputs, Error> {
436        let start = std::time::Instant::now();
437        tracing::info!(
438            "Generating ClientIVC proof ({} steps)...",
439            execution_steps.len()
440        );
441
442        let work_dir = self.ensure_working_dir().await?;
443        let chonk_dir = work_dir.join("chonk");
444        tokio::fs::create_dir_all(&chonk_dir)
445            .await
446            .map_err(|e| Error::InvalidData(format!("failed to create chonk dir: {e}")))?;
447
448        // Write execution steps as IVC inputs
449        let inputs_path = chonk_dir.join("ivc_inputs.bin");
450        let steps_data = serde_json::to_vec(execution_steps)?;
451        tokio::fs::write(&inputs_path, &steps_data)
452            .await
453            .map_err(|e| Error::InvalidData(format!("failed to write IVC inputs: {e}")))?;
454
455        let result = self.execute_chonk_proof(&inputs_path, &chonk_dir).await?;
456
457        match result {
458            BbResult::Success { duration_ms, .. } => {
459                tracing::info!("Generated ClientIVC proof in {}ms", duration_ms);
460
461                // Read proof and public inputs from output directory
462                let proof_path = chonk_dir.join("proof");
463                let pi_path = chonk_dir.join("public_inputs");
464
465                let proof = tokio::fs::read(&proof_path).await.unwrap_or_default();
466                let public_inputs = tokio::fs::read(&pi_path).await.unwrap_or_default();
467
468                // Cleanup if configured
469                if !self.config.skip_cleanup {
470                    let _ = tokio::fs::remove_dir_all(&chonk_dir).await;
471                }
472
473                Ok(ChonkProofWithPublicInputs {
474                    proof,
475                    public_inputs,
476                })
477            }
478            BbResult::Failure { reason, .. } => {
479                let elapsed = start.elapsed().as_millis();
480                tracing::error!(
481                    "ChonkProof generation failed after {}ms: {}",
482                    elapsed,
483                    reason
484                );
485                Err(Error::InvalidData(format!(
486                    "ChonkProof generation failed: {reason}"
487                )))
488            }
489        }
490    }
491
492    async fn compute_gate_count_for_circuit(
493        &self,
494        bytecode: &[u8],
495        circuit_name: &str,
496    ) -> Result<u64, Error> {
497        let work_dir = self.ensure_working_dir().await?;
498        let gates_dir = work_dir.join("gates");
499        tokio::fs::create_dir_all(&gates_dir)
500            .await
501            .map_err(|e| Error::InvalidData(format!("failed to create gates dir: {e}")))?;
502
503        let bytecode_path = gates_dir.join(format!("{circuit_name}-bytecode"));
504        tokio::fs::write(&bytecode_path, bytecode)
505            .await
506            .map_err(|e| Error::InvalidData(format!("failed to write bytecode: {e}")))?;
507
508        let bytecode_str = bytecode_path.to_str().unwrap_or("");
509        let args = vec!["--scheme", "ultra_honk", "-b", bytecode_str, "-v"];
510
511        let result = self.execute_bb("gates", &args).await?;
512
513        if !self.config.skip_cleanup {
514            let _ = tokio::fs::remove_dir_all(&gates_dir).await;
515        }
516
517        match result {
518            BbResult::Success { .. } => {
519                // In full implementation, parse stdout for circuit_size
520                Ok(0)
521            }
522            BbResult::Failure { reason, .. } => {
523                Err(Error::InvalidData(format!("gate count failed: {reason}")))
524            }
525        }
526    }
527}