libzkbob_rs/client/
state.rs

1use std::{convert::TryInto, marker::PhantomData, ops::Range};
2
3use kvdb::KeyValueDB;
4use kvdb_memorydb::InMemory as MemoryDatabase;
5#[cfg(feature = "web")]
6use kvdb_web::Database as WebDatabase;
7use libzeropool::{
8    constants,
9    fawkes_crypto::{ff_uint::Num, ff_uint::PrimeField, BorshDeserialize, BorshSerialize},
10    native::{
11        account::Account, account::Account as NativeAccount, note::Note, note::Note as NativeNote,
12        params::PoolParams,
13    },
14};
15
16use crate::{merkle::MerkleTree, sparse_array::SparseArray};
17
18pub type TxStorage<D, Fr> = SparseArray<D, Transaction<Fr>>;
19
20#[derive(BorshSerialize, BorshDeserialize, PartialEq, Eq, Debug)]
21pub enum Transaction<Fr: PrimeField> {
22    Account(NativeAccount<Fr>),
23    Note(NativeNote<Fr>),
24}
25
26pub struct State<D: KeyValueDB, P: PoolParams> {
27    pub tree: MerkleTree<D, P>,
28    /// Stores only usable (own) accounts and notes
29    pub(crate) txs: TxStorage<D, P::Fr>,
30    pub(crate) latest_account: Option<NativeAccount<P::Fr>>,
31    pub latest_account_index: Option<u64>,
32    /// Latest owned note index
33    pub latest_note_index: u64,
34    _params: PhantomData<P>,
35}
36
37#[cfg(feature = "web")]
38impl<P> State<WebDatabase, P>
39where
40    P: PoolParams,
41    P::Fr: 'static,
42{
43    pub async fn init_web(db_id: String, params: P) -> Self {
44        let merkle_db_name = format!("zkb.{}.smt", &db_id);
45        let tx_db_name = format!("zkb.{}.txs", &db_id);
46        let tree = MerkleTree::new_web(&merkle_db_name, params.clone()).await;
47        let txs = TxStorage::new_web(&tx_db_name).await;
48
49        Self::new(tree, txs)
50    }
51}
52
53impl<P> State<MemoryDatabase, P>
54where
55    P: PoolParams,
56    P::Fr: 'static,
57{
58    pub fn init_test(params: P) -> Self {
59        let tree = MerkleTree::new_test(params);
60        let txs = TxStorage::new_test();
61
62        Self::new(tree, txs)
63    }
64}
65
66impl<D, P> State<D, P>
67where
68    D: KeyValueDB,
69    P: PoolParams,
70    P::Fr: 'static,
71{
72    pub fn new(tree: MerkleTree<D, P>, txs: TxStorage<D, P::Fr>) -> Self {
73        // TODO: Cache
74        let mut latest_account_index = None;
75        let mut latest_note_index = 0;
76        let mut latest_account = None;
77        for (index, tx) in txs.iter() {
78            match tx {
79                Transaction::Account(acc) => {
80                    if index >= latest_account_index.unwrap_or(0) {
81                        latest_account_index = Some(index);
82                        latest_account = Some(acc);
83                    }
84                }
85                Transaction::Note(_) => {
86                    if index >= latest_note_index {
87                        latest_note_index = index;
88                    }
89                }
90            }
91        }
92
93        State {
94            tree,
95            txs,
96            latest_account_index,
97            latest_note_index,
98            latest_account,
99            _params: Default::default(),
100        }
101    }
102
103    /// Add OUT + 1 hashes to the tree
104    pub fn add_hashes(&mut self, at_index: u64, hashes: &[Num<P::Fr>]) {
105        // FIXME: return an error instead of asserts
106        assert_eq!(
107            at_index % (constants::OUT as u64 + 1),
108            0,
109            "index must be divisible by {}",
110            constants::OUT + 1
111        );
112
113        self.tree.add_hashes(at_index, hashes.iter().copied());
114    }
115
116    /// Add hashes, account, and notes to state
117    pub fn add_full_tx(
118        &mut self,
119        at_index: u64,
120        hashes: &[Num<P::Fr>],
121        account: Option<Account<P::Fr>>,
122        notes: &[(u64, Note<P::Fr>)],
123    ) {
124        self.add_hashes(at_index, hashes);
125
126        if let Some(acc) = account {
127            self.add_account(at_index, acc);
128        }
129
130        // Store notes
131        for (index, note) in notes {
132            self.add_note(*index, *note);
133        }
134    }
135
136    /// Cache account at specified index.
137    pub fn add_account(&mut self, at_index: u64, account: Account<P::Fr>) {
138        // Update tx storage
139        self.txs.set(at_index, &Transaction::Account(account));
140
141        if at_index >= self.latest_account_index.unwrap_or(0) {
142            self.latest_account_index = Some(at_index);
143            self.latest_account = Some(account);
144        }
145    }
146
147    /// Caches a note at specified index.
148    pub fn add_note(&mut self, at_index: u64, note: Note<P::Fr>) {
149        if self.txs.get(at_index).is_some() {
150            return;
151        }
152
153        self.txs.set(at_index, &Transaction::Note(note));
154
155        if at_index > self.latest_note_index {
156            self.latest_note_index = at_index;
157        }
158    }
159
160    pub fn get_all_txs(&self) -> Vec<(u64, Transaction<P::Fr>)> {
161        self.txs.iter().collect()
162    }
163
164    pub fn get_usable_notes(&self) -> Vec<(u64, Note<P::Fr>)> {
165        let next_usable_index = self.earliest_usable_index();
166
167        // Fetch all usable notes from the state
168        self
169        .txs
170        .iter_slice(next_usable_index..=self.latest_note_index)
171        .filter_map(|(index, tx)| match tx {
172            Transaction::Note(note) => Some((index, note)),
173            _ => None,
174        })
175        .collect()
176    }
177
178    pub fn get_account(&self, index: u64) -> Option<Account<P::Fr>> {
179        match self.txs.get(index) {
180            Some(Transaction::Account(acc)) => Some(acc),
181            _ => None,
182        }
183    }
184
185    pub fn get_previous_account(&self, index: u64) -> Option<(u64, Account<P::Fr>)> {
186        if index == 0 { return None }
187
188        let prev_acc_indexes = self.txs
189            .iter_slice(0..=(index-1))
190            .filter_map(|(idx, tx)| match tx {
191                Transaction::Account(_) => Some(idx),
192                _ => None,
193            })
194            .max();
195
196        match prev_acc_indexes {
197            Some(idx) => Some((idx, self.get_account(idx).unwrap())),
198            _ => None,
199        }
200    }
201
202    pub fn get_notes_in_range(&self, range: Range<u64>) -> Vec<(u64, Note<P::Fr>)> {
203        self.txs
204            .iter_slice(range.start..=range.end.saturating_sub(1))
205            .filter_map(|(idx, tx)| match tx {
206                Transaction::Note(note) => Some((idx, note)),
207                _ => None,
208            })
209            .collect()
210    }
211
212    /// Return an index of a earliest usable note.
213    pub fn earliest_usable_index(&self) -> u64 {
214        let latest_account_index = self
215            .latest_account
216            .map(|acc| acc.i.to_num())
217            .unwrap_or(Num::ZERO)
218            .try_into()
219            .unwrap();
220
221        self.txs
222            .iter_slice(latest_account_index..=self.latest_note_index)
223            .filter_map(|(index, tx)| match tx {
224                Transaction::Note(_) => Some(index),
225                _ => None,
226            })
227            .next()
228            .unwrap_or(latest_account_index)
229    }
230
231    /// Return an index of a earliest usable note including optimistic state
232    pub fn earliest_usable_index_optimistic(
233        &self,
234        optimistic_accounts: &[(u64, Account<P::Fr>)],
235        optimistic_notes: &[(u64, Note<P::Fr>)]
236    ) -> u64 {
237        let latest_account_index = optimistic_accounts
238            .last()
239            .map(|indexed_acc| indexed_acc.1)
240            .or(self.latest_account)
241            .map(|acc| acc.i.to_num())
242            .unwrap_or(Num::ZERO)
243            .try_into()
244            .unwrap();
245
246        
247        let latest_note_index_optimistic = optimistic_notes
248            .last()
249            .map(|indexed_note| indexed_note.0)
250            .unwrap_or(self.latest_note_index);
251
252        let optimistic_note_indices = optimistic_notes
253            .iter()
254            .map(|indexed_note| indexed_note.0)
255            .filter(move |index| (latest_account_index..=latest_note_index_optimistic).contains(index));
256
257        
258        self.txs
259            .iter_slice(latest_account_index..=latest_note_index_optimistic)
260            .filter_map(|(index, tx)| match tx {
261                Transaction::Note(_) => Some(index),
262                _ => None,
263            })
264            .chain(optimistic_note_indices)
265            .next()
266            .unwrap_or(latest_account_index)
267    }
268
269    /// Returns user's total balance (account + available notes).
270    pub fn total_balance(&self) -> Num<P::Fr> {
271        self.account_balance() + self.note_balance()
272    }
273
274    pub fn account_balance(&self) -> Num<P::Fr> {
275        self.latest_account
276            .map(|acc| acc.b.to_num())
277            .unwrap_or(Num::ZERO)
278    }
279
280    pub fn note_balance(&self) -> Num<P::Fr> {
281        let starting_index = self
282            .latest_account
283            .map(|acc| acc.i.to_num().try_into().unwrap())
284            .unwrap_or(0);
285        let mut note_balance = Num::ZERO;
286        for (_, tx) in self.txs.iter_slice(starting_index..=self.latest_note_index) {
287            if let Transaction::Note(note) = tx {
288                note_balance += note.b.to_num();
289            }
290        }
291
292        note_balance
293    }
294
295    // rollback current state, return updated next_index
296    pub fn rollback(&mut self, rollback_index: u64) -> u64 {
297        if rollback_index > self.tree.next_index() {
298            return self.tree.next_index();
299        }
300
301        let rollback_index = (rollback_index >> constants::OUTPLUSONELOG) << constants::OUTPLUSONELOG;
302        self.txs.remove_from(rollback_index);
303        self.latest_account = None;
304        self.latest_account_index = None;
305        self.latest_note_index = 0;
306        for (index, tx) in self.txs.iter() {
307            match tx {
308                Transaction::Account(acc) => {
309                    if index >= self.latest_account_index.unwrap_or(0) {
310                        self.latest_account_index = Some(index);
311                        self.latest_account = Some(acc);
312                    }
313                }
314                Transaction::Note(_) => {
315                    if index >= self.latest_note_index {
316                        self.latest_note_index = index;
317                    }
318                }
319            }
320        }
321
322        if self.tree.rollback(rollback_index).is_none() {
323            self.wipe();
324        }
325
326        self.tree.next_index()
327    }
328
329    pub fn wipe(&mut self) {
330        self.tree.wipe();
331
332        self.txs.remove_all();
333        self.latest_account_index = None;
334        self.latest_account = None;
335        self.latest_note_index = 0;
336    }
337
338}
libzkbob_rs/client/state.rs

libzkbob_rs/client/
state.rs
