1use 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#[derive(Debug, Clone, Serialize, Deserialize)]
20pub struct PrivateKernelSimulateOutput {
21 pub public_inputs: serde_json::Value,
23 #[serde(default)]
25 pub verification_key: Vec<u8>,
26 #[serde(default)]
28 pub output_witness: Vec<u8>,
29 #[serde(default)]
31 pub bytecode: Vec<u8>,
32}
33
34#[derive(Debug, Clone, Serialize, Deserialize)]
36pub struct PrivateExecutionStep {
37 pub function_name: String,
39 pub bytecode: Vec<u8>,
41 pub witness: Vec<u8>,
43 pub vk: Vec<u8>,
45 #[serde(default)]
47 pub timings: StepTimings,
48}
49
50#[derive(Debug, Clone, Default, Serialize, Deserialize)]
52pub struct StepTimings {
53 pub witgen_ms: u64,
55 pub gate_count: Option<u64>,
57 pub oracles_ms: u64,
59}
60
61#[derive(Debug, Clone, Serialize, Deserialize)]
63pub struct ChonkProofWithPublicInputs {
64 pub proof: Vec<u8>,
66 pub public_inputs: Vec<u8>,
68}
69
70#[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#[derive(Debug, Clone, Default, Serialize, Deserialize)]
86pub struct ProvingTimings {
87 pub total_ms: u64,
89 pub circuits: Vec<(String, u64)>,
91}
92
93#[async_trait]
102pub trait PrivateKernelProver: Send + Sync {
103 async fn generate_init_output(
105 &self,
106 inputs: &serde_json::Value,
107 ) -> Result<PrivateKernelSimulateOutput, Error>;
108
109 async fn simulate_init(
111 &self,
112 inputs: &serde_json::Value,
113 ) -> Result<PrivateKernelSimulateOutput, Error>;
114
115 async fn generate_inner_output(
117 &self,
118 inputs: &serde_json::Value,
119 ) -> Result<PrivateKernelSimulateOutput, Error>;
120
121 async fn simulate_inner(
123 &self,
124 inputs: &serde_json::Value,
125 ) -> Result<PrivateKernelSimulateOutput, Error>;
126
127 async fn generate_reset_output(
129 &self,
130 inputs: &serde_json::Value,
131 ) -> Result<PrivateKernelSimulateOutput, Error>;
132
133 async fn simulate_reset(
135 &self,
136 inputs: &serde_json::Value,
137 ) -> Result<PrivateKernelSimulateOutput, Error>;
138
139 async fn generate_tail_output(
141 &self,
142 inputs: &serde_json::Value,
143 ) -> Result<PrivateKernelSimulateOutput, Error>;
144
145 async fn simulate_tail(
147 &self,
148 inputs: &serde_json::Value,
149 ) -> Result<PrivateKernelSimulateOutput, Error>;
150
151 async fn generate_hiding_to_rollup_output(
153 &self,
154 inputs: &serde_json::Value,
155 ) -> Result<PrivateKernelSimulateOutput, Error>;
156
157 async fn generate_hiding_to_public_output(
159 &self,
160 inputs: &serde_json::Value,
161 ) -> Result<PrivateKernelSimulateOutput, Error>;
162
163 async fn create_chonk_proof(
165 &self,
166 execution_steps: &[PrivateExecutionStep],
167 ) -> Result<ChonkProofWithPublicInputs, Error>;
168
169 async fn compute_gate_count_for_circuit(
171 &self,
172 bytecode: &[u8],
173 circuit_name: &str,
174 ) -> Result<u64, Error>;
175}
176
177#[derive(Debug, Clone)]
183pub struct BbProverConfig {
184 pub bb_binary_path: PathBuf,
186 pub working_directory: PathBuf,
188 pub hardware_concurrency: Option<u32>,
190 pub skip_cleanup: bool,
192}
193
194impl Default for BbProverConfig {
195 fn default() -> Self {
196 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
219pub struct BbPrivateKernelProver {
224 config: BbProverConfig,
225}
226
227impl BbPrivateKernelProver {
228 pub fn new(config: BbProverConfig) -> Self {
229 Self { config }
230 }
231
232 pub fn with_defaults() -> Self {
234 Self::new(BbProverConfig::default())
235 }
236
237 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 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 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 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 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 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 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 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 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 Ok(0)
521 }
522 BbResult::Failure { reason, .. } => {
523 Err(Error::InvalidData(format!("gate count failed: {reason}")))
524 }
525 }
526 }
527}