1use alloc::boxed::Box;
2use alloc::collections::{BTreeMap, BTreeSet};
3use alloc::sync::Arc;
4use alloc::vec::Vec;
5
6use miden_protocol::Word;
7use miden_protocol::account::delta::AccountUpdateDetails;
8use miden_protocol::account::{AccountId, StorageSlot, StorageSlotContent, StorageSlotType};
9use miden_protocol::address::NetworkId;
10use miden_protocol::batch::{ProposedBatch, ProvenBatch};
11use miden_protocol::block::{BlockHeader, BlockNumber, ProvenBlock};
12use miden_protocol::crypto::merkle::mmr::{Forest, Mmr, MmrProof};
13use miden_protocol::note::{NoteAttachments, NoteHeader, NoteId, NoteScript, NoteTag};
14use miden_protocol::transaction::{ProvenTransaction, TransactionInputs};
15use miden_testing::{MockChain, MockChainNote};
16use miden_tx::utils::sync::RwLock;
17
18use crate::Client;
19use crate::rpc::domain::account::{
20 AccountDetails,
21 AccountProof,
22 AccountStorageDetails,
23 AccountStorageMapDetails,
24 AccountVaultDetails,
25 GetAccountRequest,
26 StorageMapEntries,
27 StorageMapEntry,
28 StorageMapFetch,
29};
30use crate::rpc::domain::account_vault::{AccountVaultInfo, AccountVaultUpdate};
31use crate::rpc::domain::note::{CommittedNote, FetchedNote, NoteSyncBlock};
32use crate::rpc::domain::nullifier::NullifierUpdate;
33use crate::rpc::domain::status::NetworkNoteStatusInfo;
34use crate::rpc::domain::storage_map::{StorageMapInfo, StorageMapUpdate};
35use crate::rpc::domain::sync::{ChainMmrInfo, SyncTarget};
36use crate::rpc::domain::transaction::TransactionRecord;
37use crate::rpc::{AccountStateAt, NodeRpcClient, RpcError, RpcStatusInfo};
38
39pub type MockClient<AUTH> = Client<AUTH>;
40
41#[derive(Clone)]
52pub struct MockRpcApi {
53 account_commitment_updates: Arc<RwLock<BTreeMap<BlockNumber, BTreeMap<AccountId, Word>>>>,
54 pub mock_chain: Arc<RwLock<MockChain>>,
55 historical_chains: Arc<RwLock<BTreeMap<BlockNumber, Arc<MockChain>>>>,
57 oversize_threshold: usize,
58 erased_notes: Arc<RwLock<Vec<NoteHeader>>>,
60 private_note_attachments: Arc<RwLock<BTreeMap<NoteId, NoteAttachments>>>,
64}
65
66impl Default for MockRpcApi {
67 fn default() -> Self {
68 Self::new(MockChain::new())
69 }
70}
71
72impl MockRpcApi {
73 const PAGINATION_BLOCK_LIMIT: u32 = 5;
75
76 pub fn new(mock_chain: MockChain) -> Self {
78 Self {
79 account_commitment_updates: Arc::new(RwLock::new(build_account_updates(&mock_chain))),
80 mock_chain: Arc::new(RwLock::new(mock_chain)),
81 historical_chains: Arc::new(RwLock::new(BTreeMap::new())),
82 oversize_threshold: 1000,
83 erased_notes: Arc::new(RwLock::new(Vec::new())),
84 private_note_attachments: Arc::new(RwLock::new(BTreeMap::new())),
85 }
86 }
87
88 pub fn register_private_note_attachments(&self, note_id: NoteId, attachments: NoteAttachments) {
91 self.private_note_attachments.write().insert(note_id, attachments);
92 }
93
94 #[must_use]
98 pub fn with_oversize_threshold(mut self, threshold: usize) -> Self {
99 self.oversize_threshold = threshold;
100 self
101 }
102
103 pub fn mark_note_as_erased(&self, header: NoteHeader) {
105 self.erased_notes.write().push(header);
106 }
107
108 pub fn get_mmr(&self) -> Mmr {
110 self.mock_chain.read().blockchain().as_mmr().clone()
111 }
112
113 pub fn get_chain_tip_block_num(&self) -> BlockNumber {
115 self.mock_chain.read().latest_block_header().block_num()
116 }
117
118 pub fn prove_block(&self) {
121 let proven_block = {
122 let mut mock_chain = self.mock_chain.write();
123 let historical_block_num = mock_chain.latest_block_header().block_num();
124 let snapshot = Arc::new(mock_chain.clone());
125 let proven_block = mock_chain.prove_next_block().unwrap();
126 self.historical_chains.write().insert(historical_block_num, snapshot);
127 proven_block
128 };
129 let block_num = proven_block.header().block_num();
130 let mut account_commitment_updates = self.account_commitment_updates.write();
131 let updates: BTreeMap<AccountId, Word> = proven_block
132 .body()
133 .updated_accounts()
134 .iter()
135 .map(|update| (update.account_id(), update.final_state_commitment()))
136 .collect();
137
138 if !updates.is_empty() {
139 account_commitment_updates.insert(block_num, updates);
140 }
141 }
142
143 fn get_block_by_num(&self, block_num: BlockNumber) -> BlockHeader {
145 self.mock_chain.read().block_header(block_num.as_usize())
146 }
147
148 fn get_sync_account_vault_request(
152 &self,
153 block_from: BlockNumber,
154 block_to: BlockNumber,
155 account_id: AccountId,
156 ) -> AccountVaultInfo {
157 let chain_tip = self.get_chain_tip_block_num();
158 let target_block = block_to.min(chain_tip);
159
160 let page_end_block: BlockNumber = (block_from.as_u32() + Self::PAGINATION_BLOCK_LIMIT)
161 .min(target_block.as_u32())
162 .into();
163
164 let mut updates = vec![];
165 for block in self.mock_chain.read().proven_blocks() {
166 let block_number = block.header().block_num();
167 if block_number < block_from || block_number > page_end_block {
169 continue;
170 }
171
172 for update in block
173 .body()
174 .updated_accounts()
175 .iter()
176 .filter(|block_acc_update| block_acc_update.account_id() == account_id)
177 {
178 let AccountUpdateDetails::Delta(account_delta) = update.details().clone() else {
179 continue;
180 };
181
182 let vault_delta = account_delta.vault();
183
184 for asset in vault_delta.added_assets() {
185 let account_vault_update = AccountVaultUpdate {
186 block_num: block_number,
187 asset: Some(asset),
188 vault_key: asset.vault_key(),
189 };
190 updates.push(account_vault_update);
191 }
192 }
193 }
194
195 AccountVaultInfo {
196 chain_tip,
197 block_number: page_end_block,
198 updates,
199 }
200 }
201
202 fn get_sync_transactions_request(
204 &self,
205 block_from: BlockNumber,
206 block_to: BlockNumber,
207 account_ids: &[AccountId],
208 ) -> Vec<TransactionRecord> {
209 let mut transactions = Vec::new();
210 for block in self.mock_chain.read().proven_blocks() {
211 let block_number = block.header().block_num();
212 if block_number < block_from || block_number > block_to {
213 continue;
214 }
215
216 for transaction_header in block.body().transactions().as_slice() {
217 if !account_ids.contains(&transaction_header.account_id()) {
218 continue;
219 }
220
221 let erased_output_notes = self.erased_notes.read().clone();
222
223 transactions.push(TransactionRecord {
224 block_num: block_number,
225 transaction_header: transaction_header.clone(),
226 output_notes: vec![],
227 erased_output_notes,
228 });
229 }
230 }
231
232 transactions
233 }
234
235 fn get_sync_storage_maps_request(
239 &self,
240 block_from: BlockNumber,
241 block_to: BlockNumber,
242 account_id: AccountId,
243 ) -> StorageMapInfo {
244 let chain_tip = self.get_chain_tip_block_num();
245 let target_block = block_to.min(chain_tip);
246
247 let page_end_block: BlockNumber = (block_from.as_u32() + Self::PAGINATION_BLOCK_LIMIT)
248 .min(target_block.as_u32())
249 .into();
250
251 let mut updates = vec![];
252 for block in self.mock_chain.read().proven_blocks() {
253 let block_number = block.header().block_num();
254 if block_number < block_from || block_number > page_end_block {
256 continue;
257 }
258
259 for update in block
260 .body()
261 .updated_accounts()
262 .iter()
263 .filter(|block_acc_update| block_acc_update.account_id() == account_id)
264 {
265 let AccountUpdateDetails::Delta(account_delta) = update.details().clone() else {
266 continue;
267 };
268
269 let storage_delta = account_delta.storage();
270
271 for (slot_name, map_delta) in storage_delta.maps() {
272 for (key, value) in map_delta.entries() {
273 let storage_map_info = StorageMapUpdate {
274 block_num: block_number,
275 slot_name: slot_name.clone(),
276 key: *key,
277 value: *value,
278 };
279 updates.push(storage_map_info);
280 }
281 }
282 }
283 }
284
285 StorageMapInfo {
286 chain_tip,
287 block_number: page_end_block,
288 updates,
289 }
290 }
291
292 pub fn get_available_notes(&self) -> Vec<MockChainNote> {
293 self.mock_chain.read().committed_notes().values().cloned().collect()
294 }
295
296 pub fn get_public_available_notes(&self) -> Vec<MockChainNote> {
297 self.mock_chain
298 .read()
299 .committed_notes()
300 .values()
301 .filter(|n| matches!(n, MockChainNote::Public(_, _)))
302 .cloned()
303 .collect()
304 }
305
306 pub fn get_private_available_notes(&self) -> Vec<MockChainNote> {
307 self.mock_chain
308 .read()
309 .committed_notes()
310 .values()
311 .filter(|n| matches!(n, MockChainNote::Private(_, _, _)))
312 .cloned()
313 .collect()
314 }
315
316 pub fn advance_blocks(&self, num_blocks: u32) {
317 let mut mock_chain = self.mock_chain.write();
318 let block_num = mock_chain.latest_block_header().block_num();
319 let snapshot = Arc::new(mock_chain.clone());
320 mock_chain.prove_until_block(block_num + num_blocks).unwrap();
321 self.historical_chains.write().insert(block_num, snapshot);
322 }
323}
324#[cfg_attr(not(target_arch = "wasm32"), async_trait::async_trait)]
325#[cfg_attr(target_arch = "wasm32", async_trait::async_trait(?Send))]
326impl NodeRpcClient for MockRpcApi {
327 fn has_genesis_commitment(&self) -> Option<Word> {
328 None
329 }
330
331 async fn set_genesis_commitment(&self, _commitment: Word) -> Result<(), RpcError> {
332 Ok(())
334 }
335
336 async fn sync_notes(
339 &self,
340 block_from: BlockNumber,
341 block_to: BlockNumber,
342 note_tags: &BTreeSet<NoteTag>,
343 ) -> Result<Vec<NoteSyncBlock>, RpcError> {
344 let mut blocks_with_notes: BTreeMap<BlockNumber, BTreeMap<NoteId, CommittedNote>> =
345 BTreeMap::new();
346 for note in self.mock_chain.read().committed_notes().values() {
347 let note_block = note.inclusion_proof().location().block_num();
348 if note_tags.contains(¬e.metadata().tag())
349 && note_block >= block_from
350 && note_block <= block_to
351 {
352 let committed =
353 CommittedNote::new(note.id(), *note.metadata(), note.inclusion_proof().clone());
354 blocks_with_notes.entry(note_block).or_default().insert(note.id(), committed);
355 }
356 }
357
358 Ok(blocks_with_notes
359 .into_iter()
360 .map(|(bn, notes)| {
361 let block_header = self.get_block_by_num(bn);
362 let mmr_path = self.get_mmr().open(bn.as_usize()).unwrap().merkle_path().clone();
363 NoteSyncBlock { block_header, mmr_path, notes }
364 })
365 .collect())
366 }
367
368 async fn sync_chain_mmr(
369 &self,
370 current_block_height: BlockNumber,
371 upper_bound: SyncTarget,
372 ) -> Result<ChainMmrInfo, RpcError> {
373 let chain_tip = self.get_chain_tip_block_num();
374 let target_block = match upper_bound {
376 SyncTarget::CommittedChainTip | SyncTarget::ProvenChainTip => chain_tip,
377 };
378
379 let from_forest = if current_block_height == target_block {
380 target_block.as_usize()
381 } else {
382 current_block_height.as_u32() as usize + 1
383 };
384
385 let mmr_delta = self
386 .get_mmr()
387 .get_delta(
388 Forest::new(from_forest).unwrap(),
389 Forest::new(target_block.as_usize()).unwrap(),
390 )
391 .unwrap();
392
393 let block_header = self.get_block_by_num(target_block);
394
395 Ok(ChainMmrInfo {
396 block_from: current_block_height,
397 block_to: target_block,
398 mmr_delta,
399 block_header,
400 })
401 }
402
403 async fn get_block_header_by_number(
406 &self,
407 block_num: Option<BlockNumber>,
408 include_mmr_proof: bool,
409 ) -> Result<(BlockHeader, Option<MmrProof>), RpcError> {
410 let block = if let Some(block_num) = block_num {
411 self.mock_chain.read().block_header(block_num.as_usize())
412 } else {
413 self.mock_chain.read().latest_block_header()
414 };
415
416 let mmr_proof = if include_mmr_proof {
417 Some(self.get_mmr().open(block_num.unwrap().as_usize()).unwrap())
418 } else {
419 None
420 };
421
422 Ok((block, mmr_proof))
423 }
424
425 async fn get_notes_by_id(&self, note_ids: &[NoteId]) -> Result<Vec<FetchedNote>, RpcError> {
427 let notes = self.mock_chain.read().committed_notes().clone();
429
430 let hit_notes = note_ids.iter().filter_map(|id| notes.get(id));
431 let mut return_notes = vec![];
432 for note in hit_notes {
433 let fetched_note = match note {
434 MockChainNote::Private(note_id, note_metadata, note_inclusion_proof) => {
435 let attachments = self
436 .private_note_attachments
437 .read()
438 .get(note_id)
439 .cloned()
440 .unwrap_or_else(NoteAttachments::empty);
441 FetchedNote::Private(
442 *note_id,
443 *note_metadata,
444 attachments,
445 note_inclusion_proof.clone(),
446 )
447 },
448 MockChainNote::Public(note, note_inclusion_proof) => {
449 FetchedNote::Public(note.clone(), note_inclusion_proof.clone())
450 },
451 };
452 return_notes.push(fetched_note);
453 }
454 Ok(return_notes)
455 }
456
457 async fn submit_proven_transaction(
460 &self,
461 proven_transaction: ProvenTransaction,
462 _tx_inputs: TransactionInputs, ) -> Result<BlockNumber, RpcError> {
464 {
467 let mut mock_chain = self.mock_chain.write();
468 mock_chain.add_pending_proven_transaction(proven_transaction.clone());
469 };
470
471 let block_num = self.get_chain_tip_block_num();
472
473 Ok(block_num)
474 }
475
476 async fn submit_proven_batch(
481 &self,
482 proven_batch: ProvenBatch,
483 _proposed_batch: ProposedBatch,
484 _transaction_inputs: Vec<TransactionInputs>,
485 ) -> Result<BlockNumber, RpcError> {
486 let mut mock_chain = self.mock_chain.write();
487 mock_chain.add_pending_batch(proven_batch);
488 drop(mock_chain);
489
490 let block_num = self.get_chain_tip_block_num();
491
492 Ok(block_num)
493 }
494
495 async fn get_account(
499 &self,
500 account_id: AccountId,
501 request: GetAccountRequest,
502 ) -> Result<(BlockNumber, AccountProof), RpcError> {
503 let current_chain = self.mock_chain.read();
504 let current_block_number = current_chain.latest_block_header().block_num();
505 let block_number = match request.at {
506 AccountStateAt::Block(number) => number,
507 AccountStateAt::ChainTip => current_block_number,
508 };
509 let historical_chain = match request.at {
510 AccountStateAt::Block(_) if block_number != current_block_number => Some(
511 self.historical_chains.read().get(&block_number).cloned().ok_or_else(|| {
512 RpcError::InvalidResponse(alloc::format!(
513 "no mock chain snapshot at block {block_number}"
514 ))
515 })?,
516 ),
517 AccountStateAt::ChainTip | AccountStateAt::Block(_) => None,
518 };
519 let mock_chain = historical_chain.as_deref().unwrap_or(&*current_chain);
520
521 let headers = if account_id.is_public() {
522 let account = mock_chain.committed_account(account_id).unwrap();
523
524 let requested_slots: Vec<_> = match &request.storage {
527 StorageMapFetch::Skip => Vec::new(),
528 StorageMapFetch::Slots(reqs) => reqs.inner().keys().cloned().collect(),
529 StorageMapFetch::All => account
530 .storage()
531 .to_header()
532 .slots()
533 .filter(|slot| slot.slot_type() == StorageSlotType::Map)
534 .map(|slot| slot.name().clone())
535 .collect(),
536 };
537
538 let mut map_details = vec![];
539 for slot_name in &requested_slots {
540 if let Some(StorageSlotContent::Map(storage_map)) =
541 account.storage().get(slot_name).map(StorageSlot::content)
542 {
543 let entries: Vec<StorageMapEntry> = storage_map
544 .entries()
545 .map(|(key, value)| StorageMapEntry { key: *key, value: *value })
546 .collect();
547
548 let too_many_entries = entries.len() > self.oversize_threshold;
551 let account_storage_map_detail = AccountStorageMapDetails {
552 slot_name: slot_name.clone(),
553 too_many_entries,
554 entries: StorageMapEntries::AllEntries(entries),
555 };
556
557 map_details.push(account_storage_map_detail);
558 } else {
559 panic!("Storage slot {slot_name} is not a map");
560 }
561 }
562
563 let storage_details = AccountStorageDetails {
564 header: account.storage().to_header(),
565 map_details,
566 };
567
568 let mut assets = vec![];
569 for asset in account.vault().assets() {
570 assets.push(asset);
571 }
572 let vault_details = AccountVaultDetails {
573 too_many_assets: assets.len() > self.oversize_threshold,
574 assets,
575 };
576
577 Some(AccountDetails {
578 header: account.into(),
579 storage_details,
580 code: account.code().clone(),
581 vault_details,
582 })
583 } else {
584 None
585 };
586
587 let witness = mock_chain.account_tree().open(account_id);
588
589 let proof = AccountProof::new(witness, headers).unwrap();
590
591 Ok((block_number, proof))
592 }
593
594 async fn sync_nullifiers(
597 &self,
598 prefixes: &[u16],
599 block_from: BlockNumber,
600 block_to: BlockNumber,
601 ) -> Result<Vec<NullifierUpdate>, RpcError> {
602 let nullifiers = self
603 .mock_chain
604 .read()
605 .nullifier_tree()
606 .entries()
607 .filter_map(|(nullifier, block_num)| {
608 let within_range = block_num >= block_from && block_num <= block_to;
609
610 if prefixes.contains(&nullifier.prefix()) && within_range {
611 Some(NullifierUpdate { nullifier, block_num })
612 } else {
613 None
614 }
615 })
616 .collect::<Vec<_>>();
617
618 Ok(nullifiers)
619 }
620
621 async fn get_block_by_number(
622 &self,
623 block_num: BlockNumber,
624 _include_proof: bool,
625 ) -> Result<ProvenBlock, RpcError> {
626 let block = self
627 .mock_chain
628 .read()
629 .proven_blocks()
630 .iter()
631 .find(|b| b.header().block_num() == block_num)
632 .unwrap()
633 .clone();
634
635 Ok(block)
636 }
637
638 async fn get_note_script_by_root(&self, root: Word) -> Result<Option<NoteScript>, RpcError> {
639 let script = self
640 .get_available_notes()
641 .iter()
642 .filter_map(|note| note.note())
643 .find(|n| Word::from(n.script().root()) == root)
644 .map(|n| n.script().clone());
645
646 Ok(script)
647 }
648
649 async fn sync_storage_maps(
650 &self,
651 block_from: BlockNumber,
652 block_to: BlockNumber,
653 account_id: AccountId,
654 ) -> Result<StorageMapInfo, RpcError> {
655 let mut all_updates = Vec::new();
656 let mut current_block_from = block_from;
657 let chain_tip = self.get_chain_tip_block_num();
658 let target_block = block_to.min(chain_tip);
659
660 loop {
661 let response =
662 self.get_sync_storage_maps_request(current_block_from, block_to, account_id);
663 all_updates.extend(response.updates);
664
665 if response.block_number >= target_block {
666 return Ok(StorageMapInfo {
667 chain_tip: response.chain_tip,
668 block_number: response.block_number,
669 updates: all_updates,
670 });
671 }
672
673 current_block_from = (response.block_number.as_u32() + 1).into();
674 }
675 }
676
677 async fn sync_account_vault(
678 &self,
679 block_from: BlockNumber,
680 block_to: BlockNumber,
681 account_id: AccountId,
682 ) -> Result<AccountVaultInfo, RpcError> {
683 let mut all_updates = Vec::new();
684 let mut current_block_from = block_from;
685 let chain_tip = self.get_chain_tip_block_num();
686 let target_block = block_to.min(chain_tip);
687
688 loop {
689 let response =
690 self.get_sync_account_vault_request(current_block_from, block_to, account_id);
691 all_updates.extend(response.updates);
692
693 if response.block_number >= target_block {
694 return Ok(AccountVaultInfo {
695 chain_tip: response.chain_tip,
696 block_number: response.block_number,
697 updates: all_updates,
698 });
699 }
700
701 current_block_from = (response.block_number.as_u32() + 1).into();
702 }
703 }
704
705 async fn sync_transactions(
706 &self,
707 block_from: BlockNumber,
708 block_to: BlockNumber,
709 account_ids: Vec<AccountId>,
710 ) -> Result<Vec<TransactionRecord>, RpcError> {
711 Ok(self.get_sync_transactions_request(block_from, block_to, &account_ids))
712 }
713
714 async fn get_network_id(&self) -> Result<NetworkId, RpcError> {
715 Ok(NetworkId::Testnet)
716 }
717
718 async fn get_rpc_limits(&self) -> Result<crate::rpc::RpcLimits, RpcError> {
719 Ok(crate::rpc::RpcLimits::default())
720 }
721
722 fn has_rpc_limits(&self) -> Option<crate::rpc::RpcLimits> {
723 None
724 }
725
726 async fn set_rpc_limits(&self, _limits: crate::rpc::RpcLimits) {
727 }
729
730 async fn get_status_unversioned(&self) -> Result<RpcStatusInfo, RpcError> {
731 Ok(RpcStatusInfo {
732 version: env!("CARGO_PKG_VERSION").into(),
733 genesis_commitment: None,
734 chain_tip: 0,
735 block_producer: None,
736 })
737 }
738
739 async fn get_network_note_status(
740 &self,
741 _note_id: NoteId,
742 ) -> Result<NetworkNoteStatusInfo, RpcError> {
743 todo!("We need to check if we want to implement this for the mockchain");
744 }
745}
746
747impl From<MockChain> for MockRpcApi {
751 fn from(mock_chain: MockChain) -> Self {
752 MockRpcApi::new(mock_chain)
753 }
754}
755
756fn build_account_updates(
760 mock_chain: &MockChain,
761) -> BTreeMap<BlockNumber, BTreeMap<AccountId, Word>> {
762 let mut account_commitment_updates = BTreeMap::new();
763 for block in mock_chain.proven_blocks() {
764 let block_num = block.header().block_num();
765 let mut updates = BTreeMap::new();
766
767 for update in block.body().updated_accounts() {
768 updates.insert(update.account_id(), update.final_state_commitment());
769 }
770
771 if updates.is_empty() {
772 continue;
773 }
774
775 account_commitment_updates.insert(block_num, updates);
776 }
777 account_commitment_updates
778}