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

mina_tree/scan_state/
pending_coinbase.rs
