openmina_bootstrap_sandbox/
snarked_ledger.rs

1use binprot::{BinProtRead, BinProtWrite};
2use std::{future::Future, io, pin::Pin};
3use thiserror::Error;
4
5use ledger::{Account, AccountIndex, Address, BaseLedger, Database, Mask};
6use mina_p2p_messages::{list::List, rpc::AnswerSyncLedgerQueryV2, v2};
7
8use super::client::Client;
9
10pub struct SnarkedLedger {
11    pub inner: Mask,
12    // NOTE: it is not the same as the merkle tree root
13    pub top_hash: Option<v2::LedgerHash>,
14    pub num: u32,
15}
16
17#[derive(Debug, Error)]
18pub enum Error {
19    #[error("{0}")]
20    Serde(#[from] serde_json::Error),
21    #[error("{0}")]
22    Io(#[from] io::Error),
23}
24
25impl SnarkedLedger {
26    pub fn empty() -> Self {
27        SnarkedLedger {
28            inner: Mask::new_root(Database::create(35)),
29            top_hash: None,
30            num: 0,
31        }
32    }
33
34    // for debugging
35    pub fn store_bin<W>(&self, mut writer: W) -> io::Result<()>
36    where
37        W: io::Write,
38    {
39        let accounts = self.inner.fold(vec![], |mut accounts, account| {
40            accounts.push(account.clone());
41            accounts
42        });
43        self.top_hash.binprot_write(&mut writer)?;
44        accounts.binprot_write(&mut writer)
45    }
46
47    pub fn load_bin<R>(mut reader: R) -> Result<Self, binprot::Error>
48    where
49        R: io::Read,
50    {
51        let top_hash = Option::binprot_read(&mut reader)?;
52        let accounts = List::<Account>::binprot_read(&mut reader)?;
53
54        let num = accounts.len() as _;
55        let mut inner = Mask::new_root(Database::create(35));
56        for account in accounts {
57            let account_id = account.id();
58            inner.get_or_create_account(account_id, account).unwrap();
59        }
60
61        let _ = inner.merkle_root();
62
63        Ok(SnarkedLedger {
64            inner,
65            top_hash,
66            num,
67        })
68    }
69
70    pub async fn sync_new(&mut self, client: &mut Client, root: &v2::LedgerHash) {
71        let q = v2::MinaLedgerSyncLedgerQueryStableV1::NumAccounts;
72        let r = match client
73            .rpc::<AnswerSyncLedgerQueryV2>((root.0.clone(), q))
74            .await
75            .unwrap()
76            .0
77        {
78            Ok(v) => v,
79            Err(e) => panic!("answer_sync_ledger returned error: {e}"),
80        };
81        let (num, hash) = match r {
82            v2::MinaLedgerSyncLedgerAnswerStableV2::NumAccounts(num, hash) => (num.0, hash),
83            _ => panic!(),
84        };
85        self.top_hash = Some(hash.clone());
86        self.num = num as _;
87
88        if self.inner.num_accounts() > num as _ {
89            self.inner = Mask::new_root(Database::create(35));
90        }
91
92        self.sync_at_depth_new(client, root.clone(), hash.clone(), 0, 0)
93            .await;
94        let actual_hash = self.inner.merkle_root();
95        let actual_hash = v2::LedgerHash::from(v2::MinaBaseLedgerHash0StableV1(actual_hash.into()));
96        assert_eq!(actual_hash, root.clone());
97    }
98
99    fn sync_at_depth_boxed_new<'a, 'b: 'a>(
100        &'b mut self,
101        client: &'a mut Client,
102        root: v2::LedgerHash,
103        hash: v2::LedgerHash,
104        depth: i32,
105        pos: u32,
106    ) -> Pin<Box<dyn Future<Output = ()> + 'a>> {
107        Box::pin(self.sync_at_depth_new(client, root, hash, depth, pos))
108    }
109
110    async fn sync_at_depth_new(
111        &mut self,
112        client: &mut Client,
113        root: v2::LedgerHash,
114        hash: v2::LedgerHash,
115        depth: i32,
116        pos: u32,
117    ) {
118        let addr = Address::from_index(AccountIndex(pos as _), depth as _);
119        let actual_hash = self.inner.get_inner_hash_at_addr(addr.clone()).unwrap();
120        if depth == 0 && root.0 == actual_hash.into() || depth > 0 && hash.0 == actual_hash.into() {
121            return;
122        }
123
124        const Q: i32 = 3;
125        if depth == 32 - Q {
126            let p = (pos * (1 << Q)).to_be_bytes().to_vec();
127            let q = v2::MinaLedgerSyncLedgerQueryStableV1::WhatContents(
128                v2::MerkleAddressBinableArgStableV1((depth as u64).into(), p.into()),
129            );
130            log::info!("{}", serde_json::to_string(&q).unwrap());
131            let r = client
132                .rpc::<AnswerSyncLedgerQueryV2>((root.0.clone(), q))
133                .await
134                .unwrap()
135                .0;
136            match r {
137                Err(err) => {
138                    log::error!("num: {}, error: {err}", self.num);
139                }
140                Ok(v2::MinaLedgerSyncLedgerAnswerStableV2::ContentsAre(accounts)) => {
141                    for (o, account) in accounts.into_iter().enumerate() {
142                        let account = Account::try_from(&account).unwrap();
143                        self.inner
144                            .set_at_index(
145                                AccountIndex((pos * 8 * (1 << Q)) as u64 + o as u64),
146                                Box::new(account),
147                            )
148                            .unwrap();
149                    }
150                }
151                _ => panic!(),
152            }
153        } else {
154            let b = ((depth as usize + 7) / 8).min(4);
155            let p = if depth > 0 {
156                pos * (1 << (32 - depth))
157            } else {
158                0
159            };
160            let p = p.to_be_bytes()[..b].to_vec();
161            let q = v2::MinaLedgerSyncLedgerQueryStableV1::WhatChildHashes(
162                v2::MerkleAddressBinableArgStableV1((depth as u64).into(), p.into()),
163            );
164            log::info!("{}", serde_json::to_string(&q).unwrap());
165            let r = client
166                .rpc::<AnswerSyncLedgerQueryV2>((root.0.clone(), q))
167                .await
168                .unwrap()
169                .0
170                .unwrap();
171            match r {
172                v2::MinaLedgerSyncLedgerAnswerStableV2::ChildHashesAre(l, r) => {
173                    self.sync_at_depth_boxed_new(client, root.clone(), l, depth + 1, pos * 2)
174                        .await;
175                    self.sync_at_depth_boxed_new(client, root.clone(), r, depth + 1, pos * 2 + 1)
176                        .await;
177                }
178                _ => panic!(),
179            };
180        }
181
182        let addr = Address::from_index(AccountIndex(pos as _), depth as _);
183        let actual_hash = self.inner.get_inner_hash_at_addr(addr).unwrap();
184        let actual_hash = v2::LedgerHash::from(v2::MinaBaseLedgerHash0StableV1(actual_hash.into()));
185        if depth == 0 {
186            assert_eq!(root, actual_hash);
187        } else {
188            assert_eq!(hash, actual_hash);
189        }
190    }
191
192    pub fn serve_query(
193        &mut self,
194        q: v2::MinaLedgerSyncLedgerQueryStableV1,
195    ) -> v2::MinaLedgerSyncLedgerAnswerStableV2 {
196        log::info!("query: {q:?}");
197        match q {
198            v2::MinaLedgerSyncLedgerQueryStableV1::NumAccounts => {
199                v2::MinaLedgerSyncLedgerAnswerStableV2::NumAccounts(
200                    (self.num as u64).into(),
201                    self.top_hash.as_ref().unwrap().clone(),
202                )
203            }
204            v2::MinaLedgerSyncLedgerQueryStableV1::WhatChildHashes(address) => {
205                let addr = Address::from(address);
206
207                let hash = self
208                    .inner
209                    .get_inner_hash_at_addr(addr.child_left())
210                    .unwrap();
211                let left = v2::LedgerHash::from(v2::MinaBaseLedgerHash0StableV1(hash.into()));
212
213                let hash = self
214                    .inner
215                    .get_inner_hash_at_addr(addr.child_right())
216                    .unwrap();
217                let right = v2::LedgerHash::from(v2::MinaBaseLedgerHash0StableV1(hash.into()));
218
219                v2::MinaLedgerSyncLedgerAnswerStableV2::ChildHashesAre(left, right)
220            }
221            v2::MinaLedgerSyncLedgerQueryStableV1::WhatContents(address) => {
222                let addr = Address::from(address);
223
224                let depth = addr.length();
225                let pos = addr.to_index().0;
226
227                let mut accounts = List::new();
228                let mut offset = 0;
229                let batch_length = 1u64 << (35 - depth);
230                loop {
231                    if offset == batch_length {
232                        break;
233                    }
234
235                    let pos = pos * batch_length + offset;
236                    offset += 1;
237                    if pos == self.num as u64 {
238                        break;
239                    }
240                    let addr = Address::from_index(AccountIndex(pos as _), 35);
241                    let account = self.inner.get(addr);
242                    if let Some(account) = account {
243                        accounts.push_back((&*account).into());
244                    } else {
245                        break;
246                    }
247                }
248                v2::MinaLedgerSyncLedgerAnswerStableV2::ContentsAre(accounts)
249            }
250        }
251    }
252}