feat: add btreemap example (#836)

Add BTreeMap workload with complex memory access patterns (insertions,
deletions, range scans) for profiling.

Author: zouguangxian

Cargo.lock

@@ -28,6 +28,8 @@ members = [
     "jolt-sdk",
     "jolt-sdk/macros",
     "jolt-verifier",
+    "examples/btreemap/host",
+    "examples/btreemap/guest",
     "examples/collatz",
     "examples/collatz/guest",
     "examples/fibonacci",
@@ -104,3 +106,6 @@ ark-bn254 = { git = "https://github.com/a16z/arkworks-algebra", branch = "dev/tw
 
 [workspace.metadata.cargo-machete]
 ignored = ["jolt-sdk"]
+
+[workspace.dependencies]
+spinners = { version= "4.1.1", default-features = false }

Cargo.toml

@@ -0,0 +1,10 @@
+[package]
+name = "btreemap-guest"
+version = "0.1.0"
+edition = "2021"
+
+[features]
+guest = []
+
+[dependencies]
+jolt = { package = "jolt-sdk", path = "../../../jolt-sdk" }

examples/btreemap/guest/Cargo.toml

@@ -0,0 +1,63 @@
+#![cfg_attr(feature = "guest", no_std)]
+
+extern crate alloc;
+
+use core::option::Option::None;
+
+/// Fast hash function implementation using wyhash64 algorithm.
+///
+/// This function provides a high-quality, fast hash suitable for hashmap operations.
+/// It uses the wyhash64 algorithm which offers good distribution properties and
+/// performance characteristics.
+fn wyhash64(mut x: u64) -> u64 {
+    x ^= x >> 32;
+    x = x.wrapping_mul(0xd6e8feb86659fd93);
+    x ^= x >> 32;
+    x = x.wrapping_mul(0xd6e8feb86659fd93);
+    x
+}
+
+#[jolt::provable(stack_size = 10000, memory_size = 10000000)]
+pub fn btreemap(n: u32) -> u128 {
+    use alloc::collections::BTreeMap;
+
+    let mut map = BTreeMap::new();
+    let mut inserted_keys = alloc::vec::Vec::with_capacity(n as usize);
+
+    // Phase 1: Insert N entries with high-entropy keys
+    for i in 0..n {
+        let key = wyhash64(i as u64); // Use u64 directly, not usize
+        inserted_keys.push(key);
+        map.insert(key, i as u64);
+    }
+
+    // Phase 2: Delete 25% of the inserted keys to trigger rebalancing
+    let delete_count = n / 4;
+    for i in 0..delete_count {
+        let key = inserted_keys[i as usize];
+        map.remove(&key);
+    }
+
+    // Phase 3: Insert N/2 new entries with new hashed keys
+    for i in 0..(n / 2) {
+        let key = wyhash64((i + n * 2) as u64); // Non-overlapping seed
+        map.insert(key, (i + n) as u64);
+    }
+
+    // Phase 4: Range scan over middle 25% of key space
+    let mut range_sum = 0u64;
+    if let Some((&min_key, _)) = map.first_key_value() {
+        if let Some((&max_key, _)) = map.last_key_value() {
+            let range_size = (max_key - min_key) / 4;
+            let start = min_key + range_size;
+            let end = start + range_size;
+
+            for (_, value) in map.range(start..end) {
+                range_sum = range_sum.wrapping_add(*value);
+            }
+        }
+    }
+
+    // Combine size and range sum into a single return value
+    (map.len() as u128).wrapping_add(range_sum as u128)
+}

examples/btreemap/guest/src/lib.rs

@@ -0,0 +1,5 @@
+#![cfg_attr(feature = "guest", no_std)]
+#![no_main]
+
+#[allow(unused_imports)]
+use btreemap_guest::*;

examples/btreemap/guest/src/main.rs

@@ -0,0 +1,9 @@
+[package]
+name = "btreemap"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+jolt-sdk = { path = "../../../jolt-sdk", features = ["host"] }
+guest = { package = "btreemap-guest", path = "../guest" }
+spinners = { workspace = true }

examples/btreemap/host/Cargo.toml

@@ -0,0 +1,46 @@
+use spinners::{Spinner, Spinners};
+
+macro_rules! step {
+    ($msg:expr, $action:expr) => {{
+        let mut sp = Spinner::new(Spinners::Dots9, $msg.to_string());
+        let result = $action;
+        sp.stop_with_message(format!("✓ {}", $msg));
+        result
+    }};
+}
+
+pub fn btreemap() {
+    let target_dir = "/tmp/jolt-guest-targets";
+
+    let mut program = step!("Compiling guest code", {
+        guest::compile_btreemap(target_dir)
+    });
+
+    let prover_preprocessing = step!("Preprocessing prover", {
+        guest::preprocess_prover_btreemap(&mut program)
+    });
+
+    let verifier_preprocessing = step!("Preprocessing verifier", {
+        guest::verifier_preprocessing_from_prover_btreemap(&prover_preprocessing)
+    });
+
+    let prove = step!("Building prover", {
+        guest::build_prover_btreemap(program, prover_preprocessing)
+    });
+
+    let verify = step!("Building verifier", {
+        guest::build_verifier_btreemap(verifier_preprocessing)
+    });
+
+    let n = 50;
+    let (output, proof, io_device) = step!("Proving", { prove(n) });
+    assert!(output >= 1);
+
+    let is_valid = step!("Verifying", { verify(n, output, io_device.panic,proof) });
+    assert!(is_valid);
+}
+
+fn main() {
+    println!("BTreeMap");
+    btreemap();
+}

examples/btreemap/host/src/main.rs

@@ -12,6 +12,7 @@ use rand_distr::{Distribution, Zipf};
 
 #[derive(Debug, Copy, Clone, clap::ValueEnum)]
 pub enum BenchType {
+    Btreemap,
     Fibonacci,
     Sha2,
     Sha3,
@@ -22,6 +23,7 @@ pub enum BenchType {
 
 pub fn benchmarks(bench_type: BenchType) -> Vec<(tracing::Span, Box<dyn FnOnce()>)> {
     match bench_type {
+        BenchType::Btreemap => btreemap(),
         BenchType::Sha2 => sha2(),
         BenchType::Sha3 => sha3(),
         BenchType::Sha2Chain => sha2_chain(),
@@ -115,6 +117,10 @@ fn sha3() -> Vec<(tracing::Span, Box<dyn FnOnce()>)> {
     prove_example("sha3-guest", postcard::to_stdvec(&vec![5u8; 2048]).unwrap())
 }
 
+fn btreemap() -> Vec<(tracing::Span, Box<dyn FnOnce()>)> {
+    prove_example("btreemap-guest", postcard::to_stdvec(&50u32).unwrap())
+}
+
 fn sha2_chain() -> Vec<(tracing::Span, Box<dyn FnOnce()>)> {
     let mut inputs = vec![];
     inputs.append(&mut postcard::to_stdvec(&[5u8; 32]).unwrap());