mina_tree/
tree.rs

1use crate::{
2    address::Address,
3    base::AccountIndex,
4    tree_version::{TreeVersion, V2},
5};
6use mina_curves::pasta::Fp;
7use once_cell::sync::Lazy;
8use std::{collections::BTreeMap, fmt::Debug, sync::Mutex};
9
10#[derive(Clone, Debug)]
11struct Leaf<T: TreeVersion> {
12    account: Option<Box<T::Account>>,
13}
14
15#[derive(PartialEq)]
16pub struct HashesMatrix {
17    /// 2 dimensions matrix
18    matrix: BTreeMap<u64, Fp>,
19    empty_hashes: Vec<Option<Fp>>,
20    ledger_depth: usize,
21    nhashes: usize,
22}
23
24impl Clone for HashesMatrix {
25    fn clone(&self) -> Self {
26        Self {
27            matrix: self.matrix.clone(),
28            empty_hashes: self.empty_hashes.clone(),
29            ledger_depth: self.ledger_depth,
30            nhashes: self.nhashes,
31        }
32    }
33}
34
35impl Debug for HashesMatrix {
36    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
37        f.debug_struct("HashesMatrix")
38            .field("matrix_len", &self.matrix.len())
39            .field("nhashes", &self.nhashes)
40            .finish()
41    }
42}
43
44impl HashesMatrix {
45    pub fn new(ledger_depth: usize) -> Self {
46        Self {
47            matrix: BTreeMap::new(),
48            ledger_depth,
49            empty_hashes: vec![None; ledger_depth],
50            nhashes: 0,
51        }
52    }
53
54    pub fn get(&self, addr: &Address) -> Option<&Fp> {
55        let linear: u64 = addr.to_linear_index();
56        self.matrix.get(&linear)
57    }
58
59    pub fn set(&mut self, addr: &Address, hash: Fp) {
60        let linear: u64 = addr.to_linear_index();
61        let old = self.matrix.insert(linear, hash);
62        assert!(old.is_none());
63        self.nhashes += 1;
64    }
65
66    /// Do not use directly. Used to forcefully reconstructing the hashes
67    /// matrix from raw data.
68    pub fn set_raw_index(&mut self, idx: u64, hash: Fp) {
69        let old = self.matrix.insert(idx, hash);
70        assert!(old.is_none());
71        self.nhashes += 1;
72    }
73
74    pub fn remove(&mut self, addr: &Address) {
75        let linear: u64 = addr.to_linear_index();
76        self.remove_at_index(linear);
77    }
78
79    fn remove_at_index(&mut self, linear: u64) {
80        let old = self.matrix.remove(&linear);
81        if old.is_some() {
82            self.nhashes -= 1;
83        }
84    }
85
86    pub(super) fn transfert_hashes(&mut self, hashes: HashesMatrix) {
87        for (index, hash) in hashes.matrix {
88            let old = self.matrix.insert(index, hash);
89            if old.is_none() {
90                self.nhashes += 1;
91            }
92        }
93    }
94
95    pub fn invalidate_hashes(&mut self, account_index: AccountIndex) {
96        let mut addr = Address::from_index(account_index, self.ledger_depth);
97
98        loop {
99            let index = addr.to_linear_index();
100            self.remove_at_index(index);
101            addr = match addr.parent() {
102                Some(addr) => addr,
103                None => break,
104            }
105        }
106    }
107
108    pub fn empty_hash_at_height(&mut self, height: usize) -> Fp {
109        if let Some(Some(hash)) = self.empty_hashes.get(height) {
110            return *hash;
111        };
112
113        // If `depth` is out of bound, see `HASH_EMPTIES`
114        let hash = HASH_EMPTIES.lock().unwrap()[height];
115        self.empty_hashes[height] = Some(hash);
116
117        hash
118    }
119
120    pub fn clear(&mut self) {
121        let ledger_depth = self.ledger_depth;
122        *self = Self {
123            matrix: BTreeMap::new(),
124            ledger_depth,
125            empty_hashes: vec![None; ledger_depth],
126            nhashes: 0,
127        }
128    }
129
130    pub fn take(&mut self) -> Self {
131        let Self {
132            matrix,
133            empty_hashes,
134            ledger_depth,
135            nhashes,
136        } = self;
137
138        Self {
139            matrix: std::mem::take(matrix),
140            empty_hashes: std::mem::take(empty_hashes),
141            ledger_depth: *ledger_depth,
142            nhashes: *nhashes,
143        }
144    }
145
146    pub fn get_raw_inner_hashes(&self) -> Vec<(u64, Fp)> {
147        self.matrix.clone().into_iter().collect()
148    }
149
150    pub fn set_raw_inner_hashes(&mut self, hashes: Vec<(u64, Fp)>) {
151        for (idx, hash) in hashes {
152            self.set_raw_index(idx, hash);
153        }
154    }
155}
156
157static HASH_EMPTIES: Lazy<Mutex<Vec<Fp>>> = Lazy::new(|| {
158    /// This value needs to be changed when the tree's height change
159    const RANGE_HEIGHT: std::ops::Range<usize> = 0..36;
160
161    Mutex::new((RANGE_HEIGHT).map(V2::empty_hash_at_height).collect())
162});