Initial SP1 support (WIP)

The proofs work, but it's not yet integrated with the main server.

kairos-sp1/Cargo.toml

@@ -0,0 +1,6 @@
+[workspace]
+members = [
+    "program",
+    "server"
+]
+resolver = "2"

kairos-sp1/program/Cargo.toml

@@ -0,0 +1,9 @@
+[package]
+version = "0.1.0"
+name = "kairos-batch-logic-sp1"
+edition = "2021"
+
+[dependencies]
+kairos-circuit-logic = { path = "../../kairos-prover/kairos-circuit-logic", features = ["serde", "borsh"], default-features = false }
+
+sp1-zkvm = "2.0.0"

kairos-sp1/program/src/main.rs

@@ -0,0 +1,18 @@
+#![cfg_attr(target_os = "zkvm", no_main)]
+
+#[cfg(target_os = "zkvm")]
+sp1_zkvm::entrypoint!(main);
+
+use kairos_circuit_logic::ProofInputs;
+
+pub fn main() {
+    let proof_inputs: ProofInputs = sp1_zkvm::io::read();
+
+    let output = proof_inputs
+        .run_batch_proof_logic()
+        .unwrap()
+        .borsh_serialize()
+        .unwrap();
+
+    sp1_zkvm::io::commit_slice(&output);
+}

kairos-sp1/rust-toolchain.toml

@@ -0,0 +1,3 @@
+[toolchain]
+channel = "1.79.0"
+components = ["llvm-tools", "rustc-dev"]

kairos-sp1/server/Cargo.toml

@@ -0,0 +1,37 @@
+[package]
+name = "kairos-sp1-server"
+version = "0.1.0"
+edition = "2021"
+
+[features]
+cuda = ["sp1-sdk/cuda"]
+
+[dependencies]
+sp1-sdk = "2.0.0"
+
+kairos-circuit-logic = { path = "../../kairos-prover/kairos-circuit-logic", features = ["serde", "borsh"], default-features = false }
+
+serde_json = { version = "1" }
+serde = { version = "1", default-features = false, features = ["derive"] }
+
+tokio = { version = "1", features = ["rt-multi-thread", "tracing", "macros"] }
+axum = { version = "0.7", features = ["tracing"] }
+axum-extra = { version = "0.9", features = [
+    "typed-routing",
+    "typed-header",
+    "json-deserializer",
+] }
+
+dotenvy = "0.15"
+tracing = "0.1"
+tracing-subscriber = { version = "0.3", features = ["env-filter"] }
+
+[build-dependencies]
+sp1-helper = "2.0.0"
+
+[dev-dependencies]
+kairos-trie = { git = "https://github.com/cspr-rad/kairos-trie", features = ["serde"] }
+kairos-circuit-logic = { path = "../../kairos-prover/kairos-circuit-logic", features = ["serde", "test-logic", "arbitrary" ], default-features = false }
+test-strategy = { version = "0.3" }
+proptest = { version = "1" }
+casper-types = "4.0"

kairos-sp1/server/build.rs

@@ -0,0 +1,5 @@
+use sp1_helper::build_program_with_args;
+
+fn main() {
+    build_program_with_args("../program", Default::default())
+}

kairos-sp1/server/src/main.rs

