Improved xmss

Cargo.toml

@@ -64,6 +64,7 @@ strum = { version = "0.27.2", features = ["derive"] }
 serde = { version = "1.0.228", features = ["derive"] }
 bincode = "1.3.3"
 num_enum = "0.7.5"
+owo-colors = "4.2.3"
 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" }
@@ -88,5 +89,6 @@ p3-koala-bear.workspace = true
 lean_vm.workspace = true
 multilinear-toolkit.workspace = true
 
+
 [profile.release]
 lto = "thin"

README.md

@@ -20,30 +20,38 @@ The VM design is inspired by the famous [Cairo paper](https://eprint.iacr.org/20
 
 ## Security
 
-123 bits of security. Johnson bound + degree 5 extension of koala-bear -> **no proximity gaps conjecture**. (TODO 128 bits? this would require hash digests bigger than 8 koala-bears).
+123 bits of provable security, given by Johnson bound + degree 5 extension of koala-bear. (128 bits would require hash digests of more than 8 field elements, todo?). In the benchmarks, we also display performance with conjectured security, even though leanVM targets the proven regime by default.
 
-## Benchmarks (Slightly outdated, new benchmarks incoming)
+## Benchmarks
 
 Machine: M4 Max 48GB (CPU only)
 
-| Benchmark                  | Current              | Target          |
-| -------------------------- | -------------------- | --------------- |
-| Poseidon2 (16 koala-bears) | `560K Poseidon2 / s` | n/a             |
-| 2 -> 1 Recursion           | `1.15 s`             | `0.25 s `       |
-| XMSS aggregation           | `554 XMSS / s`       | `1000 XMSS / s` |
-
 *Expect incoming perf improvements.*
 
-To reproduce:
-- `cargo run --release -- poseidon --log-n-perms 20`
-- `cargo run --release -- recursion --n 2`
-- `cargo run --release -- xmss --n-signatures 1350`
+### XMSS aggregation
+
+```
+cargo run --release -- xmss --n-signatures 1350
+```
+
+| WHIR rate \ regime | Proven               | Conjectured          |
+| ------------------ | -------------------- | -------------------- |
+| 1/2                | 530 XMSS/s - 383 KiB | 530 XMSS/s - 209 KiB |
+| 1/4                | 420 XMSS/s - 252 KiB | 420 XMSS/s - 148 KiB |
+
+(Proving throughput - proof size)
+
+### Recursion
 
-## Proof size
+```
+cargo run --release -- recursion --n 2
+```
 
-WHIR intial rate = 1/4 -> proof size ≈ 225 KiB. (150 KiB with rate 1/16, and < 100 KiB is possible with poximity gaps conjecture + rate 1/16).
+2 to 1 recursion (WHIR rate = 1/4):
 
-(TODO: remaining optimization = [2024/108](https://eprint.iacr.org/2024/108.pdf) section 3.1)
+| Proven          | Conjectured     |
+| --------------- | --------------- |
+| 1.10s - 223 KiB | 1.05s - 134 KiB |
 
 ## Credits
 

TODO.md

@@ -11,9 +11,9 @@
 - Avoid the embedding overhead in logup, when denominators = "c - index"
 - Proof size: replace all equality checks in the verifier algo by value deduction
 - Poseidon in 'Compression' mode everywhere (except in 'Sponge' mode? cf. eprint 2014/223)
-- XMSS: move from toy implem (usefull for benchmark) to a secure implem
 - Recursion: Remove the few hardcoded constants that depend on the guest execution (cycles etc)
 - About the ordering of the variables in sumchecks, currently we do as follows:
+- [2024/108](https://eprint.iacr.org/2024/108.pdf) section 3.1
 
 [a, b, c, d, e, f, g, h]                                        (1st round of sumcheck)
 [(a-r).a + r.e, (1-r).b + r.f, (1-r).c + r.g, (1-r).d + r.h]    (2nd round of sumcheck)

crates/lean_compiler/snark_lib.py

@@ -120,3 +120,6 @@ def match_range(value: int, *args):
         if value in rng:
             return fn(value)
     raise AssertionError(f"Value {value} not in any range")
+
+def hint_decompose_bits_xmss(*args):
+    _ = args

crates/lean_prover/Cargo.toml

@@ -17,6 +17,7 @@ p3-koala-bear.workspace = true
 p3-symmetric.workspace = true
 p3-util.workspace = true
 tracing.workspace = true
+owo-colors.workspace = true
 air.workspace = true
 sub_protocols.workspace = true
 lean_vm.workspace = true

crates/lean_prover/src/lib.rs

@@ -18,15 +18,16 @@ use trace_gen::*;
 // so ≈ 123.92 bits of security against collisions
 pub const SECURITY_BITS: usize = 123; // TODO 128 bits security? (with Poseidon over 20 field elements or with a more subtle soundness analysis (cf. https://eprint.iacr.org/2021/188.pdf))
 
-// Provable security (no proximity gaps conjectures)
-pub const SECURITY_REGIME: SecurityAssumption = SecurityAssumption::JohnsonBound;
-
 pub const GRINDING_BITS: usize = 18;
 
-pub fn default_whir_config(starting_log_inv_rate: usize) -> WhirConfigBuilder {
+pub fn default_whir_config(starting_log_inv_rate: usize, prox_gaps_conjecture: bool) -> WhirConfigBuilder {
     WhirConfigBuilder {
         folding_factor: FoldingFactor::new(7, 5),
-        soundness_type: SECURITY_REGIME,
+        soundness_type: if prox_gaps_conjecture {
+            SecurityAssumption::CapacityBound // TODO update formula with State of Art Conjecture
+        } else {
+            SecurityAssumption::JohnsonBound
+        },
         pow_bits: GRINDING_BITS,
         max_num_variables_to_send_coeffs: 9,
         rs_domain_initial_reduction_factor: 5,

crates/lean_prover/src/prove_execution.rs

@@ -4,6 +4,7 @@ use crate::*;
 use air::prove_air;
 use lean_vm::*;
 
+use owo_colors::OwoColorize;
 use sub_protocols::*;
 use tracing::info_span;
 use utils::build_prover_state;
@@ -62,13 +63,14 @@ pub fn prove_execution(
     let mut table_log = String::new();
     for (table, trace) in &traces {
         table_log.push_str(&format!(
-            "{}: 2^{:.2} rows |",
+            "{}: 2^{} * (1 + {:.2}) rows | ",
             table.name(),
-            f64::log2(trace.non_padded_n_rows as f64)
+            trace.log_n_rows - 1,
+            (trace.non_padded_n_rows as f64) / (1 << (trace.log_n_rows - 1)) as f64 - 1.0
         ));
     }
-    table_log.pop(); // remove last '|'
-    info_span!("Trace tables sizes: {}", table_log).in_scope(|| {});
+    table_log = table_log.trim_end_matches(" | ").to_string();
+    tracing::info!("Trace tables sizes: {}", table_log.magenta());
 
     // TODO parrallelize
     let mut memory_acc = F::zero_vec(memory.len());

crates/lean_prover/src/test_zkvm.rs

@@ -20,15 +20,15 @@ fn test_zk_vm_all_precompiles_helper(fuzzing: bool) {
     let program_str = r#"
 DIM = 5
 N = 11
-VECTOR_LEN = 8
+DIGEST_LEN = 8
 
 # Dot product precompile:
 BE = 1  # base-extension
 EE = 0  # extension-extension
 
 def main():
     pub_start = NONRESERVED_PROGRAM_INPUT_START
-    poseidon16(pub_start + 4 * VECTOR_LEN, pub_start + 5 * VECTOR_LEN, pub_start + 6 * VECTOR_LEN)
+    poseidon16(pub_start + 4 * DIGEST_LEN, pub_start + 5 * DIGEST_LEN, pub_start + 6 * DIGEST_LEN)
     dot_product(pub_start + 88, pub_start + 88 + N, pub_start + 1000, N, BE)
     dot_product(pub_start + 88 + N, pub_start + 88 + N * (DIM + 1), pub_start + 1000 + DIM, N, EE)
     c: Mut = 0
@@ -154,15 +154,15 @@ fn test_zk_vm_helper(program_str: &str, (public_input, private_input): (&[F], &[
         &bytecode,
         (public_input, private_input),
         &vec![],
-        &default_whir_config(starting_log_inv_rate),
+        &default_whir_config(starting_log_inv_rate, false),
         false,
     );
     let proof_time = time.elapsed();
     verify_execution(
         &bytecode,
         public_input,
         proof.proof.clone(),
-        &default_whir_config(starting_log_inv_rate),
+        &default_whir_config(starting_log_inv_rate, false),
     )
     .unwrap();
     println!("{}", proof.exec_summary);
@@ -183,7 +183,7 @@ fn test_zk_vm_helper(program_str: &str, (public_input, private_input): (&[F], &[
                 &bytecode,
                 public_input,
                 fuzzed_proof,
-                &default_whir_config(starting_log_inv_rate),
+                &default_whir_config(starting_log_inv_rate, false),
             );
             assert!(verify_result.is_err(), "Fuzzing failed at index {}", i);
         }

crates/lean_vm/src/isa/hint.rs

@@ -90,7 +90,7 @@ impl CustomHint {
 
     pub fn n_args_range(&self) -> Range<usize> {
         match self {
-            Self::DecomposeBitsXMSS => 3..usize::MAX,
+            Self::DecomposeBitsXMSS => 5..6,
             Self::DecomposeBits => 4..5,
             Self::LessThan => 3..4,
         }
@@ -99,15 +99,18 @@ impl CustomHint {
     pub fn execute(&self, args: &[MemOrConstant], ctx: &mut HintExecutionContext<'_>) -> Result<(), RunnerError> {
         match self {
             Self::DecomposeBitsXMSS => {
-                let decomposed = &args[0];
-                let remaining = &args[1];
-                let to_decompose = &args[2..];
-                let mut memory_index_decomposed = decomposed.read_value(ctx.memory, ctx.fp)?.to_usize();
-                let mut memory_index_remaining = remaining.read_value(ctx.memory, ctx.fp)?.to_usize();
-                for value_source in to_decompose {
-                    let value = value_source.read_value(ctx.memory, ctx.fp)?.to_usize();
-                    for i in 0..12 {
-                        let value = F::from_usize((value >> (2 * i)) & 0b11);
+                let decomposed_ptr = args[0].read_value(ctx.memory, ctx.fp)?.to_usize();
+                let remaining_ptr = args[1].read_value(ctx.memory, ctx.fp)?.to_usize();
+                let to_decompose_ptr = args[2].read_value(ctx.memory, ctx.fp)?.to_usize();
+                let num_to_decompose = args[3].read_value(ctx.memory, ctx.fp)?.to_usize();
+                let w = args[4].read_value(ctx.memory, ctx.fp)?.to_usize();
+                assert!(w == 2 || w == 3 || w == 4);
+                let mut memory_index_decomposed = decomposed_ptr;
+                let mut memory_index_remaining = remaining_ptr;
+                for i in 0..num_to_decompose {
+                    let value = ctx.memory.get(to_decompose_ptr + i)?.to_usize();
+                    for i in 0..24 / w {
+                        let value = F::from_usize((value >> (w * i)) & ((1 << w) - 1));
                         ctx.memory.set(memory_index_decomposed, value)?;
                         memory_index_decomposed += 1;
                     }

crates/rec_aggregation/hashing.py

@@ -1,7 +1,7 @@
 from snark_lib import *
 
 DIM = 5  # extension degree
-VECTOR_LEN = 8
+DIGEST_LEN = 8
 
 WHIR_MERKLE_HEIGHTS = WHIR_MERKLE_HEIGHTS_PLACEHOLDER
 WHIR_NUM_QUERIES = WHIR_NUM_QUERIES_PLACEHOLDER
@@ -44,13 +44,13 @@ def batch_hash_slice_const(num_queries, all_data_to_hash, all_resulting_hashes,
 
 @inline
 def slice_hash(data, len):
-    states = Array((len - 1) * VECTOR_LEN)
-    poseidon16(data, data + VECTOR_LEN, states)
+    states = Array((len - 1) * DIGEST_LEN)
+    poseidon16(data, data + DIGEST_LEN, states)
     state_indexes = Array(len)
     state_indexes[0] = states
     for j in unroll(1, len - 1):
-        state_indexes[j] = state_indexes[j - 1] + VECTOR_LEN
-        poseidon16(state_indexes[j - 1], data + (j + 1) * VECTOR_LEN, state_indexes[j])
+        state_indexes[j] = state_indexes[j - 1] + DIGEST_LEN
+        poseidon16(state_indexes[j - 1], data + (j + 1) * DIGEST_LEN, state_indexes[j])
     return state_indexes[len - 2]
 
 
@@ -81,7 +81,7 @@ def merkle_verif_batch_const(n_paths: Const, merkle_paths, leaves_digests, leave
     for i in unroll(0, n_paths):
         merkle_verify(
             leaves_digests[i],
-            merkle_paths + (i * height) * VECTOR_LEN,
+            merkle_paths + (i * height) * DIGEST_LEN,
             leave_positions[i],
             root,
             height,
@@ -91,7 +91,7 @@ def merkle_verif_batch_const(n_paths: Const, merkle_paths, leaves_digests, leave
 
 
 def merkle_verify(leaf_digest, merkle_path, leaf_position_bits, root, height: Const):
-    states = Array(height * VECTOR_LEN)
+    states = Array(height * DIGEST_LEN)
 
     # First merkle round
     match leaf_position_bits[0]:
@@ -104,18 +104,18 @@ def merkle_verify(leaf_digest, merkle_path, leaf_position_bits, root, height: Co
     state_indexes = Array(height)
     state_indexes[0] = states
     for j in unroll(1, height):
-        state_indexes[j] = state_indexes[j - 1] + VECTOR_LEN
+        state_indexes[j] = state_indexes[j - 1] + DIGEST_LEN
         # Warning: this works only if leaf_position_bits[i] is known to be boolean:
         match leaf_position_bits[j]:
             case 0:
                 poseidon16(
                     state_indexes[j - 1],
-                    merkle_path + j * VECTOR_LEN,
+                    merkle_path + j * DIGEST_LEN,
                     state_indexes[j],
                 )
             case 1:
                 poseidon16(
-                    merkle_path + j * VECTOR_LEN,
+                    merkle_path + j * DIGEST_LEN,
                     state_indexes[j - 1],
                     state_indexes[j],
                 )

crates/rec_aggregation/recursion.py

@@ -61,7 +61,7 @@ def recursion(inner_public_memory_log_size, inner_public_memory, proof_transcrip
     fs: Mut = fs_new(proof_transcript)
 
     # table dims
-    debug_assert(N_TABLES + 1 < VECTOR_LEN)
+    debug_assert(N_TABLES + 1 < DIGEST_LEN)
     fs, mem_and_table_dims = fs_receive_chunks(fs, 1)
     for i in unroll(N_TABLES + 1, 8):
         assert mem_and_table_dims[i] == 0

crates/rec_aggregation/src/recursion.rs

@@ -14,16 +14,14 @@ use multilinear_toolkit::prelude::symbolic::{
 use multilinear_toolkit::prelude::*;
 use utils::{BYTECODE_TABLE_INDEX, Counter, MEMORY_TABLE_INDEX};
 
-const LOG_INV_RATE: usize = 2;
-
-pub fn run_recursion_benchmark(count: usize, tracing: bool) {
+pub fn run_recursion_benchmark(count: usize, log_inv_rate: usize, prox_gaps_conjecture: bool, tracing: bool) {
     let filepath = Path::new(env!("CARGO_MANIFEST_DIR"))
         .join("recursion.py")
         .to_str()
         .unwrap()
         .to_string();
 
-    let inner_whir_config = default_whir_config(LOG_INV_RATE);
+    let inner_whir_config = default_whir_config(log_inv_rate, prox_gaps_conjecture);
     let program_to_prove = r#"
 DIM = 5
 POSEIDON_OF_ZERO = POSEIDON_OF_ZERO_PLACEHOLDER
@@ -333,15 +331,15 @@ def main():
         &recursion_bytecode,
         (&outer_public_input, &outer_private_input),
         &vec![], // TODO precompute poseidons
-        &default_whir_config(LOG_INV_RATE),
+        &default_whir_config(log_inv_rate, prox_gaps_conjecture),
         false,
     );
     let proving_time = time.elapsed();
     verify_execution(
         &recursion_bytecode,
         &outer_public_input,
         recursion_proof.proof,
-        &default_whir_config(LOG_INV_RATE),
+        &default_whir_config(log_inv_rate, prox_gaps_conjecture),
     )
     .unwrap();
     println!(
@@ -589,5 +587,5 @@ fn display_all_air_evals_in_zk_dsl() {
 
 #[test]
 fn test_end2end_recursion() {
-    run_recursion_benchmark(1, false);
+    run_recursion_benchmark(1, 2, false, false);
 }