Implement LedgerInterface for non-stark context

Author: sebastiencs

ledger/src/zkapps/intefaces.rs

@@ -477,29 +477,37 @@ where
     fn set_last_action_slot(&mut self, slot: Self::GlobalSlot);
 }
 
-pub trait LedgerInterface: LedgerIntf + Clone + ToFieldElements<Fp> {
+pub trait LedgerInterface: LedgerIntf + Clone {
     type W: WitnessGenerator<Fp>;
     type AccountUpdate: AccountUpdateInterface;
     type Account: AccountInterface;
     type Bool: BoolInterface;
     type InclusionProof;
 
-    fn empty() -> Self;
+    fn empty(depth: usize) -> Self;
     fn get_account(
         &self,
         account_update: &Self::AccountUpdate,
         w: &mut Self::W,
-    ) -> (Self::Account, Self::InclusionProof);
-    fn set_account(&mut self, account: (Self::Account, Self::InclusionProof), w: &mut Self::W);
+    ) -> Result<(Self::Account, Self::InclusionProof), String>;
+    fn set_account(
+        &mut self,
+        account: (Self::Account, Self::InclusionProof),
+        w: &mut Self::W,
+    ) -> Result<(), String>;
     fn check_inclusion(&self, account: &(Self::Account, Self::InclusionProof), w: &mut Self::W);
     fn check_account(
         public_key: &CompressedPubKey,
         token_id: &TokenId,
         account: (&Self::Account, &Self::InclusionProof),
         w: &mut Self::W,
-    ) -> Self::Bool;
-    fn exists_no_check(self, w: &mut Self::W) -> Self;
-    fn exists_no_check_on_bool(self, b: Self::Bool, w: &mut Self::W) -> Self;
+    ) -> Result<Self::Bool, String>;
+    fn exists_no_check(self, _w: &mut Self::W) -> Self {
+        self
+    }
+    fn exists_no_check_on_bool(self, _b: Self::Bool, _w: &mut Self::W) -> Self {
+        self
+    }
 }
 
 pub trait VerificationKeyHashInterface {

ledger/src/zkapps/non_snark.rs

@@ -11,23 +11,25 @@ use crate::{
         to_field_elements::ToFieldElements,
         transaction::Check,
         witness::Witness,
+        zkapp::StartDataSkeleton,
     },
     scan_state::{
         currency::{Amount, Balance, Index, Magnitude, Signed, Slot, SlotSpan, TxnVersion},
         transaction_logic::{
             account_check_timing,
             local_state::{CallStack, StackFrame},
             protocol_state::GlobalStateSkeleton,
+            set_with_location,
             zkapp_command::{
                 self, AccountPreconditions, AccountUpdate, CallForest, CheckAuthorizationResult,
                 SetOrKeep,
             },
             zkapp_statement::TransactionCommitment,
-            TimingValidation, TransactionFailure,
+            ExistingOrNew, TimingValidation, TransactionFailure,
         },
         zkapp_logic::controller_check,
     },
-    sparse_ledger::LedgerIntf,
+    sparse_ledger::{LedgerIntf, SparseLedger},
     zkapps::checks::ZkappCheck,
     Account, AccountId, Mask, MyCow, TokenId, ZkAppAccount, TXN_VERSION_CURRENT,
 };
@@ -45,7 +47,7 @@ use super::{
         TxnVersionInterface, VerificationKeyHashInterface, WitnessGenerator, ZkappApplication,
         ZkappHandler,
     },
-    zkapp_logic,
+    zkapp_logic::{self, ApplyZkappParams},
 };
 
 pub type GlobalStateForNonSnark<L> = GlobalStateSkeleton<
@@ -864,3 +866,182 @@ impl BranchInterface for NonSnarkBranch {
         BranchEvaluation::Pending(run)
     }
 }
