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