mina_tree/scan_state/transaction_logic/
local_state.rs

1use super::{
2    protocol_state::{GlobalState, ProtocolStateView},
3    transaction_applied::ZkappCommandApplied,
4    transaction_partially_applied::ZkappCommandPartiallyApplied,
5    zkapp_command::{AccountUpdate, CallForest, WithHash, ZkAppCommand},
6    TransactionFailure, TransactionStatus, WithStatus,
7};
8use crate::{
9    proofs::{
10        field::{field, Boolean, ToBoolean},
11        numbers::nat::CheckedNat,
12        to_field_elements::ToFieldElements,
13        witness::Witness,
14    },
15    scan_state::currency::{Amount, Index, Magnitude, Signed, Slot},
16    sparse_ledger::LedgerIntf,
17    zkapps::{
18        self,
19        interfaces::{
20            CallStackInterface, IndexInterface, SignedAmountInterface, StackFrameInterface,
21        },
22        non_snark::{LedgerNonSnark, ZkappNonSnark},
23    },
24    AccountId, AccountIdOrderable, AppendToInputs, ToInputs, TokenId,
25};
26use ark_ff::Zero;
27use itertools::{FoldWhile, Itertools};
28use mina_core::constants::ConstraintConstants;
29use mina_curves::pasta::Fp;
30use poseidon::hash::{hash_with_kimchi, params::MINA_ACCOUNT_UPDATE_STACK_FRAME, Inputs};
31use std::{cell::RefCell, collections::BTreeMap, rc::Rc};
32
33#[derive(Debug, Clone)]
34pub struct StackFrame {
35    pub caller: TokenId,
36    pub caller_caller: TokenId,
37    pub calls: CallForest<AccountUpdate>, // TODO
38}
39
40// <https://github.com/MinaProtocol/mina/blob/78535ae3a73e0e90c5f66155365a934a15535779/src/lib/transaction_snark/transaction_snark.ml#L1081>
41#[derive(Debug, Clone)]
42pub struct StackFrameCheckedFrame {
43    pub caller: TokenId,
44    pub caller_caller: TokenId,
45    pub calls: WithHash<CallForest<AccountUpdate>>,
46    /// Hack until we have proper cvar
47    pub is_default: bool,
48}
49
50impl ToFieldElements<Fp> for StackFrameCheckedFrame {
51    fn to_field_elements(&self, fields: &mut Vec<Fp>) {
52        let Self {
53            caller,
54            caller_caller,
55            calls,
56            is_default: _,
57        } = self;
58
59        // calls.hash().to_field_elements(fields);
60        calls.hash.to_field_elements(fields);
61        caller_caller.to_field_elements(fields);
62        caller.to_field_elements(fields);
63    }
64}
65
66enum LazyValueInner<T, D> {
67    Value(T),
68    Fun(Box<dyn FnOnce(&mut D) -> T>),
69    None,
70}
71
72impl<T, D> Default for LazyValueInner<T, D> {
73    fn default() -> Self {
74        Self::None
75    }
76}
77
78pub struct LazyValue<T, D> {
79    value: Rc<RefCell<LazyValueInner<T, D>>>,
80}
81
82impl<T, D> Clone for LazyValue<T, D> {
83    fn clone(&self) -> Self {
84        Self {
85            value: Rc::clone(&self.value),
86        }
87    }
88}
89
90impl<T: std::fmt::Debug, D> std::fmt::Debug for LazyValue<T, D> {
91    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
92        let v = self.try_get();
93        f.debug_struct("LazyValue").field("value", &v).finish()
94    }
95}
96
97impl<T, D> LazyValue<T, D> {
98    pub fn make<F>(fun: F) -> Self
99    where
100        F: FnOnce(&mut D) -> T + 'static,
101    {
102        Self {
103            value: Rc::new(RefCell::new(LazyValueInner::Fun(Box::new(fun)))),
104        }
105    }
106
107    fn get_impl(&self) -> std::cell::Ref<'_, T> {
108        use std::cell::Ref;
109
110        let inner = self.value.borrow();
111        Ref::map(inner, |inner| {
112            let LazyValueInner::Value(value) = inner else {
113                panic!("invalid state");
114            };
115            value
116        })
117    }
118
119    /// Returns the value when it already has been "computed"
120    pub fn try_get(&self) -> Option<std::cell::Ref<'_, T>> {
121        let inner = self.value.borrow();
122
123        match &*inner {
124            LazyValueInner::Value(_) => {}
125            LazyValueInner::Fun(_) => return None,
126            LazyValueInner::None => panic!("invalid state"),
127        }
128
129        Some(self.get_impl())
130    }
131
132    pub fn get(&self, data: &mut D) -> std::cell::Ref<'_, T> {
133        let v = self.value.borrow();
134
135        if let LazyValueInner::Fun(_) = &*v {
136            std::mem::drop(v);
137
138            let LazyValueInner::Fun(fun) = self.value.take() else {
139                panic!("invalid state");
140            };
141
142            let data = fun(data);
143            self.value.replace(LazyValueInner::Value(data));
144        };
145
146        self.get_impl()
147    }
148}
149
150#[derive(Clone, Debug)]
151pub struct WithLazyHash<T> {
152    pub data: T,
153    hash: LazyValue<Fp, Witness<Fp>>,
154}
155
156impl<T> WithLazyHash<T> {
157    pub fn new<F>(data: T, fun: F) -> Self
158    where
159        F: FnOnce(&mut Witness<Fp>) -> Fp + 'static,
160    {
161        Self {
162            data,
163            hash: LazyValue::make(fun),
164        }
165    }
166
167    pub fn hash(&self, w: &mut Witness<Fp>) -> Fp {
168        *self.hash.get(w)
169    }
170}
171
172impl<T> std::ops::Deref for WithLazyHash<T> {
173    type Target = T;
174
175    fn deref(&self) -> &Self::Target {
176        &self.data
177    }
178}
179
180impl<T> ToFieldElements<Fp> for WithLazyHash<T> {
181    fn to_field_elements(&self, fields: &mut Vec<Fp>) {
182        let hash = self.hash.try_get().expect("hash hasn't been computed yet");
183        hash.to_field_elements(fields)
184    }
185}
186
187// <https://github.com/MinaProtocol/mina/blob/78535ae3a73e0e90c5f66155365a934a15535779/src/lib/transaction_snark/transaction_snark.ml#L1083>
188pub type StackFrameChecked = WithLazyHash<StackFrameCheckedFrame>;
189
190impl Default for StackFrame {
191    fn default() -> Self {
192        StackFrame {
193            caller: TokenId::default(),
194            caller_caller: TokenId::default(),
195            calls: CallForest::new(),
196        }
197    }
198}
199
200impl StackFrame {
201    pub fn empty() -> Self {
202        Self {
203            caller: TokenId::default(),
204            caller_caller: TokenId::default(),
205            calls: CallForest(Vec::new()),
206        }
207    }
208
209    /// TODO: this needs to be tested
210    ///
211    /// <https://github.com/MinaProtocol/mina/blob/2ee6e004ba8c6a0541056076aab22ea162f7eb3a/src/lib/mina_base/stack_frame.ml#L90>
212    pub fn hash(&self) -> Fp {
213        let mut inputs = Inputs::new();
214
215        inputs.append_field(self.caller.0);
216        inputs.append_field(self.caller_caller.0);
217
218        self.calls.ensure_hashed();
219        let field = match self.calls.0.first() {
220            None => Fp::zero(),
221            Some(calls) => calls.stack_hash.get().unwrap(), // Never fail, we called `ensure_hashed`
222        };
223        inputs.append_field(field);
224
225        hash_with_kimchi(&MINA_ACCOUNT_UPDATE_STACK_FRAME, &inputs.to_fields())
226    }
227
228    pub fn digest(&self) -> Fp {
229        self.hash()
230    }
231
232    pub fn unhash(&self, _h: Fp, w: &mut Witness<Fp>) -> StackFrameChecked {
233        let v = self.exists_elt(w);
234        v.hash(w);
235        v
236    }
237
238    pub fn exists_elt(&self, w: &mut Witness<Fp>) -> StackFrameChecked {
239        // We decompose this way because of OCaml evaluation order
240        let calls = WithHash {
241            data: self.calls.clone(),
242            hash: w.exists(self.calls.hash()),
243        };
244        let caller_caller = w.exists(self.caller_caller.clone());
245        let caller = w.exists(self.caller.clone());
246
247        let frame = StackFrameCheckedFrame {
248            caller,
249            caller_caller,
250            calls,
251            is_default: false,
252        };
253
254        StackFrameChecked::of_frame(frame)
255    }
256}
257
258impl StackFrameCheckedFrame {
259    pub fn hash(&self, w: &mut Witness<Fp>) -> Fp {
260        let mut inputs = Inputs::new();
261
262        inputs.append(&self.caller);
263        inputs.append(&self.caller_caller.0);
264        inputs.append(&self.calls.hash);
265
266        let fields = inputs.to_fields();
267
268        if self.is_default {
269            use crate::proofs::transaction::transaction_snark::checked_hash3;
270            checked_hash3(&MINA_ACCOUNT_UPDATE_STACK_FRAME, &fields, w)
271        } else {
272            use crate::proofs::transaction::transaction_snark::checked_hash;
273            checked_hash(&MINA_ACCOUNT_UPDATE_STACK_FRAME, &fields, w)
274        }
275    }
276}
277
278impl StackFrameChecked {
279    pub fn of_frame(frame: StackFrameCheckedFrame) -> Self {
280        // TODO: Don't clone here
281        let frame2 = frame.clone();
282        let hash = LazyValue::make(move |w: &mut Witness<Fp>| frame2.hash(w));
283
284        Self { data: frame, hash }
285    }
286}
287
288#[derive(Debug, Clone)]
289pub struct CallStack(pub Vec<StackFrame>);
290
291impl Default for CallStack {
292    fn default() -> Self {
293        Self::new()
294    }
295}
296
297impl CallStack {
298    pub fn new() -> Self {
299        CallStack(Vec::new())
300    }
301
302    pub fn is_empty(&self) -> bool {
303        self.0.is_empty()
304    }
305
306    pub fn iter(&self) -> impl Iterator<Item = &StackFrame> {
307        self.0.iter().rev()
308    }
309
310    pub fn push(&self, stack_frame: &StackFrame) -> Self {
311        let mut ret = self.0.clone();
312        ret.push(stack_frame.clone());
313        Self(ret)
314    }
315
316    pub fn pop(&self) -> Option<(StackFrame, CallStack)> {
317        let mut ret = self.0.clone();
318        ret.pop().map(|frame| (frame, Self(ret)))
319    }
320
321    pub fn pop_exn(&self) -> (StackFrame, CallStack) {
322        let mut ret = self.0.clone();
323        if let Some(frame) = ret.pop() {
324            (frame, Self(ret))
325        } else {
326            panic!()
327        }
328    }
329}
330
331// NOTE: It looks like there are different instances of the polymorphic LocalEnv type
332// One with concrete types for the stack frame, call stack, and ledger. Created from the Env
333// And the other with their hashes. To differentiate them I renamed the first LocalStateEnv
334// Maybe a better solution is to keep the LocalState name and put it under a different module
335// pub type LocalStateEnv<L> = LocalStateSkeleton<
336//     L,                            // ledger
337//     StackFrame,                   // stack_frame
338//     CallStack,                    // call_stack
339//     ReceiptChainHash,             // commitments
340//     Signed<Amount>,               // excess & supply_increase
341//     Vec<Vec<TransactionFailure>>, // failure_status_tbl
342//     bool,                         // success & will_succeed
343//     Index,                        // account_update_index
344// >;
345
346pub type LocalStateEnv<L> = crate::zkapps::zkapp_logic::LocalState<ZkappNonSnark<L>>;
347
348// TODO: Dedub this with `crate::zkapps::zkapp_logic::LocalState`
349#[derive(Debug, Clone)]
350pub struct LocalStateSkeleton<
351    L: LedgerIntf + Clone,
352    StackFrame: StackFrameInterface,
353    CallStack: CallStackInterface,
354    TC,
355    SignedAmount: SignedAmountInterface,
356    FailuresTable,
357    Bool,
358    Index: IndexInterface,
359> {
360    pub stack_frame: StackFrame,
361    pub call_stack: CallStack,
362    pub transaction_commitment: TC,
363    pub full_transaction_commitment: TC,
364    pub excess: SignedAmount,
365    pub supply_increase: SignedAmount,
366    pub ledger: L,
367    pub success: Bool,
368    pub account_update_index: Index,
369    // TODO: optimize by reversing the insertion order
370    pub failure_status_tbl: FailuresTable,
371    pub will_succeed: Bool,
372}
373
374// impl<L> LocalStateEnv<L>
375// where
376//     L: LedgerNonSnark,
377// {
378//     pub fn add_new_failure_status_bucket(&self) -> Self {
379//         let mut failure_status_tbl = self.failure_status_tbl.clone();
380//         failure_status_tbl.insert(0, Vec::new());
381//         Self {
382//             failure_status_tbl,
383//             ..self.clone()
384//         }
385//     }
386
387//     pub fn add_check(&self, failure: TransactionFailure, b: bool) -> Self {
388//         let failure_status_tbl = if !b {
389//             let mut failure_status_tbl = self.failure_status_tbl.clone();
390//             failure_status_tbl[0].insert(0, failure);
391//             failure_status_tbl
392//         } else {
393//             self.failure_status_tbl.clone()
394//         };
395
396//         Self {
397//             failure_status_tbl,
398//             success: self.success && b,
399//             ..self.clone()
400//         }
401//     }
402// }
403
404#[derive(Debug, Clone, PartialEq, Eq)]
405pub struct LocalState {
406    pub stack_frame: Fp,
407    pub call_stack: Fp,
408    pub transaction_commitment: Fp,
409    pub full_transaction_commitment: Fp,
410    pub excess: Signed<Amount>,
411    pub supply_increase: Signed<Amount>,
412    pub ledger: Fp,
413    pub success: bool,
414    pub account_update_index: Index,
415    pub failure_status_tbl: Vec<Vec<TransactionFailure>>,
416    pub will_succeed: bool,
417}
418
419impl ToInputs for LocalState {
420    /// <https://github.com/MinaProtocol/mina/blob/4e0b324912017c3ff576704ee397ade3d9bda412/src/lib/mina_state/local_state.ml#L116>
421    fn to_inputs(&self, inputs: &mut Inputs) {
422        let Self {
423            stack_frame,
424            call_stack,
425            transaction_commitment,
426            full_transaction_commitment,
427            excess,
428            supply_increase,
429            ledger,
430            success,
431            account_update_index,
432            failure_status_tbl: _,
433            will_succeed,
434        } = self;
435
436        inputs.append(stack_frame);
437        inputs.append(call_stack);
438        inputs.append(transaction_commitment);
439        inputs.append(full_transaction_commitment);
440        inputs.append(excess);
441        inputs.append(supply_increase);
442        inputs.append(ledger);
443        inputs.append(account_update_index);
444        inputs.append(success);
445        inputs.append(will_succeed);
446    }
447}
448
449impl LocalState {
450    /// <https://github.com/MinaProtocol/mina/blob/436023ba41c43a50458a551b7ef7a9ae61670b25/src/lib/mina_state/local_state.ml#L65>
451    pub fn dummy() -> Self {
452        Self {
453            stack_frame: StackFrame::empty().hash(),
454            call_stack: Fp::zero(),
455            transaction_commitment: Fp::zero(),
456            full_transaction_commitment: Fp::zero(),
457            excess: Signed::<Amount>::zero(),
458            supply_increase: Signed::<Amount>::zero(),
459            ledger: Fp::zero(),
460            success: true,
461            account_update_index: <Index as Magnitude>::zero(),
462            failure_status_tbl: Vec::new(),
463            will_succeed: true,
464        }
465    }
466
467    pub fn empty() -> Self {
468        Self::dummy()
469    }
470
471    pub fn equal_without_ledger(&self, other: &Self) -> bool {
472        let Self {
473            stack_frame,
474            call_stack,
475            transaction_commitment,
476            full_transaction_commitment,
477            excess,
478            supply_increase,
479            ledger: _,
480            success,
481            account_update_index,
482            failure_status_tbl,
483            will_succeed,
484        } = self;
485
486        stack_frame == &other.stack_frame
487            && call_stack == &other.call_stack
488            && transaction_commitment == &other.transaction_commitment
489            && full_transaction_commitment == &other.full_transaction_commitment
490            && excess == &other.excess
491            && supply_increase == &other.supply_increase
492            && success == &other.success
493            && account_update_index == &other.account_update_index
494            && failure_status_tbl == &other.failure_status_tbl
495            && will_succeed == &other.will_succeed
496    }
497
498    pub fn checked_equal_prime(&self, other: &Self, w: &mut Witness<Fp>) -> [Boolean; 11] {
499        let Self {
500            stack_frame,
501            call_stack,
502            transaction_commitment,
503            full_transaction_commitment,
504            excess,
505            supply_increase,
506            ledger,
507            success,
508            account_update_index,
509            failure_status_tbl: _,
510            will_succeed,
511        } = self;
512
513        // { stack_frame : 'stack_frame
514        // ; call_stack : 'call_stack
515        // ; transaction_commitment : 'comm
516        // ; full_transaction_commitment : 'comm
517        // ; excess : 'signed_amount
518        // ; supply_increase : 'signed_amount
519        // ; ledger : 'ledger
520        // ; success : 'bool
521        // ; account_update_index : 'length
522        // ; failure_status_tbl : 'failure_status_tbl
523        // ; will_succeed : 'bool
524        // }
525
526        let mut alls = [
527            field::equal(*stack_frame, other.stack_frame, w),
528            field::equal(*call_stack, other.call_stack, w),
529            field::equal(*transaction_commitment, other.transaction_commitment, w),
530            field::equal(
531                *full_transaction_commitment,
532                other.full_transaction_commitment,
533                w,
534            ),
535            excess
536                .to_checked::<Fp>()
537                .equal(&other.excess.to_checked(), w),
538            supply_increase
539                .to_checked::<Fp>()
540                .equal(&other.supply_increase.to_checked(), w),
541            field::equal(*ledger, other.ledger, w),
542            success.to_boolean().equal(&other.success.to_boolean(), w),
543            account_update_index
544                .to_checked::<Fp>()
545                .equal(&other.account_update_index.to_checked(), w),
546            Boolean::True,
547            will_succeed
548                .to_boolean()
549                .equal(&other.will_succeed.to_boolean(), w),
550        ];
551        alls.reverse();
552        alls
553    }
554}
555
556fn step_all<A, L>(
557    _constraint_constants: &ConstraintConstants,
558    f: &impl Fn(&mut A, &GlobalState<L>, &LocalStateEnv<L>),
559    user_acc: &mut A,
560    (g_state, l_state): (&mut GlobalState<L>, &mut LocalStateEnv<L>),
561) -> Result<Vec<Vec<TransactionFailure>>, String>
562where
563    L: LedgerNonSnark,
564{
565    while !l_state.stack_frame.calls.is_empty() {
566        zkapps::non_snark::step(g_state, l_state)?;
567        f(user_acc, g_state, l_state);
568    }
569    Ok(l_state.failure_status_tbl.clone())
570}
571
572/// apply zkapp command fee payer's while stubbing out the second pass ledger
573/// CAUTION: If you use the intermediate local states, you MUST update the
574/// [`LocalStateEnv::will_succeed`] field to `false` if the `status` is [`TransactionStatus::Failed`].*)
575pub fn apply_zkapp_command_first_pass_aux<A, F, L>(
576    constraint_constants: &ConstraintConstants,
577    global_slot: Slot,
578    state_view: &ProtocolStateView,
579    init: &mut A,
580    f: F,
581    fee_excess: Option<Signed<Amount>>,
582    supply_increase: Option<Signed<Amount>>,
583    ledger: &mut L,
584    command: &ZkAppCommand,
585) -> Result<ZkappCommandPartiallyApplied<L>, String>
586where
587    L: LedgerNonSnark,
588    F: Fn(&mut A, &GlobalState<L>, &LocalStateEnv<L>),
589{
590    let fee_excess = fee_excess.unwrap_or_else(Signed::zero);
591    let supply_increase = supply_increase.unwrap_or_else(Signed::zero);
592
593    let previous_hash = ledger.merkle_root();
594    let original_first_pass_account_states = {
595        let id = command.fee_payer();
596        let location = {
597            let loc = ledger.location_of_account(&id);
598            let account = loc.as_ref().and_then(|loc| ledger.get(loc));
599            loc.zip(account)
600        };
601
602        vec![(id, location)]
603    };
604    // let perform = |eff: Eff<L>| Env::perform(eff);
605
606    let (mut global_state, mut local_state) = (
607        GlobalState {
608            protocol_state: state_view.clone(),
609            first_pass_ledger: ledger.clone(),
610            second_pass_ledger: {
611                // We stub out the second_pass_ledger initially, and then poke the
612                // correct value in place after the first pass is finished.
613                <L as LedgerIntf>::empty(0)
614            },
615            fee_excess,
616            supply_increase,
617            block_global_slot: global_slot,
618        },
619        LocalStateEnv {
620            stack_frame: StackFrame::default(),
621            call_stack: CallStack::new(),
622            transaction_commitment: Fp::zero(),
623            full_transaction_commitment: Fp::zero(),
624            excess: Signed::<Amount>::zero(),
625            supply_increase,
626            ledger: <L as LedgerIntf>::empty(0),
627            success: true,
628            account_update_index: IndexInterface::zero(),
629            failure_status_tbl: Vec::new(),
630            will_succeed: true,
631        },
632    );
633
634    f(init, &global_state, &local_state);
635    let account_updates = command.all_account_updates();
636
637    zkapps::non_snark::start(
638        &mut global_state,
639        &mut local_state,
640        zkapps::non_snark::StartData {
641            account_updates,
642            memo_hash: command.memo.hash(),
643            // It's always valid to set this value to true, and it will
644            // have no effect outside of the snark.
645            will_succeed: true,
646        },
647    )?;
648
649    let command = command.clone();
650    let constraint_constants = constraint_constants.clone();
651    let state_view = state_view.clone();
652
653    let res = ZkappCommandPartiallyApplied {
654        command,
655        previous_hash,
656        original_first_pass_account_states,
657        constraint_constants,
658        state_view,
659        global_state,
660        local_state,
661    };
662
663    Ok(res)
664}
665
666pub fn apply_zkapp_command_first_pass<L>(
667    constraint_constants: &ConstraintConstants,
668    global_slot: Slot,
669    state_view: &ProtocolStateView,
670    fee_excess: Option<Signed<Amount>>,
671    supply_increase: Option<Signed<Amount>>,
672    ledger: &mut L,
673    command: &ZkAppCommand,
674) -> Result<ZkappCommandPartiallyApplied<L>, String>
675where
676    L: LedgerNonSnark,
677{
678    let mut acc = ();
679    let partial_stmt = apply_zkapp_command_first_pass_aux(
680        constraint_constants,
681        global_slot,
682        state_view,
683        &mut acc,
684        |_acc, _g, _l| {},
685        fee_excess,
686        supply_increase,
687        ledger,
688        command,
689    )?;
690
691    Ok(partial_stmt)
692}
693
694pub fn apply_zkapp_command_second_pass_aux<A, F, L>(
695    constraint_constants: &ConstraintConstants,
696    init: &mut A,
697    f: F,
698    ledger: &mut L,
699    c: ZkappCommandPartiallyApplied<L>,
700) -> Result<ZkappCommandApplied, String>
701where
702    L: LedgerNonSnark,
703    F: Fn(&mut A, &GlobalState<L>, &LocalStateEnv<L>),
704{
705    // let perform = |eff: Eff<L>| Env::perform(eff);
706
707    let original_account_states: Vec<(AccountId, Option<_>)> = {
708        // get the original states of all the accounts in each pass.
709        // If an account updated in the first pass is referenced in account
710        // updates, then retain the value before first pass application*)
711
712        let accounts_referenced = c.command.accounts_referenced();
713
714        let mut account_states = BTreeMap::<AccountIdOrderable, Option<_>>::new();
715
716        let referenced = accounts_referenced.into_iter().map(|id| {
717            let location = {
718                let loc = ledger.location_of_account(&id);
719                let account = loc.as_ref().and_then(|loc| ledger.get(loc));
720                loc.zip(account)
721            };
722            (id, location)
723        });
724
725        c.original_first_pass_account_states
726            .into_iter()
727            .chain(referenced)
728            .for_each(|(id, acc_opt)| {
729                use std::collections::btree_map::Entry::Vacant;
730
731                let id_with_order: AccountIdOrderable = id.into();
732                if let Vacant(entry) = account_states.entry(id_with_order) {
733                    entry.insert(acc_opt);
734                };
735            });
736
737        account_states
738            .into_iter()
739            // Convert back the `AccountIdOrder` into `AccountId`, now that they are sorted
740            .map(|(id, account): (AccountIdOrderable, Option<_>)| (id.into(), account))
741            .collect()
742    };
743
744    let mut account_states_after_fee_payer = {
745        // To check if the accounts remain unchanged in the event the transaction
746        // fails. First pass updates will remain even if the transaction fails to
747        // apply zkapp account updates*)
748
749        c.command.accounts_referenced().into_iter().map(|id| {
750            let loc = ledger.location_of_account(&id);
751            let a = loc.as_ref().and_then(|loc| ledger.get(loc));
752
753            match a {
754                Some(a) => (id, Some((loc.unwrap(), a))),
755                None => (id, None),
756            }
757        })
758    };
759
760    let accounts = || {
761        original_account_states
762            .iter()
763            .map(|(id, account)| (id.clone(), account.as_ref().map(|(_loc, acc)| acc.clone())))
764            .collect::<Vec<_>>()
765    };
766
767    // Warning(OCaml): This is an abstraction leak / hack.
768    // Here, we update global second pass ledger to be the input ledger, and
769    // then update the local ledger to be the input ledger *IF AND ONLY IF*
770    // there are more transaction segments to be processed in this pass.
771
772    // TODO(OCaml): Remove this, and uplift the logic into the call in staged ledger.
773
774    let mut global_state = GlobalState {
775        second_pass_ledger: ledger.clone(),
776        ..c.global_state
777    };
778
779    let mut local_state = {
780        if c.local_state.stack_frame.calls.is_empty() {
781            // Don't mess with the local state; we've already finished the
782            // transaction after the fee payer.
783            c.local_state
784        } else {
785            // Install the ledger that should already be in the local state, but
786            // may not be in some situations depending on who the caller is.
787            LocalStateEnv {
788                ledger: global_state.second_pass_ledger(),
789                ..c.local_state
790            }
791        }
792    };
793
794    f(init, &global_state, &local_state);
795    let start = (&mut global_state, &mut local_state);
796
797    let reversed_failure_status_tbl = step_all(constraint_constants, &f, init, start)?;
798
799    let failure_status_tbl = reversed_failure_status_tbl
800        .into_iter()
801        .rev()
802        .collect::<Vec<_>>();
803
804    let account_ids_originally_not_in_ledger =
805        original_account_states
806            .iter()
807            .filter_map(|(acct_id, loc_and_acct)| {
808                if loc_and_acct.is_none() {
809                    Some(acct_id)
810                } else {
811                    None
812                }
813            });
814
815    let successfully_applied = failure_status_tbl.concat().is_empty();
816
817    // if the zkapp command fails in at least 1 account update,
818    // then all the account updates would be cancelled except
819    // the fee payer one
820    let failure_status_tbl = if successfully_applied {
821        failure_status_tbl
822    } else {
823        failure_status_tbl
824            .into_iter()
825            .enumerate()
826            .map(|(idx, fs)| {
827                if idx > 0 && fs.is_empty() {
828                    vec![TransactionFailure::Cancelled]
829                } else {
830                    fs
831                }
832            })
833            .collect()
834    };
835
836    // accounts not originally in ledger, now present in ledger
837    let new_accounts = account_ids_originally_not_in_ledger
838        .filter(|acct_id| ledger.location_of_account(acct_id).is_some())
839        .cloned()
840        .collect::<Vec<_>>();
841
842    let new_accounts_is_empty = new_accounts.is_empty();
843
844    let valid_result = Ok(ZkappCommandApplied {
845        accounts: accounts(),
846        command: WithStatus {
847            data: c.command,
848            status: if successfully_applied {
849                TransactionStatus::Applied
850            } else {
851                TransactionStatus::Failed(failure_status_tbl)
852            },
853        },
854        new_accounts,
855    });
856
857    if successfully_applied {
858        valid_result
859    } else {
860        let other_account_update_accounts_unchanged = account_states_after_fee_payer
861            .fold_while(true, |acc, (_, loc_opt)| match loc_opt {
862                Some((loc, a)) => match ledger.get(&loc) {
863                    Some(a_) if !(a == a_) => FoldWhile::Done(false),
864                    _ => FoldWhile::Continue(acc),
865                },
866                _ => FoldWhile::Continue(acc),
867            })
868            .into_inner();
869
870        // Other zkapp_command failed, therefore, updates in those should not get applied
871        if new_accounts_is_empty && other_account_update_accounts_unchanged {
872            valid_result
873        } else {
874            Err("Zkapp_command application failed but new accounts created or some of the other account_update updates applied".to_string())
875        }
876    }
877}
878
879pub fn apply_zkapp_command_second_pass<L>(
880    constraint_constants: &ConstraintConstants,
881    ledger: &mut L,
882    c: ZkappCommandPartiallyApplied<L>,
883) -> Result<ZkappCommandApplied, String>
884where
885    L: LedgerNonSnark,
886{
887    let x = apply_zkapp_command_second_pass_aux(
888        constraint_constants,
889        &mut (),
890        |_, _, _| {},
891        ledger,
892        c,
893    )?;
894    Ok(x)
895}
896
897fn apply_zkapp_command_unchecked_aux<A, F, L>(
898    constraint_constants: &ConstraintConstants,
899    global_slot: Slot,
900    state_view: &ProtocolStateView,
901    init: &mut A,
902    f: F,
903    fee_excess: Option<Signed<Amount>>,
904    supply_increase: Option<Signed<Amount>>,
905    ledger: &mut L,
906    command: &ZkAppCommand,
907) -> Result<ZkappCommandApplied, String>
908where
909    L: LedgerNonSnark,
910    F: Fn(&mut A, &GlobalState<L>, &LocalStateEnv<L>),
911{
912    let partial_stmt = apply_zkapp_command_first_pass_aux(
913        constraint_constants,
914        global_slot,
915        state_view,
916        init,
917        &f,
918        fee_excess,
919        supply_increase,
920        ledger,
921        command,
922    )?;
923
924    apply_zkapp_command_second_pass_aux(constraint_constants, init, &f, ledger, partial_stmt)
925}
926
927fn apply_zkapp_command_unchecked<L>(
928    constraint_constants: &ConstraintConstants,
929    global_slot: Slot,
930    state_view: &ProtocolStateView,
931    ledger: &mut L,
932    command: &ZkAppCommand,
933) -> Result<(ZkappCommandApplied, (LocalStateEnv<L>, Signed<Amount>)), String>
934where
935    L: LedgerNonSnark,
936{
937    let zkapp_partially_applied: ZkappCommandPartiallyApplied<L> = apply_zkapp_command_first_pass(
938        constraint_constants,
939        global_slot,
940        state_view,
941        None,
942        None,
943        ledger,
944        command,
945    )?;
946
947    let mut state_res = None;
948    let account_update_applied = apply_zkapp_command_second_pass_aux(
949        constraint_constants,
950        &mut state_res,
951        |acc, global_state, local_state| {
952            *acc = Some((local_state.clone(), global_state.fee_excess))
953        },
954        ledger,
955        zkapp_partially_applied,
956    )?;
957    let (state, amount) = state_res.unwrap();
958
959    Ok((account_update_applied, (state.clone(), amount)))
960}