From a0b9b357dafe92f25c283628514d1e333269598d Mon Sep 17 00:00:00 2001
From: Giacomo Pasini <g.pasini98@gmail.com>
Date: Mon, 25 Nov 2024 18:45:00 +0100
Subject: [PATCH] remove limitation of single ptx origin

---
 emmarin/cl/cl/src/cl/bundle.rs                | 36 +++++++++-
 .../cl/ledger/src/{bundle.rs => balance.rs}   | 16 ++---
 emmarin/cl/ledger/src/ledger.rs               | 18 ++---
 emmarin/cl/ledger/src/lib.rs                  |  2 +-
 emmarin/cl/ledger/src/partial_tx.rs           | 28 ++------
 emmarin/cl/ledger/src/zone_update.rs          | 33 ++++-----
 emmarin/cl/ledger/tests/simple_transfer.rs    | 12 ++--
 .../src/{bundle.rs => balance.rs}             |  4 +-
 .../cl/ledger_proof_statements/src/ledger.rs  | 12 +++-
 emmarin/cl/ledger_proof_statements/src/lib.rs |  2 +-
 emmarin/cl/ledger_proof_statements/src/ptx.rs |  4 +-
 .../ledger_validity_proof/ledger/src/main.rs  | 67 +++++++++----------
 emmarin/cl/risc0_proofs/Cargo.toml            |  2 +-
 .../{bundle => balance}/Cargo.toml            |  2 +-
 .../{bundle => balance}/src/main.rs           | 10 +--
 emmarin/cl/risc0_proofs/ptx/src/main.rs       | 14 ++--
 16 files changed, 145 insertions(+), 117 deletions(-)
 rename emmarin/cl/ledger/src/{bundle.rs => balance.rs} (79%)
 rename emmarin/cl/ledger_proof_statements/src/{bundle.rs => balance.rs} (83%)
 rename emmarin/cl/risc0_proofs/{bundle => balance}/Cargo.toml (97%)
 rename emmarin/cl/risc0_proofs/{bundle => balance}/src/main.rs (57%)

diff --git a/emmarin/cl/cl/src/cl/bundle.rs b/emmarin/cl/cl/src/cl/bundle.rs
index c9166cf..b43e734 100644
--- a/emmarin/cl/cl/src/cl/bundle.rs
+++ b/emmarin/cl/cl/src/cl/bundle.rs
@@ -1,16 +1,50 @@
 use serde::{Deserialize, Serialize};
 
-use crate::cl::{partial_tx::PartialTx, BalanceWitness, PartialTxWitness};
+use crate::{
+    cl::{partial_tx::PartialTx, BalanceWitness, PartialTxWitness},
+    zone_layer::notes::ZoneId,
+};
+use sha2::{Digest, Sha256};
+use std::collections::HashSet;
 
 /// The transaction bundle is a collection of partial transactions.
 /// The goal in bundling transactions is to produce a set of partial transactions
 /// that balance each other.
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
+pub struct BundleId(pub [u8; 32]);
+
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
 pub struct Bundle {
     pub partials: Vec<PartialTx>,
 }
 