+
+fn get_with_location<L>(
+    ledger: &L,
+    account_id: &AccountId,
+) -> Result<(ExistingOrNew<L::Location>, Box<Account>), String>
+where
+    L: LedgerIntf,
+{
+    match ledger.location_of_account(account_id) {
+        Some(location) => match ledger.get(&location) {
+            Some(account) => Ok((ExistingOrNew::Existing(location), account)),
+            None => Err("Ledger location with no account".to_string()),
+        },
+        None => Ok((
+            ExistingOrNew::New,
+            Box::new(Account::create_with(account_id.clone(), Balance::zero())),
+        )),
+    }
+}
+
+mod ledger {
+    use super::*;
+
+    type InclusionProof<L> = ExistingOrNew<<L as LedgerIntf>::Location>;
+
+    pub(super) fn get_account<L: LedgerIntf>(
+        ledger: &L,
+        account_update: &AccountUpdate,
+    ) -> Result<(Account, InclusionProof<L>), String> {
+        let (loc, account) = get_with_location(ledger, &account_update.account_id())?;
+        Ok((*account, loc))
+    }
+    pub(super) fn set_account<L: LedgerIntf>(
+        ledger: &mut L,
+        account: (Account, InclusionProof<L>),
+    ) -> Result<(), String> {
+        let (account, location) = account;
+        set_with_location(ledger, &location, Box::new(account))
+    }
+    pub(super) fn check_account<L: LedgerIntf>(
+        public_key: &CompressedPubKey,
+        token_id: &TokenId,
+        account: (&Account, &InclusionProof<L>),
+    ) -> Result<bool, String> {
+        let (account, loc) = account;
+        if public_key != &account.public_key {
+            return Err("check_account: public_key != &account.public_key".to_string());
+        }
+        if token_id != &account.token_id {
+            return Err("check_account: token_id != &account.token_id".to_string());
+        }
+        match loc {
+            ExistingOrNew::Existing(_) => Ok(false),
+            ExistingOrNew::New => Ok(true),
+        }
+    }
+}
+
+impl LedgerInterface for Mask {
+    type W = ();
+    type AccountUpdate = AccountUpdate;
+    type Account = Account;
+    type Bool = bool;
+    type InclusionProof = ExistingOrNew<<Mask as LedgerIntf>::Location>;
+
+    fn empty(depth: usize) -> Self {
+        <Self as LedgerIntf>::empty(depth)
+    }
+    fn get_account(
+        &self,
+        account_update: &Self::AccountUpdate,
+        _w: &mut Self::W,
+    ) -> Result<(Self::Account, Self::InclusionProof), String> {
+        ledger::get_account(self, account_update)
+    }
+    fn set_account(
+        &mut self,
+        account: (Self::Account, Self::InclusionProof),
+        _w: &mut Self::W,
+    ) -> Result<(), String> {
+        ledger::set_account(self, account)
+    }
+    fn check_inclusion(&self, _account: &(Self::Account, Self::InclusionProof), _w: &mut Self::W) {
+        // Nothing
+    }
+    fn check_account(
+        public_key: &CompressedPubKey,
+        token_id: &TokenId,
+        account: (&Self::Account, &Self::InclusionProof),
+        _w: &mut Self::W,
+    ) -> Result<Self::Bool, String> {
+        ledger::check_account::<Self>(public_key, token_id, account)
+    }
+}
+impl LedgerInterface for SparseLedger {
+    type W = ();
+    type AccountUpdate = AccountUpdate;
+    type Account = Account;
+    type Bool = bool;
+    type InclusionProof = ExistingOrNew<<Mask as LedgerIntf>::Location>;
+
+    fn empty(depth: usize) -> Self {
+        <Self as LedgerIntf>::empty(depth)
+    }
+    fn get_account(
+        &self,
+        account_update: &Self::AccountUpdate,
+        _w: &mut Self::W,
+    ) -> Result<(Self::Account, Self::InclusionProof), String> {
+        ledger::get_account(self, account_update)
+    }
+    fn set_account(
+        &mut self,
+        account: (Self::Account, Self::InclusionProof),
+        _w: &mut Self::W,
+    ) -> Result<(), String> {
+        ledger::set_account(self, account)
+    }
+    fn check_inclusion(&self, _account: &(Self::Account, Self::InclusionProof), _w: &mut Self::W) {
+        // Nothing
+    }
+    fn check_account(
+        public_key: &CompressedPubKey,
+        token_id: &TokenId,
+        account: (&Self::Account, &Self::InclusionProof),
+        _w: &mut Self::W,
+    ) -> Result<Self::Bool, String> {
+        ledger::check_account::<Self>(public_key, token_id, account)
+    }
+}
+
+pub fn step<L>(
+    global_state: &mut GlobalStateForNonSnark<L>,
+    local_state: &mut zkapp_logic::LocalState<ZkappNonSnark<L>>,
+) -> Result<(), String>
+where
+    L: LedgerInterface<
+        W = (),
+        AccountUpdate = AccountUpdate,
+        Account = crate::Account,
+        Bool = bool,
+    >,
+{
+    zkapp_logic::apply(
+        ApplyZkappParams {
+            is_start: zkapp_logic::IsStart::No,
+            global_state,
+            local_state,
+            single_data: (),
+        },
+        &mut (),
+    )
+}
+
+type StartData = StartDataSkeleton<CallForest<AccountUpdate>, bool>;
+
+pub fn start<L>(
+    global_state: &mut GlobalStateForNonSnark<L>,
+    local_state: &mut zkapp_logic::LocalState<ZkappNonSnark<L>>,
+    start_data: StartData,
+) -> Result<(), String>
+where
+    L: LedgerInterface<
+        W = (),
+        AccountUpdate = AccountUpdate,
+        Account = crate::Account,
+        Bool = bool,
+    >,
+{
+    zkapp_logic::apply(
+        ApplyZkappParams {
+            is_start: zkapp_logic::IsStart::Yes(start_data),
+            global_state,
+            local_state,
+            single_data: (),
+        },
+        &mut (),
+    )
+}

ledger/src/zkapps/snark.rs