@@ -0,0 +1,262 @@
+use std::time::Instant;
+
+use axum::{http::StatusCode, Json};
+use axum_extra::routing::{RouterExt, TypedPath};
+use kairos_circuit_logic::{ProofInputs, ProofOutputs};
+
+use sp1_sdk::{SP1ProofWithPublicValues, SP1Stdin};
+// These constants represent the RISC-V ELF and the image ID generated by risc0-build.
+// The ELF is used for proving and the ID is used for verification.
+use tracing::level_filters::LevelFilter;
+use tracing_subscriber::{prelude::*, EnvFilter};
+
+pub const PROVE_BATCH_ELF: &[u8] = include_bytes!("../../elf/riscv32im-succinct-zkvm-elf");
+
+#[tokio::main]
+async fn main() {
+    let _ = dotenvy::dotenv();
+    // Initialize tracing. In order to view logs, run `RUST_LOG=info cargo run`
+    tracing_subscriber::registry()
+        .with(
+            EnvFilter::try_from_default_env()
+                .unwrap_or_else(|_| "info,kairos_sp1_server=info".into()),
+        )
+        .with(tracing_subscriber::fmt::layer())
+        .init();
+
+    let socket_addr = std::env::var("KAIROS_SP1_PROVER_SERVER_SOCKET_ADDR")
+        .expect("Failed to fetch environment variable KAIROS_SP1_PROVER_SERVER_SOCKET_ADDR");
+    let socket_addr = socket_addr
+        .parse::<std::net::SocketAddr>()
+        .expect("Failed to parse KAIROS_SP1_PROVER_SERVER_SOCKET_ADDR");
+
+    let app = axum::Router::new()
+        .typed_post(prove_batch_route)
+        .with_state(());
+
+    tracing::info!("starting http server on `{}`", socket_addr);
+    let listener = tokio::net::TcpListener::bind(socket_addr).await.unwrap();
+    tracing::info!("listening on `{}`", socket_addr);
+    axum::serve(listener, app).await.unwrap()
+}
+
+#[derive(TypedPath, Debug, Clone, Copy)]
+#[typed_path("/api/v1/prove/batch")]
+pub struct ProveBatch;
+
+pub async fn prove_batch_route(
+    _: ProveBatch,
+    proof_inputs: Json<ProofInputs>,
+) -> Result<Json<(ProofOutputs, SP1ProofWithPublicValues)>, (StatusCode, String)> {
+    let proof = tokio::task::spawn_blocking(move || prove_execution(proof_inputs.0))
+        .await
+        .map_err(|e| {
+            let e = format!("Error while joining proving task: {e}");
+            tracing::error!(e);
+            (StatusCode::INTERNAL_SERVER_ERROR, e)
+        })?;
+
+    let proof = proof.map_err(|e| {
+        let e = format!("Error while proving batch: {e}");
+        tracing::error!(e);
+        (StatusCode::INTERNAL_SERVER_ERROR, e)
+    })?;
+
+    Ok(Json(proof))
+}
+
+fn prove_execution(
+    proof_inputs: kairos_circuit_logic::ProofInputs,
+) -> Result<(ProofOutputs, SP1ProofWithPublicValues), String> {
+    let timestamp = Instant::now();
+
+    let client = sp1_sdk::ProverClient::local();
+
+    // Setup the inputs.
+    let mut stdin = SP1Stdin::new();
+    stdin.write(&proof_inputs);
+
+    let (pk, vk) = client.setup(PROVE_BATCH_ELF);
+
+    let proof = client
+        .prove(&pk, stdin)
+        .run()
+        .map_err(|e| format!("Failed to prove execution: {e}"))?;
+
+    tracing::info!("Proved batch: {}s", timestamp.elapsed().as_secs_f64());
+
+    let timestamp = Instant::now();
+
+    tracing::info!("Verified batch: {}s", timestamp.elapsed().as_secs_f64());
+
+    client
+        .verify(&proof, &vk)
+        .map_err(|e| format!("Failed to verify proof: {e}"))?;
+
+    // this panics if deserialization fails
+    let proof_outputs: &[u8] = proof.public_values.as_slice();
+    let proof_outputs = ProofOutputs::borsh_deserialize(proof_outputs)?;
+
+    Ok((proof_outputs, proof))
+}
+
+pub fn test_setup() {
+    let _ = tracing_subscriber::fmt()
+        .with_env_filter(
+            EnvFilter::builder()
+                .with_default_directive(LevelFilter::INFO.into())
+                .from_env_lossy(),
+        )
+        .try_init();
+
+    if cfg!(feature = "disable-dev-mode") {
+        std::env::set_var("RISC0_DEV_MODE", "0");
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::rc::Rc;
+
+    use casper_types::{bytesrepr::ToBytes, AsymmetricType};
+    use kairos_trie::{stored::memory_db::MemoryDb, TrieRoot};
+    use proptest::prelude::*;
+
+    use kairos_circuit_logic::{
+        account_trie::{test_logic::test_prove_batch, Account},
+        transactions::{
+            arbitrary::RandomBatches, KairosTransaction, L1Deposit, Signed, Transfer, Withdraw,
+        },
+    };
+
+    use crate::test_setup;
+
+    #[test]
+    fn test_prove_simple_batches() {
+        test_setup();
+
+        let alice_public_key = include_bytes!("../../../testdata/users/user-2/public_key_hex");
+        let alice_public_key = casper_types::PublicKey::from_hex(alice_public_key)
+            .expect("Failed to parse public key")
+            .to_bytes()
+            .expect("Failed to convert public key to bytes");
+
+        let bob_public_key = include_bytes!("../../../testdata/users/user-3/public_key_hex");
+        let bob_public_key = casper_types::PublicKey::from_hex(bob_public_key)
+            .expect("Failed to parse public key")
+            .to_bytes()
+            .expect("Failed to convert public key to bytes");
+
+        let batches = vec![
+            vec![
+                KairosTransaction::Deposit(L1Deposit {
+                    recipient: alice_public_key.clone(),
+                    amount: 10,
+                }),
+                KairosTransaction::Transfer(Signed {
+                    public_key: alice_public_key.clone(),
+                    transaction: Transfer {
+                        recipient: bob_public_key.clone(),
+                        amount: 5,
+                    },
+                    nonce: 0,
+                }),
+                KairosTransaction::Withdraw(Signed {
+                    public_key: alice_public_key.clone(),
+                    transaction: Withdraw { amount: 5 },
+                    nonce: 1,
+                }),
+            ],
+            vec![
+                KairosTransaction::Transfer(Signed {
+                    public_key: bob_public_key.clone(),
+                    transaction: Transfer {
+                        recipient: alice_public_key.clone(),
+                        amount: 2,
+                    },
+                    nonce: 0,
+                }),
+                KairosTransaction::Withdraw(Signed {
+                    public_key: bob_public_key.clone(),
+                    transaction: Withdraw { amount: 3 },
+                    nonce: 1,
+                }),
+                KairosTransaction::Withdraw(Signed {
+                    public_key: alice_public_key.clone(),
+                    transaction: Withdraw { amount: 2 },
+                    nonce: 2,
+                }),
+            ],
+        ];
+
+        test_prove_batch(
+            TrieRoot::Empty,
+            Rc::new(MemoryDb::<Account>::empty()),
+            batches,
+            |(batch_num, proof_inputs)| {
+                let (proof_outputs, receipt) =
+                    crate::prove_execution(proof_inputs).expect("Failed to prove execution");
+
+                if cfg!(feature = "write-test-proofs") {
+                    let proof_file = std::fs::File::create(format!(
+                        "test_prove_simple_batches_{batch_num}.json"
+                    ))
+                    .expect("Failed to create proof file");
+
+                    serde_json::to_writer(proof_file, &receipt)
+                        .expect("Failed to write proof file");
+                };
+
+                Ok(proof_outputs)
+            },
+        );
+    }
+
+    #[test_strategy::proptest(ProptestConfig::default(), cases = if cfg!(feature = "disable-dev-mode") { 2 } else { 40 })]
+    fn proptest_prove_batches(
+        #[any(batch_size = 5..=10, batch_count = 2..=4, initial_l2_accounts = 10_000..=100_000)]
+        args: RandomBatches,
+    ) {
+        test_setup();
+        let batches = args.filter_success();
+
+        proptest::prop_assume!(batches.len() >= 2);
+
+        test_prove_batch(
+            args.initial_trie,
+            args.trie_db,
+            batches,
+            |(_, proof_inputs)| {
+                tracing::info!(
+                    "Proving batch of size: {}, over trie of {} accounts.",
+                    proof_inputs.transactions.len(),
+                    args.initial_state.l2.len()
+                );
+
+                let (proof_outputs, receipt) =
+                    crate::prove_execution(proof_inputs).expect("Failed to prove execution");
+
+                // if cfg!(feature = "write-test-proofs") {
+                //     use std::hash::{DefaultHasher, Hasher};
+
+                //     let mut hasher = DefaultHasher::new();
+                //     hasher.write(&receipt.journal.bytes);
+                //     let journal_hash = hasher.finish();
+
+                //     let proof_path = format!(
+                //         "proptest_prove_batches-proof-journal-{:x}.json",
+                //         journal_hash
+                //     );
+
+                //     let proof_file =
+                //         std::fs::File::create(proof_path).expect("Failed to create proof file");
+
+                //     serde_json::to_writer(proof_file, &receipt)
+                //         .expect("Failed to write proof file");
+                // }
+
+                Ok(proof_outputs)
+            },
+        )
+    }
+}