+impl Bundle {
+    pub fn zones(&self) -> HashSet<ZoneId> {
+        self.partials
+            .iter()
+            .flat_map(|ptx| {
+                ptx.inputs
+                    .iter()
+                    .map(|i| i.zone_id)
+                    .chain(ptx.outputs.iter().map(|o| o.zone_id))
+            })
+            .collect()
+    }
+
+    ///
+    pub fn id(&self) -> BundleId {
+        // TODO: change to merkle root
+        let mut hasher = Sha256::new();
+        hasher.update(b"NOMOS_CL_BUNDLE_ID");
+        for ptx in &self.partials {
+            hasher.update(&ptx.root().0);
+        }
+
+        BundleId(hasher.finalize().into())
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 pub struct BundleWitness {
     pub partials: Vec<PartialTxWitness>,
diff --git a/emmarin/cl/ledger/src/bundle.rs b/emmarin/cl/ledger/src/balance.rs
similarity index 79%
rename from emmarin/cl/ledger/src/bundle.rs
rename to emmarin/cl/ledger/src/balance.rs
index e8d4e8e..453ae8c 100644
--- a/emmarin/cl/ledger/src/bundle.rs
+++ b/emmarin/cl/ledger/src/balance.rs
@@ -1,19 +1,19 @@
 use crate::error::{Error, Result};
 use cl::cl::BundleWitness;
-use ledger_proof_statements::bundle::{BundlePrivate, BundlePublic};
+use ledger_proof_statements::balance::{BalancePrivate, BalancePublic};
 
 #[derive(Debug, Clone)]
-pub struct ProvedBundle {
-    pub bundle: BundlePublic,
+pub struct ProvedBalance {
+    pub bundle: BalancePublic,
     pub risc0_receipt: risc0_zkvm::Receipt,
 }
 
-impl ProvedBundle {
+impl ProvedBalance {
     pub fn prove(bundle_witness: &BundleWitness) -> Result<Self> {
         // need to show that bundle is balanced.
         // i.e. the sum of ptx balances is 0
 
-        let bundle_private = BundlePrivate {
+        let bundle_private = BalancePrivate {
             balances: bundle_witness
                 .partials
                 .iter()
@@ -33,7 +33,7 @@ impl ProvedBundle {
 
         let opts = risc0_zkvm::ProverOpts::succinct();
         let prove_info = prover
-            .prove_with_opts(env, nomos_cl_risc0_proofs::BUNDLE_ELF, &opts)
+            .prove_with_opts(env, nomos_cl_risc0_proofs::BALANCE_ELF, &opts)
             .map_err(|_| Error::Risc0ProofFailed)?;
 
         println!(
@@ -50,7 +50,7 @@ impl ProvedBundle {
         })
     }
 
-    pub fn public(&self) -> Result<ledger_proof_statements::bundle::BundlePublic> {
+    pub fn public(&self) -> Result<ledger_proof_statements::balance::BalancePublic> {
         Ok(self.risc0_receipt.journal.decode()?)
     }
 
@@ -62,7 +62,7 @@ impl ProvedBundle {
         // Vec::from_iter(self.bundle.partials.iter().map(|ptx| ptx.balance)) == bundle_public.balances
         // &&
         self.risc0_receipt
-            .verify(nomos_cl_risc0_proofs::BUNDLE_ID)
+            .verify(nomos_cl_risc0_proofs::BALANCE_ID)
             .is_ok()
     }
 }
diff --git a/emmarin/cl/ledger/src/ledger.rs b/emmarin/cl/ledger/src/ledger.rs
index 7920972..d7cdf68 100644
--- a/emmarin/cl/ledger/src/ledger.rs
+++ b/emmarin/cl/ledger/src/ledger.rs
@@ -4,7 +4,7 @@ use ledger_proof_statements::{
 };
 
 use crate::{
-    bundle::ProvedBundle,
+    balance::ProvedBalance,
     constraint::ConstraintProof,
     error::{Error, Result},
     partial_tx::ProvedPartialTx,
@@ -19,14 +19,14 @@ pub struct ProvedLedgerTransition {
 
 // TODO: find a better name
 #[derive(Debug, Clone)]
-pub struct ProvedZoneTx {
-    pub bundle: ProvedBundle,
+pub struct ProvedBundle {
+    pub bundle: ProvedBalance,
     pub ptxs: Vec<ProvedPartialTx>,
 }
 
-impl ProvedZoneTx {
+impl ProvedBundle {
     fn to_public(&self) -> Vec<PtxPublic> {
-        self.ptxs.iter().map(|p| p.public().unwrap()).collect()
+        self.ptxs.iter().map(|p| p.public.clone()).collect()
     }
 
     fn proofs(&self) -> Vec<risc0_zkvm::Receipt> {
@@ -40,19 +40,19 @@ impl ProvedLedgerTransition {
     pub fn prove(
         ledger: LedgerWitness,
         zone_id: ZoneId,
-        ptxs: Vec<ProvedZoneTx>,
+        bundles: Vec<ProvedBundle>,
         constraints: Vec<ConstraintProof>,
     ) -> Result<Self> {
         let witness = LedgerProofPrivate {
-            bundles: ptxs.iter().map(|p| p.to_public()).collect(),
+            bundles: bundles.iter().map(|p| p.to_public()).collect(),
             ledger,
             id: zone_id,
         };
 
         let mut env = risc0_zkvm::ExecutorEnv::builder();
 
-        for ptx in ptxs {
-            for proof in ptx.proofs() {
+        for bundle in bundles {
+            for proof in bundle.proofs() {
                 env.add_assumption(proof);
             }
         }
diff --git a/emmarin/cl/ledger/src/lib.rs b/emmarin/cl/ledger/src/lib.rs
index ad327b3..d42f46c 100644
--- a/emmarin/cl/ledger/src/lib.rs
+++ b/emmarin/cl/ledger/src/lib.rs
@@ -1,4 +1,4 @@
-pub mod bundle;
+pub mod balance;
 pub mod constraint;
 pub mod error;
 pub mod ledger;
diff --git a/emmarin/cl/ledger/src/partial_tx.rs b/emmarin/cl/ledger/src/partial_tx.rs
index d49891f..d57bea5 100644
--- a/emmarin/cl/ledger/src/partial_tx.rs
+++ b/emmarin/cl/ledger/src/partial_tx.rs
@@ -1,13 +1,12 @@
 use ledger_proof_statements::ptx::{PtxPrivate, PtxPublic};
 
 use crate::error::{Error, Result};
-use cl::cl::{merkle, PartialTx, PartialTxWitness};
+use cl::cl::{merkle, PartialTxWitness};
 use cl::zone_layer::notes::ZoneId;
 
 #[derive(Debug, Clone)]
 pub struct ProvedPartialTx {
-    pub ptx: PartialTx,
-    pub cm_root: [u8; 32],
+    pub public: PtxPublic,
     pub risc0_receipt: risc0_zkvm::Receipt,
 }
 
@@ -15,15 +14,14 @@ impl ProvedPartialTx {
     pub fn prove(
         ptx_witness: PartialTxWitness,
         input_cm_paths: Vec<Vec<merkle::PathNode>>,
-        cm_root: [u8; 32],
+        cm_roots: Vec<[u8; 32]>,
         from: Vec<ZoneId>,
         to: Vec<ZoneId>,
     ) -> Result<ProvedPartialTx> {
-        let ptx = ptx_witness.commit(&from, &to);
         let ptx_private = PtxPrivate {
             ptx: ptx_witness,
             input_cm_paths,
-            cm_root,
+            cm_roots: cm_roots.clone(),
             from,
             to,
         };
@@ -53,28 +51,12 @@ impl ProvedPartialTx {
         );
 
         Ok(Self {
-            ptx,
-            cm_root,
+            public: prove_info.receipt.journal.decode()?,
             risc0_receipt: prove_info.receipt,
         })
     }
 
-    pub fn public(&self) -> Result<PtxPublic> {
-        Ok(self.risc0_receipt.journal.decode()?)
-    }
-
     pub fn verify(&self) -> bool {
-        let Ok(proved_ptx_inputs) = self.public() else {
-            return false;
-        };
-        let expected_ptx_inputs = PtxPublic {
-            ptx: self.ptx.clone(),
-            cm_root: self.cm_root,
-        };
-        if expected_ptx_inputs != proved_ptx_inputs {
-            return false;
-        }
-
         self.risc0_receipt
             .verify(nomos_cl_risc0_proofs::PTX_ID)
             .is_ok()
diff --git a/emmarin/cl/ledger/src/zone_update.rs b/emmarin/cl/ledger/src/zone_update.rs
index fb145cd..2bf7baa 100644
--- a/emmarin/cl/ledger/src/zone_update.rs
+++ b/emmarin/cl/ledger/src/zone_update.rs
@@ -1,7 +1,7 @@
 pub use crate::error::{Error, Result};
 use crate::{ledger::ProvedLedgerTransition, stf::StfProof};
 use cl::zone_layer::tx::UpdateBundle;
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 
 pub struct ProvedUpdateBundle {
     pub bundle: UpdateBundle,
@@ -11,28 +11,31 @@ pub struct ProvedUpdateBundle {
 
 impl ProvedUpdateBundle {
     pub fn verify(&self) -> bool {
-        let mut consumed_commitments = HashSet::new();
-        let mut produced_commitments = HashSet::new();
+        let mut expected_zones = HashMap::new();
+        let mut actual_zones = HashMap::new();
         for proof in &self.ledger_proofs {
             if !proof.verify() {
                 return false;
             }
 
-            for comm in &proof.public.cross_out {
-                if produced_commitments.insert(comm) {
-                    // already in?
-                }
-            }
-            for comm in &proof.public.cross_in {
-                if consumed_commitments.insert(comm) {
-                    // already in?
-                }
+            for bundle in &proof.public.cross_bundles {
+                expected_zones.insert(bundle.id, HashSet::from_iter(bundle.zones.clone()));
+                actual_zones
+                    .entry(bundle.id)
+                    .or_insert_with(|| HashSet::new())
+                    .insert(proof.public.id);
             }
         }
 
-        // check that cross zone transactions match
-        if consumed_commitments != produced_commitments {
-            return false;
+        println!("{:?} | {:?}", expected_zones, actual_zones);
+        for (bundle, expected) in expected_zones.iter() {
+            if let Some(actual) = actual_zones.get(bundle) {
+                if actual != expected {
+                    panic!("{:?} | {:?}", actual, expected);
+                }
+            } else {
+                panic!();
+            }
         }
 
         for ((update, stf_proof), ledger_proof) in self
diff --git a/emmarin/cl/ledger/tests/simple_transfer.rs b/emmarin/cl/ledger/tests/simple_transfer.rs
index 8918fb6..b9858c3 100644
--- a/emmarin/cl/ledger/tests/simple_transfer.rs
+++ b/emmarin/cl/ledger/tests/simple_transfer.rs
@@ -10,9 +10,9 @@ use cl::{
     },
 };
 use ledger::{
-    bundle::ProvedBundle,
+    balance::ProvedBalance,
     constraint::ConstraintProof,
-    ledger::{ProvedLedgerTransition, ProvedZoneTx},
+    ledger::{ProvedBundle, ProvedLedgerTransition},
     partial_tx::ProvedPartialTx,
     stf::StfProof,
     zone_update::ProvedUpdateBundle,
@@ -73,25 +73,25 @@ fn cross_transfer_transition(
     let proved_ptx = ProvedPartialTx::prove(
         ptx_witness.clone(),
         vec![ledger_a.cm_path(&input.note_commitment(&zone_a)).unwrap()],
-        ledger_a.cm_root(),
+        vec![ledger_a.cm_root()],
         vec![zone_a],
         vec![zone_b, zone_a],
     )
     .unwrap();
 
-    let bundle = ProvedBundle::prove(&BundleWitness {
+    let bundle = ProvedBalance::prove(&BundleWitness {
         partials: vec![ptx_witness],
     })
     .unwrap();
 
-    let zone_tx = ProvedZoneTx {
+    let zone_tx = ProvedBundle {
         ptxs: vec![proved_ptx.clone()],
         bundle,
     };
 
     // Prove the constraints for alices input (she uses the no-op constraint)
     let constraint_proof =
-        ConstraintProof::prove_nop(input.nullifier(&zone_a), proved_ptx.ptx.root());
+        ConstraintProof::prove_nop(input.nullifier(&zone_a), proved_ptx.public.ptx.root());
 
     let ledger_a_transition = ProvedLedgerTransition::prove(
         ledger_a,
diff --git a/emmarin/cl/ledger_proof_statements/src/bundle.rs b/emmarin/cl/ledger_proof_statements/src/balance.rs
similarity index 83%
rename from emmarin/cl/ledger_proof_statements/src/bundle.rs
rename to emmarin/cl/ledger_proof_statements/src/balance.rs
index f952a4c..509b1b5 100644
--- a/emmarin/cl/ledger_proof_statements/src/bundle.rs
+++ b/emmarin/cl/ledger_proof_statements/src/balance.rs
@@ -2,11 +2,11 @@ use cl::cl::{Balance, BalanceWitness};
 use serde::{Deserialize, Serialize};
 
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
-pub struct BundlePublic {
+pub struct BalancePublic {
     pub balances: Vec<Balance>,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
-pub struct BundlePrivate {
+pub struct BalancePrivate {
     pub balances: Vec<BalanceWitness>,
 }
diff --git a/emmarin/cl/ledger_proof_statements/src/ledger.rs b/emmarin/cl/ledger_proof_statements/src/ledger.rs
index 824246e..a7feda6 100644
--- a/emmarin/cl/ledger_proof_statements/src/ledger.rs
+++ b/emmarin/cl/ledger_proof_statements/src/ledger.rs
@@ -1,5 +1,5 @@
 use crate::ptx::PtxPublic;
-use cl::cl::Output;
+use cl::cl::{bundle::BundleId, Output};
 use cl::zone_layer::{
     ledger::{Ledger, LedgerWitness},
     notes::ZoneId,
@@ -11,8 +11,8 @@ pub struct LedgerProofPublic {
     pub old_ledger: Ledger,
     pub ledger: Ledger,
     pub id: ZoneId,
-    pub cross_in: Vec<Output>,
-    pub cross_out: Vec<Output>,
+    pub cross_bundles: Vec<CrossZoneBundle>,
+    pub outputs: Vec<Output>,
 }
 
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
@@ -21,3 +21,9 @@ pub struct LedgerProofPrivate {
     pub id: ZoneId,
     pub bundles: Vec<Vec<PtxPublic>>,
 }
+
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub struct CrossZoneBundle {
+    pub id: BundleId,
+    pub zones: Vec<ZoneId>,
+}
diff --git a/emmarin/cl/ledger_proof_statements/src/lib.rs b/emmarin/cl/ledger_proof_statements/src/lib.rs
index 5d907dc..2000d4b 100644
--- a/emmarin/cl/ledger_proof_statements/src/lib.rs
+++ b/emmarin/cl/ledger_proof_statements/src/lib.rs
@@ -1,4 +1,4 @@
-pub mod bundle;
+pub mod balance;
 pub mod constraint;
 pub mod ledger;
 pub mod ptx;
diff --git a/emmarin/cl/ledger_proof_statements/src/ptx.rs b/emmarin/cl/ledger_proof_statements/src/ptx.rs
index fe256cc..7f53548 100644
--- a/emmarin/cl/ledger_proof_statements/src/ptx.rs
+++ b/emmarin/cl/ledger_proof_statements/src/ptx.rs
@@ -7,14 +7,14 @@ use serde::{Deserialize, Serialize};
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 pub struct PtxPublic {
     pub ptx: PartialTx,
-    pub cm_root: [u8; 32],
+    pub cm_roots: Vec<[u8; 32]>,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
 pub struct PtxPrivate {
     pub ptx: PartialTxWitness,
     pub input_cm_paths: Vec<Vec<merkle::PathNode>>,
-    pub cm_root: [u8; 32],
+    pub cm_roots: Vec<[u8; 32]>,
     pub from: Vec<ZoneId>,
     pub to: Vec<ZoneId>,
 }
diff --git a/emmarin/cl/ledger_validity_proof/ledger/src/main.rs b/emmarin/cl/ledger_validity_proof/ledger/src/main.rs
index b71212c..d4b644a 100644
--- a/emmarin/cl/ledger_validity_proof/ledger/src/main.rs
+++ b/emmarin/cl/ledger_validity_proof/ledger/src/main.rs
@@ -1,11 +1,11 @@
 use cl::{
-    cl::Output,
+    cl::{Bundle, Output},
     zone_layer::{ledger::LedgerWitness, notes::ZoneId},
 };
 use ledger_proof_statements::{
-    bundle::BundlePublic,
+    balance::BalancePublic,
     constraint::ConstraintPublic,
-    ledger::{LedgerProofPrivate, LedgerProofPublic},
+    ledger::{CrossZoneBundle, LedgerProofPrivate, LedgerProofPublic},
     ptx::PtxPublic,
 };
 use risc0_zkvm::{guest::env, serde};
@@ -18,29 +18,37 @@ fn main() {
     } = env::read();
 
     let old_ledger = ledger.commit();
+    let mut cross_bundles = vec![];
+    let mut outputs = vec![];
 
     let cm_root = ledger.cm_root();
 
-    let mut cross_in = vec![];
-    let mut cross_out = vec![];
-
     for bundle in bundles {
-        let bundle_public = BundlePublic {
+        let balance_public = BalancePublic {
             balances: bundle.iter().map(|ptx| ptx.ptx.balance).collect::<Vec<_>>(),
         };
         // verify bundle is balanced
         env::verify(
-            nomos_cl_risc0_proofs::BUNDLE_ID,
-            &serde::to_vec(&bundle_public).unwrap(),
+            nomos_cl_risc0_proofs::BALANCE_ID,
+            &serde::to_vec(&balance_public).unwrap(),
         )
         .unwrap();
 
         for ptx in &bundle {
-            let (new_ledger, consumed_commitments, produced_commitments) =
-                process_ptx(ledger, ptx, id, cm_root);
-            cross_in.extend(consumed_commitments);
-            cross_out.extend(produced_commitments);
+            let (new_ledger, ptx_outputs) = process_ptx(ledger, ptx, id, cm_root);
             ledger = new_ledger;
+            outputs.extend(ptx_outputs);
+        }
+
+        let bundle = Bundle {
+            partials: bundle.into_iter().map(|ptx| ptx.ptx).collect(),
+        };
+        let zones = bundle.zones();
+        if zones.len() > 1 {
+            cross_bundles.push(CrossZoneBundle {
+                id: bundle.id(),
+                zones: zones.into_iter().collect(),
+            });
         }
     }
 
@@ -48,8 +56,8 @@ fn main() {
         old_ledger,
         ledger: ledger.commit(),
         id,
-        cross_in,
-        cross_out,
+        cross_bundles,
+        outputs,
     });
 }
 
@@ -58,21 +66,18 @@ fn process_ptx(
     ptx: &PtxPublic,
     zone_id: ZoneId,
     cm_root: [u8; 32],
-) -> (LedgerWitness, Vec<Output>, Vec<Output>) {
-    let mut cross_in = vec![];
-    let mut cross_out = vec![];
-
+) -> (LedgerWitness, Vec<Output>) {
+    // always verify the ptx to ensure outputs were derived with the correct zone id
     env::verify(nomos_cl_risc0_proofs::PTX_ID, &serde::to_vec(&ptx).unwrap()).unwrap();
 
-    let ptx_cm_root = ptx.cm_root;
+    let cm_roots = &ptx.cm_roots;
     let ptx = &ptx.ptx;
 
-    // TODO: accept inputs from multiple zones
-    let check_inputs = ptx.inputs.iter().all(|input| input.zone_id == zone_id);
+    let mut outputs = vec![];
 
-    if check_inputs {
-        assert_eq!(ptx_cm_root, cm_root);
-        for input in &ptx.inputs {
+    for (input, input_cm_root) in ptx.inputs.iter().zip(cm_roots) {
+        if input.zone_id == zone_id {
+            assert_eq!(*input_cm_root, cm_root);
             assert!(!ledger.nullifiers.contains(&input.nullifier));
             ledger.nullifiers.push(input.nullifier);
 
@@ -91,17 +96,9 @@ fn process_ptx(
     for output in &ptx.outputs {
         if output.zone_id == zone_id {
             ledger.commitments.push(output.note_comm);
-            // if this output was not originating from this zone, it is a cross zone transaction
-            if !check_inputs {
-                cross_in.push(*output);
-            }
-        } else {
-            // if this output is not going to this zone but originated from this zone, it is a cross zone transaction
-            if check_inputs {
-                cross_out.push(*output);
-            }
+            outputs.push(*output);
         }
     }
 
-    (ledger, cross_in, cross_out)
+    (ledger, outputs)
 }
diff --git a/emmarin/cl/risc0_proofs/Cargo.toml b/emmarin/cl/risc0_proofs/Cargo.toml
index e544320..19d562a 100644
--- a/emmarin/cl/risc0_proofs/Cargo.toml
+++ b/emmarin/cl/risc0_proofs/Cargo.toml
@@ -7,5 +7,5 @@ edition = "2021"
 risc0-build = { version = "1.0" }
 
 [package.metadata.risc0]
-methods = ["bundle", "constraint_nop", "ptx", "stf_nop"]
+methods = ["balance", "constraint_nop", "ptx", "stf_nop"]
 
diff --git a/emmarin/cl/risc0_proofs/bundle/Cargo.toml b/emmarin/cl/risc0_proofs/balance/Cargo.toml
similarity index 97%
rename from emmarin/cl/risc0_proofs/bundle/Cargo.toml
rename to emmarin/cl/risc0_proofs/balance/Cargo.toml
index f207737..161f1c8 100644
--- a/emmarin/cl/risc0_proofs/bundle/Cargo.toml
+++ b/emmarin/cl/risc0_proofs/balance/Cargo.toml
@@ -1,5 +1,5 @@
 [package]
-name = "bundle"
+name = "balance"
 version = "0.1.0"
 edition = "2021"
 
diff --git a/emmarin/cl/risc0_proofs/bundle/src/main.rs b/emmarin/cl/risc0_proofs/balance/src/main.rs
similarity index 57%
rename from emmarin/cl/risc0_proofs/bundle/src/main.rs
rename to emmarin/cl/risc0_proofs/balance/src/main.rs
index 30d2f59..3342001 100644
--- a/emmarin/cl/risc0_proofs/bundle/src/main.rs
+++ b/emmarin/cl/risc0_proofs/balance/src/main.rs
@@ -12,13 +12,13 @@ use cl::cl::BalanceWitness;
 use risc0_zkvm::guest::env;
 
 fn main() {
-    let bundle_private: ledger_proof_statements::bundle::BundlePrivate = env::read();
+    let balance_private: ledger_proof_statements::balance::BalancePrivate = env::read();
 
-    let bundle_public = ledger_proof_statements::bundle::BundlePublic {
-        balances: Vec::from_iter(bundle_private.balances.iter().map(|b| b.commit())),
+    let balance_public = ledger_proof_statements::balance::BalancePublic {
+        balances: Vec::from_iter(balance_private.balances.iter().map(|b| b.commit())),
     };
 
-    assert!(BalanceWitness::combine(bundle_private.balances, [0u8; 16]).is_zero());
+    assert!(BalanceWitness::combine(balance_private.balances, [0u8; 16]).is_zero());
 
-    env::commit(&bundle_public);
+    env::commit(&balance_public);
 }
diff --git a/emmarin/cl/risc0_proofs/ptx/src/main.rs b/emmarin/cl/risc0_proofs/ptx/src/main.rs
index af4c2b1..83a8cee 100644
--- a/emmarin/cl/risc0_proofs/ptx/src/main.rs
+++ b/emmarin/cl/risc0_proofs/ptx/src/main.rs
@@ -7,16 +7,22 @@ fn main() {
     let PtxPrivate {
         ptx,
         input_cm_paths,
-        cm_root,
+        cm_roots,
         from,
         to,
     } = env::read();
 
     assert_eq!(ptx.inputs.len(), input_cm_paths.len());
-    for ((input, cm_path), zone_id) in ptx.inputs.iter().zip(input_cm_paths).zip(&from) {
+    for (((input, cm_path), zone_id), cm_root) in ptx
+        .inputs
+        .iter()
+        .zip(input_cm_paths)
+        .zip(&from)
+        .zip(&cm_roots)
+    {
         let note_cm = input.note_commitment(zone_id);
         let cm_leaf = merkle::leaf(note_cm.as_bytes());
-        assert_eq!(cm_root, merkle::path_root(cm_leaf, &cm_path));
+        assert_eq!(*cm_root, merkle::path_root(cm_leaf, &cm_path));
     }
 
     for output in ptx.outputs.iter() {
@@ -25,6 +31,6 @@ fn main() {
 
     env::commit(&PtxPublic {
         ptx: ptx.commit(&from, &to),
-        cm_root,
+        cm_roots,
     });
 }