Add zone logic support for withdrawal txs

goas/zone/common/Cargo.toml

@@ -5,3 +5,5 @@ edition = "2021"
 
 [dependencies]
 serde = { version = "1", features = ["derive"] }
+cl = { path = "../../cl/cl" }
+proof_statements = { path = "../../cl/proof_statements" }

goas/zone/common/src/lib.rs

@@ -1,28 +1,106 @@
-use serde::{Serialize, Deserialize};
+use cl::nullifier::{Nullifier, NullifierCommitment};
+use serde::{Deserialize, Serialize};
 use std::collections::BTreeMap;
 
+// TODO: sparse merkle tree
+pub const MAX_BALANCES: usize = 1 << 8;
+pub const MAX_TXS: usize = 1 << 8;
+pub const MAX_EVENTS: usize = 1 << 8;
+
 // state of the zone
-pub type State = BTreeMap<u32, u32>;
-// list of all inputs that were executed up to this point
-pub type Journal = Vec<Input>;
-
 #[derive(Clone, Copy, Serialize, Deserialize)]
-pub enum Input {
-    Transfer { from: u32, to: u32, amount: u32 },
-    None,
+pub struct StateCommitment([u8; 32]);
+
+#[derive(Clone, Serialize, Deserialize)]
+pub struct StateWitness {
+    pub balances: BTreeMap<u32, u32>,
+    pub included_txs: Vec<Input>,
+    pub output_events: Vec<Event>,
 }
 
-
-/// State transition function of the zone
-pub fn stf(mut state: State, input: Input) -> State {
-    match input {
-        Input::Transfer { from, to, amount } => {
-            // compute transfer
-            let from = state.entry(from).or_insert(0);
-            *from = from.checked_sub(amount).unwrap();
-            *state.entry(to).or_insert(0) += amount;
-        }
-        Input::None => {}
+impl StateWitness {
+    pub fn commit(&self) -> StateCommitment {
+        let root = self.balances_root();
+        let root = cl::merkle::node(self.events_root(), root);
+        let root = cl::merkle::node(self.included_txs_root(), root);
+        StateCommitment(root)
+    }
+
+    fn events_root(&self) -> [u8; 32] {
+        let event_bytes = Vec::from_iter(
+            self.output_events
+                .iter()
+                .map(Event::to_bytes)
+                .map(Vec::from_iter),
+        );
+        let event_merkle_leaves = cl::merkle::padded_leaves(&event_bytes);
+        cl::merkle::root::<MAX_EVENTS>(event_merkle_leaves)
+    }
+
+    fn included_txs_root(&self) -> [u8; 32] {
+        let tx_bytes = Vec::from_iter(
+            self.included_txs
+                .iter()
+                .map(Input::to_bytes)
+                .map(Vec::from_iter),
+        );
+        let tx_merkle_leaves = cl::merkle::padded_leaves(&tx_bytes);
+        cl::merkle::root::<MAX_TXS>(tx_merkle_leaves)
+    }
+
+    fn balances_root(&self) -> [u8; 32] {
+        let balance_bytes = Vec::from_iter(self.balances.iter().map(|(k, v)| {
+            let mut bytes = [0; 8];
+            bytes.copy_from_slice(&k.to_le_bytes());
+            bytes[8..].copy_from_slice(&v.to_le_bytes());
+            bytes.to_vec()
+        }));
+        let balance_merkle_leaves = cl::merkle::padded_leaves(&balance_bytes);
+        cl::merkle::root::<MAX_BALANCES>(balance_merkle_leaves)
+    }
+}
+
+#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
+pub struct Withdraw {
+    pub from: u32,
+    pub amount: u32,
+    pub to: NullifierCommitment,
+    pub nf: Nullifier,
+}
+
+impl Withdraw {
+    pub fn to_bytes(&self) -> [u8; 72] {
+        let mut bytes = [0; 72];
+        bytes[0..4].copy_from_slice(&self.from.to_le_bytes());
+        bytes[4..8].copy_from_slice(&self.amount.to_le_bytes());
+        bytes[8..40].copy_from_slice(self.to.as_bytes());
+        bytes[40..72].copy_from_slice(self.nf.as_bytes());
+        bytes
+    }
+}
+
+#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
+pub enum Input {
+    Withdraw(Withdraw),
+}
+
+impl Input {
+    pub fn to_bytes(&self) -> Vec<u8> {
+        match self {
+            Input::Withdraw(withdraw) => withdraw.to_bytes().to_vec(),
+        }
+    }
+}
+
+#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
+pub enum Event {
+    Spend(proof_statements::zone_funds::Spend),
+}
+
+impl Event {
+    pub fn to_bytes(&self) -> Vec<u8> {
+        match self {
+            Event::Spend(spend) => spend.to_bytes().to_vec(),
+        }
     }
-    state
 }

goas/zone/host/src/main.rs

@@ -1,77 +1,28 @@
 // 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 blake2::{Blake2s256, Digest};
 use common::*;
 use risc0_zkvm::{default_prover, ExecutorEnv};
+use std::path::PathBuf;
 
 use clap::Parser;
 
 #[derive(Parser, Debug)]
 #[command(version, about, long_about = None)]
 enum Action {
-    Stf,
+    Stf {
+        // path to bincode-encoded state witness
+        state: PathBuf,
+        // path to bincode-encoded inputs
+        inputs: PathBuf,
+    },
 }
 
-fn stf_prove_stark() {
-    let mut rng = rand::thread_rng();
-
-    let state: State = [(0, 1000)].into_iter().collect();
-    let journal = vec![];
-    let zone_input = Input::Transfer {
-        from: 0,
-        to: 1,
-        amount: 10,
-    };
-
-    let in_state_cm = calculate_state_hash(&state);
-    let in_journal_cm = calculate_journal_hash(&journal);
-    let in_state_root = cl::merkle::node(in_state_cm, in_journal_cm);
-    let in_note = cl::NoteWitness::new(1, "ZONE", in_state_root, &mut rng);
-
-    let mut out_journal = journal.clone();
-    out_journal.push(zone_input);
-
-    let out_state_cm = calculate_state_hash(&stf(state.clone(), zone_input));
-    let out_journal_cm = calculate_journal_hash(&out_journal);
-    let out_state_root = cl::merkle::node(out_state_cm, out_journal_cm);
-    let out_note = cl::NoteWitness::new(1, "ZONE", out_state_root, &mut rng);
-
-    let input = cl::InputWitness::random(in_note, &mut rng);
-    let output = cl::OutputWitness::random(
-        out_note,
-        cl::NullifierSecret::random(&mut rng).commit(),
-        &mut rng,
-    );
-    let ptx = cl::PartialTx::from_witness(cl::PartialTxWitness {
-        inputs: vec![input.clone()],
-        outputs: vec![output.clone()],
-    });
-
-    let ptx_root = ptx.root().0;
-    let in_ptx_path = ptx.input_merkle_path(0);
-    let out_ptx_path = ptx.output_merkle_path(0);
-
+fn stf_prove_stark(state: StateWitness, inputs: Vec<Input>) {
     let env = ExecutorEnv::builder()
-        .write(&ptx_root)
-        .unwrap()
-        .write(&ptx.input_root())
-        .unwrap()
-        .write(&ptx.output_root())
-        .unwrap()
-        .write(&in_ptx_path)
-        .unwrap()
-        .write(&out_ptx_path)
-        .unwrap()
-        .write(&input)
-        .unwrap()
-        .write(&output)
-        .unwrap()
-        .write(&zone_input)
+        .write(&inputs)
         .unwrap()
         .write(&state)
         .unwrap()
-        .write(&journal)
-        .unwrap()
         .build()
         .unwrap();
 
@@ -109,16 +60,10 @@ fn main() {
     let action = Action::parse();
 
     match action {
-        Action::Stf => stf_prove_stark(),
+        Action::Stf { state, inputs } => {
+            let state = bincode::deserialize(&std::fs::read(state).unwrap()).unwrap();
+            let inputs = bincode::deserialize(&std::fs::read(inputs).unwrap()).unwrap();
+            stf_prove_stark(state, inputs);
+        }
     }
 }
-
-fn calculate_state_hash(state: &State) -> [u8; 32] {
-    let bytes = bincode::serialize(state).unwrap();
-    Blake2s256::digest(&bytes).into()
-}
-
-fn calculate_journal_hash(journal: &Journal) -> [u8; 32] {
-    let bytes = bincode::serialize(journal).unwrap();
-    Blake2s256::digest(&bytes).into()
-}

goas/zone/methods/guest/Cargo.toml

@@ -14,6 +14,7 @@ serde = { version = "1.0", features = ["derive"] }
 bincode = "1"
 common = { path = "../../common" }
 cl = { path = "../../../cl/cl" }
+proof_statements = { path = "../../../cl/proof_statements" }
 
 [patch.crates-io]
 # Placing these patch statement in the workspace Cargo.toml will add RISC Zero SHA-256 and bigint

goas/zone/methods/guest/src/main.rs

@@ -1,113 +1,40 @@
-use blake2::{Blake2s256, Digest};
-use cl::input::InputWitness;
-use cl::merkle;
-use cl::output::OutputWitness;
 use common::*;
+use proof_statements::zone_funds::Spend;
 use risc0_zkvm::guest::env;
 
-/// Public Inputs:
-/// * ptx_root: the root of the partial tx merkle tree of inputs/outputs
-/// Private inputs:
-/// TODO
+fn withdraw(mut state: StateWitness, withdraw: Withdraw) -> StateWitness {
+    state.included_txs.push(Input::Withdraw(withdraw));
 
-fn execute(
-    ptx_root: [u8; 32],
-    input_root: [u8; 32],
-    output_root: [u8; 32],
-    in_ptx_path: Vec<merkle::PathNode>,
-    out_ptx_path: Vec<merkle::PathNode>,
-    in_note: InputWitness,
-    out_note: OutputWitness,
-    input: Input,
-    state: State,
-    mut journal: Journal,
-) {
-    // verify ptx/cl preconditions
-    eprintln!("start exec: {}", env::cycle_count());
-    assert_eq!(ptx_root, merkle::node(input_root, output_root));
-    eprintln!("ptx_root: {}", env::cycle_count());
+    let Withdraw {
+        from,
+        amount,
+        to,
+        nf,
+    } = withdraw;
 
-    // Glue the zone and the cl together, specifically, it verifies the note requesting
-    // a transfer is included as part of the same transaction in the cl
-    let in_comm = in_note.commit().to_bytes();
-    eprintln!("input comm: {}", env::cycle_count());
+    let from = state.balances.entry(from).or_insert(0);
+    *from = from.checked_sub(amount).unwrap();
+    let spend_auth = Spend {
+        amount: amount.into(),
+        to,
+        nf,
+    };
 
-    assert_eq!(
-        merkle::path_root(merkle::leaf(&in_comm), &in_ptx_path),
-        input_root
-    );
-    eprintln!("input merkle path: {}", env::cycle_count());
-
-    // check the commitments match the actual data
-    let state_cm = calculate_state_hash(&state);
-    let journal_cm = calculate_journal_hash(&journal);
-    let state_root = merkle::node(state_cm, journal_cm);
-    assert_eq!(state_root, in_note.note.state);
-    eprintln!("input state root: {}", env::cycle_count());
-
-    // then run the state transition function
-    let state = stf(state, input);
-    journal.push(input);
-    eprintln!("stf: {}", env::cycle_count());
-
-    // verifying ptx/cl postconditions
-
-    let out_state_cm = calculate_state_hash(&state);
-    let out_journal_cm = calculate_journal_hash(&journal);
-    let out_state_root = merkle::node(out_state_cm, out_journal_cm);
-    // TODO: verify death constraints are propagated
-    assert_eq!(out_state_root, out_note.note.state);
-    eprintln!("out state root: {}", env::cycle_count());
-
-    // Glue the zone and the cl together, specifically, it verifies an output note
-    // containing the zone state is included as part of the same transaction in the cl
-    // (this is done in the death condition to disallow burning)
-    let out_comm = out_note.commit().to_bytes();
-    eprintln!("output comm: {}", env::cycle_count());
-
-    assert_eq!(
-        merkle::path_root(merkle::leaf(&out_comm), &out_ptx_path),
-        output_root
-    );
-    eprintln!("out merkle proof: {}", env::cycle_count());
+    state.output_events.push(Event::Spend(spend_auth));
+    state
 }
 
 fn main() {
-    // public input
-    let ptx_root: [u8; 32] = env::read();
+    let inputs: Vec<Input> = env::read();
+    let mut state: StateWitness = env::read();
 
-    // private input
-    let input_root: [u8; 32] = env::read();
-    let output_root: [u8; 32] = env::read();
-    let in_ptx_path: Vec<merkle::PathNode> = env::read();
-    let out_ptx_path: Vec<merkle::PathNode> = env::read();
-    let in_note: InputWitness = env::read();
-    let out_note: OutputWitness = env::read();
-    let input: Input = env::read();
-    let state: State = env::read();
-    let journal: Journal = env::read();
+    for input in inputs {
+        match input {
+            Input::Withdraw(input) => {
+                state = withdraw(state, input);
+            }
+        }
+    }
 
-    eprintln!("parse input: {}", env::cycle_count());
-    execute(
-        ptx_root,
-        input_root,
-        output_root,
-        in_ptx_path,
-        out_ptx_path,
-        in_note,
-        out_note,
-        input,
-        state,
-        journal,
-    );
-}
-
-fn calculate_state_hash(state: &State) -> [u8; 32] {
-    let bytes = bincode::serialize(state).unwrap();
-    Blake2s256::digest(&bytes).into()
-}
-
-fn calculate_journal_hash(journal: &Journal) -> [u8; 32] {
-    let bytes = bincode::serialize(journal).unwrap();
-    Blake2s256::digest(&bytes).into()
+    env::commit(&state.commit());
 }