fix: roots mismatch between program and backend

Author: MarcosNicolau

examples/L2/Makefile

@@ -2,8 +2,5 @@
 
 SHELL := /bin/bash
 
-# Either sp1|risc0
-ZKVM=sp1
-
 run:
-	cargo run --release
+	@cd crates/l2 && cargo run --release

examples/L2/crates/l2/Cargo.toml

@@ -18,7 +18,6 @@ futures-util = "0.3"
 tokio = "1.44"
 alloy = { version = "0.14", features = ["full", "providers"] }
 
-
 [build-dependencies]
 sp1-build = { version = "4.1.3" } 
 

examples/L2/crates/l2/src/lib.rs

@@ -1,5 +1,7 @@
 use aligned::{send_proof_to_be_verified_on_aligned, wait_until_proof_is_aggregated};
+use alloy::hex;
 use db::{generate_random_transfers, DB};
+use primitive_types::U256;
 use sp1_state_transition_program::ProgramOutput;
 use zk::{prove_state_transition, PROGRAM_ELF};
 
@@ -14,30 +16,39 @@ pub async fn start_l2(
     wallet: aligned_sdk::core::types::Wallet<aligned_sdk::core::types::SigningKey>,
 ) {
     // 0. Load merkle tree file, if not created, create initial state
-    let mut db = DB::new("./db".to_string()).expect("create db");
+    let mut db = DB::new("./db".to_string());
 
     // 1. Create random transfers
-    let account_updates = generate_random_transfers(&db, 10);
+    let transfers = generate_random_transfers(&db, 10);
 
     // 2. Call zkvm and pass (MerkleTree, Updates to perform)
-    let (mut proof, _vk) = prove_state_transition(&mut db, account_updates.clone());
+    let (mut proof, _vk) = prove_state_transition(&db, transfers.clone());
     let ProgramOutput {
         initial_state_merkle_root,
         post_state_merkle_root,
     } = proof.public_values.read::<ProgramOutput>();
 
     // 3. If the proving went alright, update the db and verify that the merkle root matches
     assert!(db.commitment() == initial_state_merkle_root);
-    for update in account_updates {
-        db.user_states
-            .get_mut(&update.from)
-            .expect("User should exist")
-            .balance -= update.amount;
-
-        db.user_states
-            .get_mut(&update.to)
-            .expect("User should exist")
-            .balance += update.amount;
+    for transfer in transfers {
+        let mut user_from = db
+            .user_states
+            .get(&transfer.from)
+            .expect("User must exist in state")
+            .clone();
+
+        let mut user_to = db
+            .user_states
+            .get(&transfer.to)
+            .expect("User must exist in state")
+            .clone();
+
+        user_from.balance -= transfer.amount;
+        user_from.nonce += U256::one();
+        user_to.balance += transfer.amount;
+
+        db.user_states.insert(transfer.from, user_from);
+        db.user_states.insert(transfer.to, user_to);
     }
     assert!(db.commitment() == post_state_merkle_root);
 

examples/L2/crates/l2/src/zk.rs

@@ -1,3 +1,4 @@
+use sp1_sdk::Prover;
 use types::Transfer;
 
 use crate::db::DB;
@@ -6,7 +7,7 @@ pub const PROGRAM_ELF: &[u8] =
     include_bytes!("../zkvm_programs/sp1/elf/sp1_state_transition_program");
 
 pub fn prove_state_transition(
-    db: &mut DB,
+    db: &DB,
     transfers: Vec<Transfer>,
 ) -> (sp1_sdk::SP1ProofWithPublicValues, sp1_sdk::SP1VerifyingKey) {
     let mut stdin = sp1_sdk::SP1Stdin::new();

examples/L2/crates/l2/zkvm_programs/sp1/elf/sp1_state_transition_program

@@ -1,11 +1,11 @@
 use primitive_types::H160;
 use serde::{Deserialize, Serialize};
-use std::collections::HashMap;
+use std::collections::BTreeMap;
 use types::{Transfer, UserState};
 
 #[derive(Deserialize, Serialize)]
 pub struct ProgramInput {
-    pub user_states: HashMap<H160, UserState>,
+    pub user_states: BTreeMap<H160, UserState>,
     pub transfers: Vec<Transfer>,
 }
 

examples/L2/crates/l2/zkvm_programs/sp1/src/lib.rs

@@ -2,6 +2,7 @@
 sp1_zkvm::entrypoint!(main);
 
 use lambdaworks_crypto::merkle_tree::merkle::MerkleTree;
+use primitive_types::U256;
 use sp1_state_transition_program::{ProgramInput, ProgramOutput};
 use types::UserState;
 
@@ -26,15 +27,15 @@ pub fn main() {
             .clone();
 
         if user_from.balance >= transfer.amount {
-            user_from.balance -= transfer.amount
+            user_from.balance -= transfer.amount;
+            user_from.nonce += U256::one();
+            user_to.balance += transfer.amount;
         } else {
-            panic!("User does not have enough balance to perform the transfer");
+            panic!("User does not have enough balance to perform the transfer",);
         }
 
-        user_to.balance += transfer.amount;
-
-        input.user_states.insert(transfer.from, user_from.clone());
-        input.user_states.insert(transfer.to, user_to.clone());
+        input.user_states.insert(transfer.from, user_from);
+        input.user_states.insert(transfer.to, user_to);
     }
 
     let post_state: Vec<UserState> = input.user_states.clone().into_values().collect();

examples/L2/crates/l2/zkvm_programs/sp1/src/main.rs

@@ -3,7 +3,7 @@ use primitive_types::{H160, U256};
 use serde::{Deserialize, Serialize};
 use sha3::{Digest, Keccak256};
 
-#[derive(Clone, Default, Serialize, Deserialize)]
+#[derive(Clone, Debug, Default, Serialize, Deserialize)]
 pub struct UserState {
     pub address: H160,
     pub balance: U256,
@@ -36,7 +36,7 @@ impl IsMerkleTreeBackend for UserState {
     }
 }
 
-#[derive(Serialize, Deserialize, Clone)]
+#[derive(Serialize, Deserialize, Clone, Debug)]
 pub struct Transfer {
     pub from: H160,
     pub to: H160,