Merge branch 'better_xmss'

Author: TomWambsgans

Cargo.lock

@@ -63,15 +63,14 @@ sha3 = "0.10.8"
 derive_more = { version = "2.0.1", features = ["full"] }
 pest = "2.7"
 pest_derive = "2.7"
-itertools = "0.10.5"
+itertools = "0.14.0"
 colored = "3.0.0"
 tracing = "0.1.26"
 serde_json = "1.0.145"
 num_enum = "0.7.5"
 tracing-subscriber = { version = "0.3.19", features = ["std", "env-filter"] }
 tracing-forest = { version = "0.3.0", features = ["ansi", "smallvec"] }
 p3-koala-bear = { git = "https://github.com/TomWambsgans/Plonky3.git", branch = "lean-multisig" }
-p3-baby-bear = { git = "https://github.com/TomWambsgans/Plonky3.git", branch = "lean-multisig" }
 p3-poseidon2 = { git = "https://github.com/TomWambsgans/Plonky3.git", branch = "lean-multisig" }
 p3-symmetric = { git = "https://github.com/TomWambsgans/Plonky3.git", branch = "lean-multisig" }
 p3-air = { git = "https://github.com/TomWambsgans/Plonky3.git", branch = "lean-multisig" }
@@ -86,7 +85,11 @@ multilinear-toolkit = { git = "https://github.com/leanEthereum/multilinear-toolk
 [dependencies]
 clap.workspace = true
 rec_aggregation.workspace = true
+xmss.workspace = true
 poseidon_circuit.workspace = true
+p3-koala-bear.workspace = true
+multilinear-toolkit.workspace = true
+whir-p3.workspace = true
 
 # [patch."https://github.com/TomWambsgans/Plonky3.git"]
 # p3-koala-bear = { path = "../zk/Plonky3/koala-bear" }

Cargo.toml

@@ -46,7 +46,7 @@ RUSTFLAGS='-C target-cpu=native' cargo run --release -- recursion
 ### XMSS aggregation
 
 ```console
-RUSTFLAGS='-C target-cpu=native' cargo run --release -- xmss --n-signatures 990
+RUSTFLAGS='-C target-cpu=native' cargo run --release -- xmss --n-signatures 1000
 ```
 
 [Trivial encoding](docs/XMSS_trivial_encoding.pdf) (for now).

README.md

@@ -108,7 +108,7 @@ fn test_air_helper<const VIRTUAL_COLUMN: bool>() {
     const N_PREPROCESSED_COLUMNS: usize = 5;
     let log_n_rows = 12;
     let n_rows = 1 << log_n_rows;
-    let mut prover_state = build_prover_state::<EF>();
+    let mut prover_state = build_prover_state::<EF>(false);
 
     let (columns_plus_one_f, columns_plus_one_ef) = generate_trace::<N_COLUMNS, N_PREPROCESSED_COLUMNS>(n_rows + 1);
     let columns_ref_f = columns_plus_one_f.iter().map(|col| &col[..n_rows]).collect::<Vec<_>>();

crates/air/tests/complex_air.rs

@@ -60,7 +60,7 @@ fn generate_trace(n_rows: usize) -> (Vec<F>, Vec<EF>) {
 fn test_air_fibonacci() {
     let log_n_rows = 14;
     let n_rows = 1 << log_n_rows;
-    let mut prover_state = build_prover_state::<EF>();
+    let mut prover_state = build_prover_state::<EF>(false);
 
     let (columns_plus_one_f, columns_plus_one_ef) = generate_trace(n_rows + 1);
     let columns_ref_f = vec![&columns_plus_one_f[..n_rows]];

crates/air/tests/fib_air.rs

@@ -1632,6 +1632,11 @@ fn handle_const_arguments_helper(
                 // TODO we should unroll before const arguments handling
                 handle_const_arguments_helper(body, constant_functions, new_functions);
             }
+            Line::Match { arms, .. } => {
+                for (_, arm) in arms {
+                    handle_const_arguments_helper(arm, constant_functions, new_functions);
+                }
+            }
             _ => {}
         }
     }

crates/lean_compiler/src/a_simplify_lang.rs

@@ -23,7 +23,7 @@ pub fn prove_execution(
     no_vec_runtime_memory: usize, // size of the "non-vectorized" runtime memory
     vm_profiler: bool,
     (poseidons_16_precomputed, poseidons_24_precomputed): (&Poseidon16History, &Poseidon24History),
-) -> (Vec<PF<EF>>, usize, String) {
+) -> (Vec<PF<EF>>, String) {
     let mut exec_summary = String::new();
     let ExecutionTrace {
         traces,
@@ -82,7 +82,7 @@ pub fn prove_execution(
         })
     });
 
-    let mut prover_state = build_prover_state::<EF>();
+    let mut prover_state = build_prover_state::<EF>(false);
     prover_state.add_base_scalars(
         &[
             vec![private_memory.len()],
@@ -380,11 +380,7 @@ pub fn prove_execution(
         &packed_pcs_witness_extension.packed_polynomial.by_ref(),
     );
 
-    (
-        prover_state.proof_data().to_vec(),
-        prover_state.proof_size(),
-        exec_summary,
-    )
+    (prover_state.proof_data().to_vec(), exec_summary)
 }
 
 fn prove_bus_and_air(

crates/lean_prover/src/prove_execution.rs

@@ -24,7 +24,7 @@ pub fn verify_execution(
     proof_data: Vec<PF<EF>>,
     whir_config_builder: WhirConfigBuilder,
 ) -> Result<(), ProofError> {
-    let mut verifier_state = VerifierState::new(proof_data, build_challenger());
+    let mut verifier_state = VerifierState::new(proof_data, build_challenger(), false);
 
     let p16_gkr_layers = PoseidonGKRLayers::<16, N_COMMITED_CUBES_P16>::build(Some(VECTOR_LEN));
     let p24_gkr_layers = PoseidonGKRLayers::<24, N_COMMITED_CUBES_P24>::build(None);

crates/lean_prover/src/verify_execution.rs

@@ -115,7 +115,7 @@ fn test_zk_vm_helper(program_str: &str, (public_input, private_input): (&[F], &[
     utils::init_tracing();
     let bytecode = compile_program(program_str.to_string());
     let time = std::time::Instant::now();
-    let (proof_data, _, summary) = prove_execution(
+    let (proof_data, summary) = prove_execution(
         &bytecode,
         (public_input, private_input),
         whir_config_builder(),

crates/lean_prover/tests/test_zkvm.rs

@@ -257,22 +257,6 @@ fn execute_bytecode_helper(
         poseidons_24_precomputed.len(),
         "Warning: not all precomputed Poseidon24 were used"
     );
-    tracing::info!(
-        "{}% of Poseidon16 precomputed",
-        if poseidons_16_precomputed.is_empty() {
-            0.
-        } else {
-            (n_poseidon16_precomputed_used as f64 / poseidons_16_precomputed.len() as f64) * 100.0
-        }
-    );
-    tracing::info!(
-        "{}% of Poseidon24 precomputed",
-        if poseidons_24_precomputed.is_empty() {
-            0.
-        } else {
-            (n_poseidon24_precomputed_used as f64 / poseidons_24_precomputed.len() as f64) * 100.0
-        }
-    );
 
     assert_eq!(pc, ENDING_PC);
     pcs.push(pc);
@@ -330,6 +314,18 @@ fn execute_bytecode_helper(
         used_memory_cells as f64 / runtime_memory_size as f64 * 100.0
     ));
 
+    // precomputed poseidons
+    summary.push_str(&format!(
+        "Poseidon2_16 precomputed used: {}/{}\n",
+        pretty_integer(n_poseidon16_precomputed_used),
+        pretty_integer(poseidons_16_precomputed.len())
+    ));
+    summary.push_str(&format!(
+        "Poseidon2_24 precomputed used: {}/{}\n",
+        pretty_integer(n_poseidon24_precomputed_used),
+        pretty_integer(poseidons_24_precomputed.len())
+    ));
+
     summary.push('\n');
 
     if traces[Table::poseidon16().index()].base[0].len() + traces[Table::poseidon24().index()].base[0].len() > 0 {