@@ -992,16 +992,16 @@ impl LedgerInterface for LedgerWithHash {
     type Bool = SnarkBool;
     type InclusionProof = Vec<(Boolean, Fp)>;
 
-    fn empty() -> Self {
-        let mut ledger = <SparseLedger as crate::sparse_ledger::LedgerIntf>::empty(0);
+    fn empty(depth: usize) -> Self {
+        let mut ledger = <SparseLedger as crate::sparse_ledger::LedgerIntf>::empty(depth);
         let hash = ledger.merkle_root();
         Self { ledger, hash }
     }
     fn get_account(
         &self,
         account_update: &Self::AccountUpdate,
         w: &mut Self::W,
-    ) -> (Self::Account, Self::InclusionProof) {
+    ) -> Result<(Self::Account, Self::InclusionProof), String> {
         let Self {
             ledger,
             hash: _root,
@@ -1026,14 +1026,19 @@ impl LedgerInterface for LedgerWithHash {
                 })
                 .collect::<Vec<_>>(),
         );
-        (account, inclusion)
+        Ok((account, inclusion))
     }
-    fn set_account(&mut self, (a, incl): (Self::Account, Self::InclusionProof), w: &mut Self::W) {
+    fn set_account(
+        &mut self,
+        (a, incl): (Self::Account, Self::InclusionProof),
+        w: &mut Self::W,
+    ) -> Result<(), String> {
         let Self { ledger, hash } = self;
         let new_hash = implied_root(&a, &incl, w);
         let idx = ledger.find_index_exn(a.id());
         ledger.set_exn(idx, a.data);
         *hash = new_hash;
+        Ok(())
     }
     fn check_inclusion(
         &self,
@@ -1047,7 +1052,7 @@ impl LedgerInterface for LedgerWithHash {
         token_id: &TokenId,
         account: (&Self::Account, &Self::InclusionProof),
         w: &mut Self::W,
-    ) -> Self::Bool {
+    ) -> Result<Self::Bool, String> {
         let (WithLazyHash { data: account, .. }, _) = account;
         let is_new = checked_equal_compressed_key_const_and(
             &account.public_key,
@@ -1058,7 +1063,7 @@ impl LedgerInterface for LedgerWithHash {
         Boolean::assert_any(&[is_new, is_same], w);
         let is_same_token = field::equal(token_id.0, account.token_id.0, w);
         Boolean::assert_any(&[is_new, is_same_token], w);
-        is_new.var()
+        Ok(is_new.var())
     }
     fn exists_no_check(self, w: &mut Self::W) -> Self {
         w.exists_no_check(self.hash);

ledger/src/zkapps/zkapp_logic.rs

@@ -28,12 +28,6 @@ pub enum IsStart<T> {
     Compute(T),
 }
 
-enum PerformResult<Z: ZkappApplication> {
-    None,
-    Bool(Z::Bool),
-    Account(Z::Account),
-}
-
 struct GetNextAccountUpdateResult<Z: ZkappApplication> {
     account_update: Z::AccountUpdate,
     caller_id: TokenId,
@@ -367,13 +361,11 @@ where
         IsStart::Yes(start_data) => start_data.will_succeed,
         IsStart::No => local_state.will_succeed,
     };
-    local_state.ledger = w.exists_no_check_on_bool(
-        is_start2,
-        match is_start2.as_boolean() {
-            Boolean::True => global_state.first_pass_ledger(),
-            Boolean::False => local_state.ledger.clone(),
-        },
-    );
+    local_state.ledger = match is_start2.as_boolean() {
+        Boolean::True => global_state.first_pass_ledger(),
+        Boolean::False => local_state.ledger.clone(),
+    }
+    .exists_no_check_on_bool(is_start2, w);
     local_state.will_succeed = will_succeed;
 
     let ((account_update, remaining, call_stack), account_update_forest, (mut a, inclusion_proof)) = {
@@ -440,7 +432,7 @@ where
             );
         };
 
-        let acct = local_state.ledger.get_account(&account_update, w);
+        let acct = local_state.ledger.get_account(&account_update, w)?;
         local_state.ledger.check_inclusion(&acct, w);
 
         let (transaction_commitment, full_transaction_commitment) = match is_start {
@@ -490,7 +482,7 @@ where
         &account_update.body().token_id,
         (&a, &inclusion_proof),
         w,
-    );
+    )?;
 
     {
         let self_delegate = {
@@ -1211,7 +1203,7 @@ where
         overflowed.neg(),
         w,
     );
-    local_state.ledger.set_account((a, inclusion_proof), w);
+    local_state.ledger.set_account((a, inclusion_proof), w)?;
 
     let is_last_account_update = Z::CallForest::is_empty(Z::StackFrame::calls(&remaining), w);
     // We decompose this way because of OCaml evaluation order
@@ -1347,7 +1339,7 @@ where
             Boolean::False => local_state.success,
         });
         let ledger = match is_last_account_update.as_boolean() {
-            Boolean::True => Z::Ledger::empty(),
+            Boolean::True => Z::Ledger::empty(0),
             Boolean::False => local_state.ledger.clone(),
         }
         .exists_no_check(w);