mina_tree/scan_state/
pending_coinbase.rs

1///
2/// Pending_coinbase is to keep track of all the coinbase transactions that have been
3/// applied to the ledger but for which there is no ledger proof yet. Every ledger
4/// proof corresponds to a sequence of coinbase transactions which is part of all the
5/// transactions it proves. Each of these sequences[Stack] are stored using the merkle
6/// tree representation. The stacks are operated in a FIFO manner by keeping track of
7/// its positions in the merkle tree. Whenever a ledger proof is emitted, the oldest
8/// stack is removed from the tree and when a new coinbase is applied, the latest stack
9/// is updated with the new coinbase.
10///
11/// The operations on the merkle tree of coinbase stacks include:
12/// 1) adding a new singleton stack
13/// 2) updating the latest stack when a new coinbase is added to it
14/// 2) deleting the oldest stack
15///
16/// A stack can be either be created or modified by pushing a coinbase on to it.
17///
18/// This module also provides an interface for the checked computations required required to prove it in snark
19///
20/// Stack operations are done for transaction snarks and tree operations are done for the blockchain snark*)
21use std::{collections::HashMap, fmt::Write, marker::PhantomData};
22
23use ark_ff::{fields::arithmetic::InvalidBigInt, Zero};
24use mina_curves::pasta::Fp;
25use mina_signer::CompressedPubKey;
26use openmina_core::constants::constraint_constants;
27use poseidon::hash::{
28    hash_noinputs, hash_with_kimchi,
29    params::{
30        get_coinbase_param_for_height, COINBASE_STACK, MINA_PROTO_STATE, NO_INPUT_COINBASE_STACK,
31    },
32    Inputs,
33};
34use sha2::{Digest, Sha256};
35
36use crate::{
37    proofs::{
38        field::{field, Boolean},
39        numbers::{
40            currency::{CheckedAmount, CheckedCurrency},
41            nat::{CheckedNat, CheckedSlot},
42        },
43        transaction::transaction_snark::checked_hash,
44        witness::Witness,
45    },
46    staged_ledger::hash::PendingCoinbaseAux,
47    Address, AppendToInputs as _, MerklePath, ToInputs,
48};
49
50use self::merkle_tree::MiniMerkleTree;
51
52use super::{
53    currency::{Amount, Magnitude, Slot},
54    transaction_logic::Coinbase,
55};
56
57#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
58pub struct StackId(u64);
59
60impl std::fmt::Debug for StackId {
61    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
62        f.write_fmt(format_args!("StackId({})", self.0))
63    }
64}
65
66impl StackId {
67    pub fn incr_by_one(&self) -> Self {
68        self.0
69            .checked_add(1)
70            .map(Self)
71            .ok_or_else(|| "Stack_id overflow".to_string())
72            .unwrap()
73    }
74
75    pub fn zero() -> Self {
76        Self(0)
77    }
78
79    pub(super) fn new(number: u64) -> Self {
80        Self(number)
81    }
82
83    pub(super) fn as_u64(&self) -> u64 {
84        self.0
85    }
86}
87
88struct CoinbaseData {
89    receiver: CompressedPubKey,
90    amount: Amount,
91}
92
93impl ToInputs for CoinbaseData {
94    fn to_inputs(&self, inputs: &mut Inputs) {
95        let Self { receiver, amount } = self;
96        inputs.append(receiver);
97        inputs.append(amount);
98    }
99}
100
101impl CoinbaseData {
102    pub fn empty() -> Self {
103        Self {
104            receiver: CompressedPubKey::empty(),
105            amount: Amount::zero(),
106        }
107    }
108
109    pub fn of_coinbase(cb: Coinbase) -> Self {
110        let Coinbase {
111            receiver,
112            amount,
113            fee_transfer: _,
114        } = cb;
115        Self { receiver, amount }
116    }
117
118    pub fn genesis() -> Self {
119        Self::empty()
120    }
121}
122
123#[derive(Clone, Debug)]
124pub struct StackState {
125    pub source: Stack,
126    pub target: Stack,
127}
128
129#[derive(Clone, PartialEq, Eq)]
130pub struct CoinbaseStack(pub Fp);
131
132impl ToInputs for CoinbaseStack {
133    fn to_inputs(&self, inputs: &mut Inputs) {
134        let Self(fp) = self;
135
136        inputs.append(fp)
137    }
138}
139
140impl std::fmt::Debug for CoinbaseStack {
141    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
142        if self.is_empty() {
143            f.write_fmt(format_args!("CoinbaseStack(Empty)"))
144        } else {
145            f.debug_tuple("CoinbaseStack").field(&self.0).finish()
146        }
147    }
148}
149
150impl CoinbaseStack {
151    /// <https://github.com/MinaProtocol/mina/blob/2ee6e004ba8c6a0541056076aab22ea162f7eb3a/src/lib/mina_base/pending_coinbase.ml#L180>
152    pub fn push(&self, cb: Coinbase) -> Self {
153        let mut inputs = Inputs::new();
154
155        inputs.append(&CoinbaseData::of_coinbase(cb));
156        inputs.append_field(self.0);
157
158        let hash = hash_with_kimchi(&COINBASE_STACK, &inputs.to_fields());
159        Self(hash)
160    }
161
162    pub fn checked_push(&self, cb: Coinbase, w: &mut Witness<Fp>) -> Self {
163        let mut inputs = Inputs::new();
164
165        inputs.append(&CoinbaseData::of_coinbase(cb));
166        inputs.append_field(self.0);
167
168        let hash = checked_hash(&COINBASE_STACK, &inputs.to_fields(), w);
169        Self(hash)
170    }
171
172    fn check_merge(
173        (_, t1): (&Self, &Self),
174        (s2, _): (&Self, &Self),
175        w: &mut Witness<Fp>,
176    ) -> Boolean {
177        field::equal(t1.0, s2.0, w)
178    }
179
180    /// <https://github.com/MinaProtocol/mina/blob/2ee6e004ba8c6a0541056076aab22ea162f7eb3a/src/lib/mina_base/pending_coinbase.ml#L188>
181    pub fn empty() -> Self {
182        Self(hash_noinputs(&NO_INPUT_COINBASE_STACK))
183    }
184
185    /// Used for tests/debug only
186    fn is_empty(&self) -> bool {
187        self == &Self::empty()
188    }
189}
190
191type StackHash = Fp;
192
193#[derive(Clone, PartialEq, Eq)]
194pub struct StateStack {
195    pub init: StackHash,
196    pub curr: StackHash,
197}
198
199impl ToInputs for StateStack {
200    /// <https://github.com/MinaProtocol/mina/blob/4e0b324912017c3ff576704ee397ade3d9bda412/src/lib/mina_base/pending_coinbase.ml#L271>
201    fn to_inputs(&self, inputs: &mut Inputs) {
202        let Self { init, curr } = self;
203        inputs.append(init);
204        inputs.append(curr);
205    }
206}
207
208impl std::fmt::Debug for StateStack {
209    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
210        if self.is_empty() {
211            f.write_fmt(format_args!("StateStack(Empty)"))
212        } else {
213            f.debug_struct("StateStack")
214                .field("init", &self.init)
215                .field("curr", &self.curr)
216                .finish()
217        }
218    }
219}
220
221impl StateStack {
222    fn push(&self, state_body_hash: Fp, global_slot: Slot) -> Self {
223        let mut inputs = Inputs::new();
224
225        inputs.append_field(self.curr);
226        inputs.append_field(state_body_hash);
227        inputs.append_field(global_slot.to_field());
228
229        let hash = hash_with_kimchi(&MINA_PROTO_STATE, &inputs.to_fields());
230
231        Self {
232            init: self.init,
233            curr: hash,
234        }
235    }
236
237    fn checked_push(
238        &self,
239        state_body_hash: Fp,
240        global_slot: CheckedSlot<Fp>,
241        w: &mut Witness<Fp>,
242    ) -> Self {
243        let mut inputs = Inputs::new();
244
245        inputs.append_field(self.curr);
246        inputs.append_field(state_body_hash);
247        inputs.append_field(global_slot.to_field());
248
249        let hash = checked_hash(&MINA_PROTO_STATE, &inputs.to_fields(), w);
250
251        Self {
252            init: self.init,
253            curr: hash,
254        }
255    }
256
257    fn equal_var(&self, other: &Self, w: &mut Witness<Fp>) -> Boolean {
258        let b1 = field::equal(self.init, other.init, w);
259        let b2 = field::equal(self.curr, other.curr, w);
260        b1.and(&b2, w)
261    }
262
263    fn check_merge(
264        (s1, t1): (&Self, &Self),
265        (s2, t2): (&Self, &Self),
266        w: &mut Witness<Fp>,
267    ) -> Boolean {
268        let eq_src = s1.equal_var(s2, w);
269        let eq_target = t1.equal_var(t2, w);
270        let correct_transition = t1.equal_var(s2, w);
271        let same_update = eq_src.and(&eq_target, w);
272        Boolean::any(&[same_update, correct_transition], w)
273    }
274
275    fn empty() -> Self {
276        Self {
277            init: Fp::zero(),
278            curr: Fp::zero(),
279        }
280    }
281
282    /// Used for tests/debug only
283    fn is_empty(&self) -> bool {
284        self.curr.is_zero() && self.init.is_zero()
285    }
286
287    fn create(init: StackHash) -> Self {
288        Self { init, curr: init }
289    }
290}
291
292pub mod update {
293    use crate::scan_state::currency::{Amount, Magnitude};
294
295    #[derive(Clone, Debug)]
296    pub enum Action {
297        None,
298        One,
299        TwoCoinbaseInFirst,
300        TwoCoinbaseInSecond,
301    }
302
303    #[derive(Debug)]
304    pub enum StackUpdate {
305        None,
306        One(super::Stack),
307        Two((super::Stack, super::Stack)),
308    }
309
310    #[derive(Clone, Debug)]
311    pub struct Update {
312        pub action: Action,
313        pub coinbase_amount: Amount,
314    }
315
316    impl Update {
317        fn genesis() -> Self {
318            Self {
319                action: Action::None,
320                coinbase_amount: Amount::zero(),
321            }
322        }
323    }
324}
325
326#[derive(Clone, PartialEq, Eq)]
327pub struct Stack {
328    pub data: CoinbaseStack,
329    pub state: StateStack,
330}
331
332impl ToInputs for Stack {
333    /// <https://github.com/MinaProtocol/mina/blob/4e0b324912017c3ff576704ee397ade3d9bda412/src/lib/mina_base/pending_coinbase.ml#L591>
334    fn to_inputs(&self, inputs: &mut Inputs) {
335        let Self { data, state } = self;
336        inputs.append(data);
337        inputs.append(state);
338    }
339}
340
341impl std::fmt::Debug for Stack {
342    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
343        if self.data.is_empty() && self.state.is_empty() {
344            f.write_fmt(format_args!("Stack(Empty)"))
345        } else {
346            f.debug_struct("Stack")
347                .field("data", &self.data)
348                .field("state", &self.state)
349                .finish()
350        }
351    }
352}
353
354impl Stack {
355    pub fn empty() -> Self {
356        Self {
357            data: CoinbaseStack::empty(),
358            state: StateStack::empty(),
359        }
360    }
361
362    pub fn push_coinbase(&self, cb: Coinbase) -> Self {
363        Self {
364            data: self.data.push(cb),
365            state: self.state.clone(),
366        }
367    }
368
369    pub fn push_state(&self, state_body_hash: Fp, global_slot: Slot) -> Self {
370        Self {
371            data: self.data.clone(),
372            state: self.state.push(state_body_hash, global_slot),
373        }
374    }
375
376    pub fn checked_push_coinbase(&self, cb: Coinbase, w: &mut Witness<Fp>) -> Self {
377        Self {
378            data: self.data.checked_push(cb, w),
379            state: self.state.clone(),
380        }
381    }
382
383    pub fn checked_push_state(
384        &self,
385        state_body_hash: Fp,
386        global_slot: CheckedSlot<Fp>,
387        w: &mut Witness<Fp>,
388    ) -> Self {
389        Self {
390            data: self.data.clone(),
391            state: self.state.checked_push(state_body_hash, global_slot, w),
392        }
393    }
394
395    pub fn equal_var(&self, other: &Self, w: &mut Witness<Fp>) -> Boolean {
396        let b1 = field::equal(self.data.0, other.data.0, w);
397        let b2 = {
398            let b1 = field::equal(self.state.init, other.state.init, w);
399            let b2 = field::equal(self.state.curr, other.state.curr, w);
400            b1.and(&b2, w)
401        };
402        b1.and(&b2, w)
403    }
404
405    pub fn check_merge(
406        transition1: (&Self, &Self),
407        transition2: (&Self, &Self),
408        w: &mut Witness<Fp>,
409    ) -> Boolean {
410        let (s, t) = transition1;
411        let (s2, t2) = transition2;
412
413        let valid_coinbase_stacks =
414            CoinbaseStack::check_merge((&s.data, &t.data), (&s2.data, &t2.data), w);
415        let valid_state_stacks =
416            StateStack::check_merge((&s.state, &t.state), (&s2.state, &t2.state), w);
417
418        valid_coinbase_stacks.and(&valid_state_stacks, w)
419    }
420
421    /// <https://github.com/MinaProtocol/mina/blob/05c2f73d0f6e4f1341286843814ce02dcb3919e0/src/lib/mina_base/pending_coinbase.ml#L651>
422    pub fn create_with(other: &Self) -> Self {
423        Self {
424            state: StateStack::create(other.state.curr),
425            ..Self::empty()
426        }
427    }
428
429    /// <https://github.com/MinaProtocol/mina/blob/f5b013880dede0e2ef04cebf4b0213b850a85548/src/lib/mina_base/pending_coinbase.ml#L738>
430    pub fn var_create_with(other: &Self) -> Self {
431        // Note: Here we use `init`
432        Self {
433            state: StateStack::create(other.state.init),
434            ..Self::empty()
435        }
436    }
437
438    /// <https://github.com/MinaProtocol/mina/blob/2ee6e004ba8c6a0541056076aab22ea162f7eb3a/src/lib/mina_base/pending_coinbase.ml#L658>
439    pub fn connected(first: &Self, second: &Self, prev: Option<&Self>) -> bool {
440        // same as old stack or second could be a new stack with empty data
441        let coinbase_stack_connected =
442            (first.data == second.data) || { second.data == CoinbaseStack::empty() };
443
444        // 1. same as old stack or
445        // 2. new stack initialized with the stack state of last block. Not possible to know this unless we track
446        //    all the stack states because they are updated once per block (init=curr)
447        // 3. [second] could be a new stack initialized with the latest state of [first] or
448        // 4. [second] starts from the previous state of [first]. This is not available in either [first] or [second] *)
449        let state_stack_connected = first.state == second.state
450            || second.state.init == second.state.curr
451            || first.state.curr == second.state.curr
452            || prev
453                .map(|prev| prev.state.curr == second.state.curr)
454                .unwrap_or(true);
455
456        coinbase_stack_connected && state_stack_connected
457    }
458
459    fn hash_var(&self, w: &mut Witness<Fp>) -> Fp {
460        checked_hash(&COINBASE_STACK, &self.to_inputs_owned().to_fields(), w)
461    }
462}
463
464#[derive(Clone, Debug)]
465pub struct PendingCoinbase {
466    pub(super) tree: merkle_tree::MiniMerkleTree<StackId, Stack, StackHasher>,
467    pub(super) pos_list: Vec<StackId>,
468    pub(super) new_pos: StackId,
469}
470
471#[derive(Clone, Debug)]
472pub(super) struct StackHasher;
473
474impl merkle_tree::TreeHasher<Stack> for StackHasher {
475    fn hash_value(value: &Stack) -> Fp {
476        let mut inputs = Inputs::new();
477
478        inputs.append_field(value.data.0);
479
480        inputs.append_field(value.state.init);
481        inputs.append_field(value.state.curr);
482
483        hash_with_kimchi(&COINBASE_STACK, &inputs.to_fields())
484    }
485
486    fn merge_hash(height: usize, left: Fp, right: Fp) -> Fp {
487        let param = get_coinbase_param_for_height(height);
488        poseidon::hash::hash_with_kimchi(param, &[left, right])
489    }
490
491    fn empty_value() -> Stack {
492        Stack::empty()
493    }
494}
495
496impl PendingCoinbase {
497    pub fn create(depth: usize) -> Self {
498        let mut tree = MiniMerkleTree::create(depth);
499
500        let nstacks = 2u64.pow(depth as u32);
501        let mut stack_id = StackId::zero();
502
503        assert_eq!(depth, 5);
504
505        tree.fill_with((0..nstacks).map(|_| {
506            let this_id = stack_id;
507            stack_id = stack_id.incr_by_one();
508            (this_id, Stack::empty())
509        }));
510
511        Self {
512            tree,
513            pos_list: Vec::with_capacity(128),
514            new_pos: StackId::zero(),
515        }
516    }
517
518    pub fn merkle_root(&mut self) -> Fp {
519        self.tree.merkle_root()
520    }
521
522    fn get_stack(&self, addr: Address) -> &Stack {
523        self.tree.get_exn(addr)
524    }
525
526    fn path(&mut self, addr: Address) -> Vec<MerklePath> {
527        self.tree.path_exn(addr)
528    }
529
530    fn find_index(&self, key: StackId) -> Address {
531        self.tree.find_index_exn(key)
532    }
533
534    fn next_index(&self, depth: usize) -> StackId {
535        if self.new_pos.0 == (2u64.pow(depth as u32) - 1) {
536            StackId::zero()
537        } else {
538            self.new_pos.incr_by_one()
539        }
540    }
541
542    fn next_stack_id(&self, depth: usize, is_new_stack: bool) -> StackId {
543        if is_new_stack {
544            self.next_index(depth)
545        } else {
546            self.new_pos
547        }
548    }
549
550    fn incr_index(&mut self, depth: usize, is_new_stack: bool) {
551        if is_new_stack {
552            let new_pos = self.next_index(depth);
553            self.pos_list.push(self.new_pos);
554            self.new_pos = new_pos;
555        }
556    }
557
558    fn set_stack(&mut self, depth: usize, addr: Address, stack: Stack, is_new_stack: bool) {
559        self.tree.set_exn(addr, stack);
560        self.incr_index(depth, is_new_stack);
561    }
562
563    fn latest_stack_id(&self, is_new_stack: bool) -> StackId {
564        if is_new_stack {
565            self.new_pos
566        } else {
567            self.pos_list.last().cloned().unwrap_or_else(StackId::zero)
568        }
569    }
570
571    fn current_stack_id(&self) -> Option<StackId> {
572        self.pos_list.last().cloned()
573    }
574
575    fn current_stack(&self) -> &Stack {
576        let prev_stack_id = self.current_stack_id().unwrap_or_else(StackId::zero);
577        let addr = self.tree.find_index_exn(prev_stack_id);
578        self.tree.get_exn(addr)
579    }
580
581    pub fn latest_stack(&self, is_new_stack: bool) -> Stack {
582        let key = self.latest_stack_id(is_new_stack);
583        let addr = self.tree.find_index_exn(key);
584        let mut res = self.tree.get_exn(addr).clone();
585        if is_new_stack {
586            let prev_stack = self.current_stack();
587            res.state = StateStack::create(prev_stack.state.curr);
588        }
589        res
590    }
591
592    fn oldest_stack_id(&self) -> Option<StackId> {
593        self.pos_list.first().cloned()
594    }
595
596    fn remove_oldest_stack_id(&mut self) {
597        todo!()
598    }
599
600    fn oldest_stack(&self) -> &Stack {
601        let key = self.oldest_stack_id().unwrap_or_else(StackId::zero);
602        let addr = self.find_index(key);
603        self.get_stack(addr)
604    }
605
606    fn update_stack<F>(&mut self, depth: usize, is_new_stack: bool, fun: F)
607    where
608        F: FnOnce(&Stack) -> Stack,
609    {
610        let key = self.latest_stack_id(is_new_stack);
611        let stack_addr = self.find_index(key);
612        let stack_before = self.get_stack(stack_addr.clone());
613        let stack_after = fun(stack_before);
614        // state hash in "after" stack becomes previous state hash at top level
615        self.set_stack(depth, stack_addr, stack_after, is_new_stack);
616    }
617
618    fn add_coinbase(&mut self, depth: usize, coinbase: Coinbase, is_new_stack: bool) {
619        self.update_stack(depth, is_new_stack, |stack| stack.push_coinbase(coinbase))
620    }
621
622    fn add_state(
623        &mut self,
624        depth: usize,
625        state_body_hash: Fp,
626        global_slot: Slot,
627        is_new_stack: bool,
628    ) {
629        self.update_stack(depth, is_new_stack, |stack| {
630            stack.push_state(state_body_hash, global_slot)
631        })
632    }
633
634    pub fn update_coinbase_stack(
635        &mut self,
636        depth: usize,
637        stack: Stack,
638        is_new_stack: bool,
639    ) -> Result<(), String> {
640        self.update_stack(depth, is_new_stack, |_| stack);
641        Ok(())
642    }
643
644    pub fn remove_coinbase_stack(&mut self, depth: usize) -> Result<Stack, String> {
645        let oldest_stack_id = if !self.pos_list.is_empty() {
646            self.pos_list.remove(0) // TODO: Use `VecDeque`
647        } else {
648            return Err("No coinbase stack-with-state-hash to pop".to_string());
649        };
650        let stack_addr = self.find_index(oldest_stack_id);
651        let stack = self.get_stack(stack_addr.clone()).clone();
652        self.set_stack(depth, stack_addr, Stack::empty(), false);
653        Ok(stack)
654    }
655
656    pub fn hash_extra(&self) -> PendingCoinbaseAux {
657        let mut s = String::with_capacity(64 * 1024);
658        for pos in self.pos_list.iter().rev() {
659            write!(&mut s, "{}", pos.0).unwrap();
660        }
661
662        let mut sha = Sha256::new();
663        sha.update(s.as_bytes());
664
665        s.clear();
666        write!(&mut s, "{}", self.new_pos.0).unwrap();
667        sha.update(s);
668
669        let digest = sha.finalize();
670        PendingCoinbaseAux(digest.into())
671    }
672
673    pub fn pop_coinbases(
674        proof_emitted: Boolean,
675        pending_coinbase_witness: &mut PendingCoinbaseWitness,
676        w: &mut Witness<Fp>,
677    ) -> Result<(Fp, Stack), InvalidBigInt> {
678        let addr = w.exists(pending_coinbase_witness.find_index_of_oldest_stack());
679        let (prev, prev_path) = w.exists(pending_coinbase_witness.get_coinbase_stack(addr.clone()));
680
681        checked_verify_merkle_path(&prev, &prev_path, w);
682
683        let next = w.exists_no_check(match proof_emitted {
684            Boolean::True => Stack::empty(),
685            Boolean::False => prev.clone(),
686        });
687
688        pending_coinbase_witness.set_oldest_coinbase_stack(addr, next.clone());
689
690        // Note: in OCaml hashing of `next` is made before `set_oldest_coinbase_stack`
691        let new_root = checked_verify_merkle_path(&next, &prev_path, w);
692
693        Ok((new_root, prev))
694    }
695
696    pub fn add_coinbase_checked(
697        update: &update::Update,
698        coinbase_receiver: &CompressedPubKey,
699        supercharge_coinbase: Boolean,
700        state_body_hash: Fp,
701        global_slot: &CheckedSlot<Fp>,
702        pending_coinbase_witness: &mut PendingCoinbaseWitness,
703        w: &mut Witness<Fp>,
704    ) -> Result<Fp, InvalidBigInt> {
705        let no_update = |[b0, b1]: &[Boolean; 2], w: &mut Witness<Fp>| b0.neg().and(&b1.neg(), w);
706        let update_two_stacks_coinbase_in_first =
707            |[b0, b1]: &[Boolean; 2], w: &mut Witness<Fp>| b0.neg().and(b1, w);
708        let update_two_stacks_coinbase_in_second =
709            |[b0, b1]: &[Boolean; 2], w: &mut Witness<Fp>| b0.and(b1, w);
710
711        let update::Update {
712            action,
713            coinbase_amount: amount,
714        } = update;
715
716        let amount = Amount::from_u64(amount.as_u64()).to_checked();
717        let (addr1, addr2) = w.exists(pending_coinbase_witness.find_index_of_newest_stacks());
718
719        let action = {
720            use update::Action::*;
721            match action {
722                None => [Boolean::False, Boolean::False],
723                One => [Boolean::True, Boolean::False],
724                TwoCoinbaseInFirst => [Boolean::False, Boolean::True],
725                TwoCoinbaseInSecond => [Boolean::True, Boolean::True],
726            }
727        };
728
729        let no_update = no_update(&action, w);
730
731        let update_state_stack = |stack: Stack,
732                                  pending_coinbase_witness: &mut PendingCoinbaseWitness,
733                                  w: &mut Witness<Fp>| {
734            let previous_state_stack = w.exists(pending_coinbase_witness.get_previous_stack());
735            let stack_initialized = Stack {
736                state: previous_state_stack,
737                ..stack.clone()
738            };
739            let stack_with_state_hash =
740                stack_initialized.checked_push_state(state_body_hash, *global_slot, w);
741            w.exists_no_check(match no_update {
742                Boolean::True => stack,
743                Boolean::False => stack_with_state_hash,
744            })
745        };
746
747        let update_stack1 = |stack: Stack,
748                             pending_coinbase_witness: &mut PendingCoinbaseWitness,
749                             w: &mut Witness<Fp>| {
750            let stack = update_state_stack(stack, pending_coinbase_witness, w);
751            let total_coinbase_amount = {
752                let coinbase_amount = Amount::from_u64(constraint_constants().coinbase_amount);
753                let superchaged_coinbase = coinbase_amount
754                    .scale(constraint_constants().supercharged_coinbase_factor)
755                    .unwrap()
756                    .to_checked::<Fp>();
757                let coinbase_amount = coinbase_amount.to_checked::<Fp>();
758
759                match supercharge_coinbase {
760                    Boolean::True => superchaged_coinbase,
761                    Boolean::False => coinbase_amount,
762                }
763            };
764
765            let rem_amount = total_coinbase_amount.sub(&amount, w);
766            let no_coinbase_in_this_stack = update_two_stacks_coinbase_in_second(&action, w);
767
768            let amount1_equal_to_zero = amount.equal(&CheckedAmount::zero(), w);
769            let amount2_equal_to_zero = rem_amount.equal(&CheckedAmount::zero(), w);
770
771            no_update.equal(&amount1_equal_to_zero, w);
772
773            let no_coinbase = no_update.or(&no_coinbase_in_this_stack, w);
774
775            let stack_with_amount1 = stack.checked_push_coinbase(
776                Coinbase {
777                    receiver: coinbase_receiver.clone(),
778                    amount: amount.to_inner(), // TODO: Overflow ?
779                    fee_transfer: None,
780                },
781                w,
782            );
783
784            let stack_with_amount2 = stack_with_amount1.checked_push_coinbase(
785                Coinbase {
786                    receiver: coinbase_receiver.clone(),
787                    amount: rem_amount.to_inner(), // TODO: Overflow ?
788                    fee_transfer: None,
789                },
790                w,
791            );
792
793            let on_false = {
794                w.exists_no_check(match amount2_equal_to_zero {
795                    Boolean::True => stack_with_amount1,
796                    Boolean::False => stack_with_amount2,
797                })
798            };
799
800            w.exists_no_check(match no_coinbase {
801                Boolean::True => stack,
802                Boolean::False => on_false,
803            })
804        };
805
806        let update_stack2 = |init_stack: Stack, stack0: Stack, w: &mut Witness<Fp>| {
807            let add_coinbase = update_two_stacks_coinbase_in_second(&action, w);
808            let update_state = {
809                let update_second_stack = update_two_stacks_coinbase_in_first(&action, w);
810                update_second_stack.or(&add_coinbase, w)
811            };
812
813            let stack = {
814                let stack_with_state = Stack {
815                    state: StateStack::create(init_stack.state.curr),
816                    ..stack0.clone()
817                }
818                .checked_push_state(state_body_hash, *global_slot, w);
819                w.exists_no_check(match update_state {
820                    Boolean::True => stack_with_state,
821                    Boolean::False => stack0,
822                })
823            };
824
825            let stack_with_coinbase = stack.checked_push_coinbase(
826                Coinbase {
827                    receiver: coinbase_receiver.clone(),
828                    amount: amount.to_inner(), // TODO: Overflow ?
829                    fee_transfer: None,
830                },
831                w,
832            );
833
834            w.exists_no_check(match add_coinbase {
835                Boolean::True => stack_with_coinbase,
836                Boolean::False => stack,
837            })
838        };
839
840        let (_new_root, prev, _updated_stack1) = {
841            let (stack, path) = w.exists_no_check({
842                let pc = &mut pending_coinbase_witness.pending_coinbase;
843                let stack = pc.get_stack(addr1.clone()).clone();
844                let path = pc.path(addr1.clone());
845                (stack, path)
846            });
847            checked_verify_merkle_path(&stack, &path, w);
848
849            let next = update_stack1(stack.clone(), pending_coinbase_witness, w);
850
851            pending_coinbase_witness.set_coinbase_stack(addr1, next.clone());
852            let new_root = checked_verify_merkle_path(&next, &path, w);
853            (new_root, stack, next)
854        };
855
856        let (root, _, _) = {
857            let (stack, path) = w.exists_no_check({
858                let pc = &mut pending_coinbase_witness.pending_coinbase;
859                let stack = pc.get_stack(addr2.clone()).clone();
860                let path = pc.path(addr2.clone());
861                (stack, path)
862            });
863            checked_verify_merkle_path(&stack, &path, w);
864
865            let next = update_stack2(prev, stack.clone(), w);
866
867            pending_coinbase_witness.set_coinbase_stack(addr2, next.clone());
868            let new_root = checked_verify_merkle_path(&next, &path, w);
869            (new_root, stack, next)
870        };
871
872        Ok(root)
873    }
874}
875
876/// `implied_root` in OCaml
877pub fn checked_verify_merkle_path(
878    account: &Stack,
879    merkle_path: &[MerklePath],
880    w: &mut Witness<Fp>,
881) -> Fp {
882    let account_hash = account.hash_var(w);
883
884    let hash = merkle_path
885        .iter()
886        .enumerate()
887        .fold(account_hash, |accum, (height, path)| {
888            let hashes = match path {
889                MerklePath::Left(right) => [accum, *right],
890                MerklePath::Right(left) => [*left, accum],
891            };
892            w.exists(hashes);
893
894            let param = get_coinbase_param_for_height(height);
895            checked_hash(param, &hashes, w)
896        });
897
898    hash
899}
900
901pub struct PendingCoinbaseWitness {
902    pub pending_coinbase: PendingCoinbase,
903    pub is_new_stack: bool,
904}
905
906impl PendingCoinbaseWitness {
907    fn coinbase_stack_path_exn(&mut self, idx: Address) -> Vec<MerklePath> {
908        self.pending_coinbase.path(idx)
909    }
910
911    fn find_index_of_oldest_stack(&self) -> Address {
912        let stack_id = self
913            .pending_coinbase
914            .oldest_stack_id()
915            .unwrap_or_else(StackId::zero);
916        self.pending_coinbase.find_index(stack_id)
917    }
918
919    fn get_coinbase_stack(&mut self, idx: Address) -> (Stack, Vec<MerklePath>) {
920        let elt = self.pending_coinbase.get_stack(idx.clone()).clone();
921        let path = self.coinbase_stack_path_exn(idx);
922        (elt, path)
923    }
924
925    fn set_oldest_coinbase_stack(&mut self, idx: Address, stack: Stack) {
926        let depth = constraint_constants().pending_coinbase_depth;
927        self.pending_coinbase.set_stack(depth, idx, stack, false);
928    }
929
930    fn find_index_of_newest_stacks(&self) -> (Address, Address) {
931        let depth = constraint_constants().pending_coinbase_depth;
932
933        let index1 = {
934            let stack_id = self.pending_coinbase.latest_stack_id(self.is_new_stack);
935            self.pending_coinbase.find_index(stack_id)
936        };
937
938        let index2 = {
939            let stack_id = self
940                .pending_coinbase
941                .next_stack_id(depth, self.is_new_stack);
942            self.pending_coinbase.find_index(stack_id)
943        };
944
945        (index1, index2)
946    }
947
948    fn get_previous_stack(&self) -> StateStack {
949        if self.is_new_stack {
950            let stack = self.pending_coinbase.current_stack();
951            StateStack {
952                init: stack.state.curr,
953                curr: stack.state.curr,
954            }
955        } else {
956            let stack = self.pending_coinbase.latest_stack(self.is_new_stack);
957            stack.state
958        }
959    }
960
961    fn set_coinbase_stack(&mut self, idx: Address, stack: Stack) {
962        let depth = constraint_constants().pending_coinbase_depth;
963        self.pending_coinbase
964            .set_stack(depth, idx, stack, self.is_new_stack);
965    }
966}
967
968/// Keep it a bit generic, in case we need a merkle tree somewhere else
969pub mod merkle_tree {
970    use crate::{AccountIndex, AddressIterator, Direction, HashesMatrix};
971
972    use super::*;
973
974    pub trait TreeHasher<V> {
975        fn hash_value(value: &V) -> Fp;
976        fn empty_value() -> V;
977        fn merge_hash(depth: usize, left: Fp, right: Fp) -> Fp;
978    }
979
980    #[derive(Clone)]
981    pub struct MiniMerkleTree<K, V, H> {
982        pub values: Vec<V>,
983        pub indexes: HashMap<K, Address>,
984        pub hashes_matrix: HashesMatrix,
985        pub depth: usize,
986        pub _hasher: PhantomData<H>,
987    }
988
989    impl<K, V, H> std::fmt::Debug for MiniMerkleTree<K, V, H>
990    where
991        K: std::fmt::Debug + Ord,
992        V: std::fmt::Debug,
993    {
994        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
995            let mut indexes = self.indexes.iter().collect::<Vec<_>>();
996            indexes.sort_by_key(|(key, _addr)| *key);
997            // indexes.sort_by_key(|(_key, addr)| addr.to_index());
998
999            f.debug_struct("MiniMerkleTree")
1000                .field("values", &self.values)
1001                .field("indexes_len", &indexes.len())
1002                .field("indexes", &indexes)
1003                // .field("hashes_matrix", &self.hashes_matrix)
1004                .field("depth", &self.depth)
1005                .finish()
1006        }
1007    }
1008
1009    impl<K, V, H> MiniMerkleTree<K, V, H>
1010    where
1011        K: Eq + std::hash::Hash,
1012        H: TreeHasher<V>,
1013        V: PartialEq,
1014    {
1015        pub fn create(depth: usize) -> Self {
1016            let max_values = 2u64.pow(depth as u32) as usize;
1017
1018            Self {
1019                values: Vec::with_capacity(max_values),
1020                indexes: HashMap::new(),
1021                depth,
1022                hashes_matrix: HashesMatrix::new(depth),
1023                _hasher: PhantomData,
1024            }
1025        }
1026
1027        pub fn fill_with<I>(&mut self, data: I)
1028        where
1029            I: Iterator<Item = (K, V)>,
1030        {
1031            assert!(self.values.is_empty());
1032
1033            assert_eq!(self.depth, 5);
1034
1035            // OCaml uses those indexes
1036            let indexes: HashMap<usize, usize> = [
1037                (31, 31),
1038                (30, 15),
1039                (29, 23),
1040                (28, 7),
1041                (27, 27),
1042                (26, 11),
1043                (25, 19),
1044                (24, 3),
1045                (23, 29),
1046                (22, 13),
1047                (21, 21),
1048                (20, 5),
1049                (19, 25),
1050                (18, 9),
1051                (17, 17),
1052                (16, 1),
1053                (15, 30),
1054                (14, 14),
1055                (13, 22),
1056                (12, 6),
1057                (11, 26),
1058                (10, 10),
1059                (9, 18),
1060                (8, 2),
1061                (7, 28),
1062                (6, 12),
1063                (5, 20),
1064                (4, 4),
1065                (3, 24),
1066                (2, 8),
1067                (1, 16),
1068                (0, 0),
1069            ]
1070            .iter()
1071            .copied()
1072            .collect();
1073
1074            for (index, (key, value)) in data.enumerate() {
1075                self.values.push(value);
1076
1077                let index = indexes
1078                    .get(&index)
1079                    .copied()
1080                    .map(AccountIndex::from)
1081                    .unwrap();
1082                self.indexes
1083                    .insert(key, Address::from_index(index, self.depth));
1084            }
1085        }
1086
1087        fn get(&self, addr: Address) -> Option<&V> {
1088            assert_eq!(addr.length(), self.depth);
1089
1090            let index = addr.to_index().0 as usize;
1091            self.values.get(index)
1092        }
1093
1094        pub fn get_exn(&self, addr: Address) -> &V {
1095            self.get(addr).unwrap()
1096        }
1097
1098        pub fn set_exn(&mut self, addr: Address, value: V) {
1099            use std::cmp::Ordering::*;
1100
1101            assert_eq!(addr.length(), self.depth);
1102            let index = addr.to_index().0 as usize;
1103
1104            let mut invalidate = true;
1105
1106            match index.cmp(&self.values.len()) {
1107                Less => {
1108                    invalidate = self.values[index] != value;
1109                    self.values[index] = value
1110                }
1111                Equal => self.values.push(value),
1112                Greater => panic!("wrong use of `set_exn`"),
1113            }
1114
1115            if invalidate {
1116                self.hashes_matrix.invalidate_hashes(addr.to_index());
1117            }
1118        }
1119
1120        pub fn find_index_exn(&self, key: K) -> Address {
1121            self.indexes.get(&key).cloned().unwrap()
1122        }
1123
1124        pub fn path_exn(&mut self, addr: Address) -> Vec<MerklePath> {
1125            let mut merkle_path = Vec::with_capacity(addr.length());
1126            let mut path_to_addr = addr.into_iter();
1127            let root = Address::root();
1128
1129            self.emulate_tree_to_get_path(root, &mut path_to_addr, &mut merkle_path);
1130
1131            merkle_path
1132        }
1133
1134        fn get_value_hash(&mut self, addr: Address) -> Fp {
1135            if let Some(hash) = self.hashes_matrix.get(&addr) {
1136                return *hash;
1137            }
1138
1139            let hash = match self.get(addr.clone()) {
1140                Some(value) => H::hash_value(value),
1141                None => H::hash_value(&H::empty_value()),
1142            };
1143
1144            self.hashes_matrix.set(&addr, hash);
1145
1146            hash
1147        }
1148
1149        fn get_node_hash(&mut self, addr: &Address, left: Fp, right: Fp) -> Fp {
1150            if let Some(hash) = self.hashes_matrix.get(addr) {
1151                return *hash;
1152            };
1153
1154            let depth_in_tree = self.depth - addr.length();
1155
1156            let hash = H::merge_hash(depth_in_tree - 1, left, right);
1157            self.hashes_matrix.set(addr, hash);
1158            hash
1159        }
1160
1161        fn emulate_tree_to_get_path(
1162            &mut self,
1163            addr: Address,
1164            path: &mut AddressIterator,
1165            merkle_path: &mut Vec<MerklePath>,
1166        ) -> Fp {
1167            if addr.length() == self.depth {
1168                return self.get_value_hash(addr);
1169            }
1170
1171            let next_direction = path.next();
1172
1173            // We go until the end of the path
1174            if let Some(direction) = next_direction.as_ref() {
1175                let child = match direction {
1176                    Direction::Left => addr.child_left(),
1177                    Direction::Right => addr.child_right(),
1178                };
1179                self.emulate_tree_to_get_path(child, path, merkle_path);
1180            };
1181
1182            let mut get_child_hash = |addr: Address| match self.hashes_matrix.get(&addr) {
1183                Some(hash) => *hash,
1184                None => {
1185                    if let Some(hash) = self.hashes_matrix.get(&addr) {
1186                        *hash
1187                    } else {
1188                        self.emulate_tree_to_get_path(addr, path, merkle_path)
1189                    }
1190                }
1191            };
1192
1193            let left = get_child_hash(addr.child_left());
1194            let right = get_child_hash(addr.child_right());
1195
1196            if let Some(direction) = next_direction {
1197                let hash = match direction {
1198                    Direction::Left => MerklePath::Left(right),
1199                    Direction::Right => MerklePath::Right(left),
1200                };
1201                merkle_path.push(hash);
1202            };
1203
1204            self.get_node_hash(&addr, left, right)
1205        }
1206
1207        pub fn merkle_root(&mut self) -> Fp {
1208            let root = Address::root();
1209
1210            if let Some(hash) = self.hashes_matrix.get(&root) {
1211                return *hash;
1212            };
1213
1214            self.emulate_tree_merkle_root(root)
1215        }
1216
1217        pub fn emulate_tree_merkle_root(&mut self, addr: Address) -> Fp {
1218            let current_depth = self.depth - addr.length();
1219
1220            if current_depth == 0 {
1221                return self.get_value_hash(addr);
1222            }
1223
1224            let mut get_child_hash = |addr: Address| {
1225                if let Some(hash) = self.hashes_matrix.get(&addr) {
1226                    *hash
1227                } else {
1228                    self.emulate_tree_merkle_root(addr)
1229                }
1230            };
1231
1232            let left_hash = get_child_hash(addr.child_left());
1233            let right_hash = get_child_hash(addr.child_right());
1234
1235            self.get_node_hash(&addr, left_hash, right_hash)
1236        }
1237    }
1238
1239    //   [%%define_locally
1240    //   M.
1241    //     ( of_hash
1242    //     , get_exn
1243    //     , path_exn
1244    //     , set_exn
1245    //     , find_index_exn
1246    //     , add_path
1247    //     , merkle_root )]
1248    // end
1249}
1250
1251#[cfg(test)]
1252mod tests {
1253    use crate::FpExt;
1254
1255    use super::*;
1256
1257    #[cfg(target_family = "wasm")]
1258    use wasm_bindgen_test::wasm_bindgen_test as test;
1259
1260    #[test]
1261    fn test_merkle_tree() {
1262        {
1263            const DEPTH: usize = 3;
1264            let mut tree = MiniMerkleTree::<StackId, Stack, StackHasher>::create(DEPTH);
1265            let merkle_root = tree.merkle_root();
1266            assert_eq!(
1267                merkle_root.to_decimal(),
1268                "9939061863620980199451530646711695641079091335264396436068661296746064363179"
1269            );
1270        }
1271
1272        {
1273            const DEPTH: usize = 5;
1274            let mut tree = MiniMerkleTree::<StackId, Stack, StackHasher>::create(DEPTH);
1275            let merkle_root = tree.merkle_root();
1276            assert_eq!(
1277                merkle_root.to_decimal(),
1278                "25504365445533103805898245102289650498571312278321176071043666991586378788150"
1279            );
1280        }
1281    }
1282}
mina_tree/scan_state/pending_coinbase.rs

mina_tree/scan_state/
pending_coinbase.rs
