1use std::{
2 collections::HashMap,
3 sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard},
4};
5
6use alloy::primitives::{Address, Bytes as AlloyBytes, B256, U256};
7use revm::{
8 context::DBErrorMarker,
9 primitives::KECCAK_EMPTY,
10 state::{AccountInfo, Bytecode},
11 DatabaseRef,
12};
13use thiserror::Error;
14use tracing::{debug, error, instrument, warn};
15use tycho_client::feed::BlockHeader;
16
17use crate::evm::{
18 account_storage::{AccountStorage, StateUpdate},
19 engine_db::engine_db_interface::EngineDatabaseInterface,
20 tycho_models::{AccountUpdate, ChangeType},
21};
22
23#[derive(Error, Debug)]
24pub enum TychoClientError {
25 #[error("Failed to parse URI: {0}. Error: {1}")]
26 UriParsing(String, String),
27 #[error("Failed to format request: {0}")]
28 FormatRequest(String),
29 #[error("Unexpected HTTP client error: {0}")]
30 HttpClient(String),
31 #[error("Failed to parse response: {0}")]
32 ParseResponse(String),
33}
34
35#[derive(Error, Debug)]
36pub enum PreCachedDBError {
37 #[error("Account {0} not found")]
38 MissingAccount(Address),
39 #[error("Bad account update: {0} - {1:?}")]
40 BadUpdate(String, Box<AccountUpdate>),
41 #[error("Block needs to be set")]
42 BlockNotSet(),
43 #[error("Tycho Client error: {0}")]
44 TychoClientError(#[from] TychoClientError),
45 #[error("{0}")]
46 Fatal(String),
47}
48
49impl DBErrorMarker for PreCachedDBError {}
50
51#[derive(Clone, Debug)]
52pub struct PreCachedDBInner {
53 accounts: AccountStorage,
55 block: Option<BlockHeader>,
57}
58
59#[derive(Clone, Debug)]
60pub struct PreCachedDB {
61 pub inner: Arc<RwLock<PreCachedDBInner>>,
67}
68
69impl PreCachedDB {
70 pub fn new() -> Result<Self, PreCachedDBError> {
72 Ok(PreCachedDB {
73 inner: Arc::new(RwLock::new(PreCachedDBInner {
74 accounts: AccountStorage::new(),
75 block: None,
76 })),
77 })
78 }
79
80 #[instrument(skip_all)]
90 pub fn update(
91 &self,
92 account_updates: Vec<AccountUpdate>,
93 block: Option<BlockHeader>,
94 ) -> Result<(), PreCachedDBError> {
95 let mut write_guard = self.write_inner()?;
98
99 write_guard.block = block;
100
101 for update in account_updates {
102 match update.change {
103 ChangeType::Update => {
104 debug!(%update.address, "Updating account");
105
106 write_guard.accounts.update_account(
109 &update.address,
110 &StateUpdate {
111 storage: Some(update.slots.clone()),
112 balance: update.balance,
113 },
114 );
115 }
116 ChangeType::Deletion => {
117 debug!(%update.address, "Deleting account");
118
119 warn!(%update.address, "Deletion not implemented");
120 }
121 ChangeType::Creation => {
122 debug!(%update.address, "Creating account");
123
124 let code = Bytecode::new_raw(AlloyBytes::from(
126 update.code.clone().ok_or_else(|| {
127 error!(%update.address, "MissingCode");
128 PreCachedDBError::BadUpdate(
129 "MissingCode".into(),
130 Box::new(update.clone()),
131 )
132 })?,
133 ));
134 let balance = update.balance.unwrap_or(U256::ZERO);
136
137 write_guard.accounts.overwrite_account(
140 update.address,
141 AccountInfo::new(balance, 0, code.hash_slow(), code),
142 Some(update.slots.clone()),
143 true, );
146 }
147 ChangeType::Unspecified => {
148 warn!(%update.address, "Unspecified change type");
149 }
150 }
151 }
152 Ok(())
153 }
154
155 pub fn force_update_accounts(
159 &self,
160 account_updates: Vec<AccountUpdate>,
161 ) -> Result<(), PreCachedDBError> {
162 let mut write_guard = self.write_inner()?;
163
164 for update in account_updates {
165 if matches!(update.change, ChangeType::Creation) {
166 let code =
167 Bytecode::new_raw(AlloyBytes::from(update.code.clone().ok_or_else(|| {
168 error!(%update.address, "MissingCode");
169 PreCachedDBError::BadUpdate("MissingCode".into(), Box::new(update.clone()))
170 })?));
171 let balance = update.balance.unwrap_or(U256::ZERO);
172
173 write_guard.accounts.overwrite_account(
174 update.address,
175 AccountInfo::new(balance, 0, code.hash_slow(), code),
176 Some(update.slots.clone()),
177 true,
178 );
179 } else {
180 warn!(%update.address, "force_update_accounts called with non-Creation update; ignoring");
181 }
182 }
183 Ok(())
184 }
185
186 pub fn get_storage(&self, address: &Address, index: &U256) -> Option<U256> {
202 self.inner
203 .read()
204 .unwrap()
205 .accounts
206 .get_storage(address, index)
207 }
208
209 pub fn update_state(
219 &mut self,
220 updates: &HashMap<Address, StateUpdate>,
221 block: BlockHeader,
222 ) -> Result<HashMap<Address, StateUpdate>, PreCachedDBError> {
223 let mut write_guard = self.write_inner()?;
226
227 let mut revert_updates = HashMap::new();
228 write_guard.block = Some(block);
229
230 for (address, update_info) in updates.iter() {
231 let mut revert_entry = StateUpdate::default();
232
233 if let Some(current_account) = write_guard
234 .accounts
235 .get_account_info(address)
236 {
237 revert_entry.balance = Some(current_account.balance);
238 }
239
240 if let Some(storage) = &update_info.storage {
241 let mut revert_storage = HashMap::default();
242 for index in storage.keys() {
243 if let Some(s) = write_guard
244 .accounts
245 .get_storage(address, index)
246 {
247 revert_storage.insert(*index, s);
248 }
249 }
250 revert_entry.storage = Some(revert_storage);
251 }
252 revert_updates.insert(*address, revert_entry);
253 write_guard
254 .accounts
255 .update_account(address, update_info);
256 }
257
258 Ok(revert_updates)
259 }
260
261 #[cfg(test)]
262 pub fn get_account_storage(&self) -> Result<AccountStorage, PreCachedDBError> {
263 self.read_inner()
264 .map(|guard| guard.accounts.clone())
265 }
266
267 pub fn block_number(&self) -> Result<Option<u64>, PreCachedDBError> {
269 self.read_inner().map(|guard| {
270 guard
271 .block
272 .as_ref()
273 .map(|header| header.number)
274 })
275 }
276
277 pub fn clear(&self) -> Result<(), PreCachedDBError> {
279 let mut write_guard = self.write_inner()?;
280 write_guard.accounts.clear();
281 write_guard.block = None;
282 Ok(())
283 }
284
285 fn read_inner(&self) -> Result<RwLockReadGuard<'_, PreCachedDBInner>, PreCachedDBError> {
286 self.inner
287 .read()
288 .map_err(|_| PreCachedDBError::Fatal("Tycho state db lock poisoned".into()))
289 }
290
291 fn write_inner(&self) -> Result<RwLockWriteGuard<'_, PreCachedDBInner>, PreCachedDBError> {
292 self.inner
293 .write()
294 .map_err(|_| PreCachedDBError::Fatal("Tycho state db lock poisoned".into()))
295 }
296}
297
298impl EngineDatabaseInterface for PreCachedDB {
299 type Error = PreCachedDBError;
300
301 fn init_account(
312 &self,
313 address: Address,
314 account: AccountInfo,
315 permanent_storage: Option<HashMap<U256, U256>>,
316 _mocked: bool,
317 ) -> Result<(), <Self as EngineDatabaseInterface>::Error> {
318 if account.code.is_none() && account.code_hash != KECCAK_EMPTY {
319 warn!("Code is None for account {address} but code hash is not KECCAK_EMPTY");
320 } else if account.code.is_some() && account.code_hash == KECCAK_EMPTY {
321 warn!("Code is Some for account {address} but code hash is KECCAK_EMPTY");
322 }
323
324 self.write_inner()?
325 .accounts
326 .init_account(address, account, permanent_storage, true);
327
328 Ok(())
329 }
330
331 fn clear_temp_storage(&mut self) -> Result<(), <Self as EngineDatabaseInterface>::Error> {
333 debug!("Temp storage in TychoDB is never set, nothing to clear");
334
335 Ok(())
336 }
337
338 fn get_current_block(&self) -> Option<BlockHeader> {
339 self.inner.read().unwrap().block.clone()
340 }
341}
342
343impl DatabaseRef for PreCachedDB {
344 type Error = PreCachedDBError;
345 fn basic_ref(&self, address: Address) -> Result<Option<AccountInfo>, Self::Error> {
358 self.read_inner()?
359 .accounts
360 .get_account_info(&address)
361 .map(|acc| Some(acc.clone()))
362 .ok_or(PreCachedDBError::MissingAccount(address))
363 }
364
365 fn code_by_hash_ref(&self, code_hash: B256) -> Result<Bytecode, Self::Error> {
366 Err(PreCachedDBError::Fatal(format!("Code by hash not supported: {code_hash}")))
367 }
368
369 fn storage_ref(&self, address: Address, index: U256) -> Result<U256, Self::Error> {
384 debug!(%address, %index, "Requested storage of account");
385 let read_guard = self.read_inner()?;
386 if let Some(storage_value) = read_guard
387 .accounts
388 .get_storage(&address, &index)
389 {
390 debug!(%address, %index, %storage_value, "Got value locally");
391 Ok(storage_value)
392 } else {
393 if read_guard
395 .accounts
396 .account_present(&address)
397 {
398 debug!(%address, %index, "Account found, but slot is zero");
401 Ok(U256::ZERO)
402 } else {
403 debug!(%address, %index, "Account not found");
405 Err(PreCachedDBError::MissingAccount(address))
406 }
407 }
408 }
409
410 fn block_hash_ref(&self, _number: u64) -> Result<B256, Self::Error> {
412 match self.read_inner()?.block.clone() {
413 Some(header) => Ok(B256::from_slice(&header.hash)),
414 None => Ok(B256::default()),
415 }
416 }
417}
418
419#[cfg(test)]
420mod tests {
421 use std::{error::Error, str::FromStr};
422
423 use revm::primitives::U256;
424 use rstest::{fixture, rstest};
425 use tycho_common::Bytes;
426
427 use super::*;
428 use crate::evm::tycho_models::{AccountUpdate, Chain, ChangeType};
429
430 #[fixture]
431 pub fn mock_db() -> PreCachedDB {
432 PreCachedDB {
433 inner: Arc::new(RwLock::new(PreCachedDBInner {
434 accounts: AccountStorage::new(),
435 block: None,
436 })),
437 }
438 }
439
440 #[rstest]
441 #[tokio::test]
442 async fn test_account_get_acc_info(mock_db: PreCachedDB) -> Result<(), Box<dyn Error>> {
443 let mock_acc_address = Address::from_str("0xb4e16d0168e52d35cacd2c6185b44281ec28c9dc")?;
446 mock_db
447 .init_account(mock_acc_address, AccountInfo::default(), None, false)
448 .expect("Account init should succeed");
449
450 let acc_info = mock_db
451 .basic_ref(mock_acc_address)
452 .unwrap()
453 .unwrap();
454
455 assert_eq!(
456 mock_db
457 .basic_ref(mock_acc_address)
458 .unwrap()
459 .unwrap(),
460 acc_info
461 );
462 Ok(())
463 }
464
465 #[rstest]
466 fn test_account_storage(mock_db: PreCachedDB) -> Result<(), Box<dyn Error>> {
467 let mock_acc_address = Address::from_str("0xb4e16d0168e52d35cacd2c6185b44281ec28c9dc")?;
468 let storage_address = U256::from(1);
469 let mut permanent_storage: HashMap<U256, U256> = HashMap::new();
470 permanent_storage.insert(storage_address, U256::from(10));
471 mock_db
472 .init_account(mock_acc_address, AccountInfo::default(), Some(permanent_storage), false)
473 .expect("Account init should succeed");
474
475 let storage = mock_db
476 .storage_ref(mock_acc_address, storage_address)
477 .unwrap();
478
479 assert_eq!(storage, U256::from(10));
480 Ok(())
481 }
482
483 #[rstest]
484 fn test_account_storage_zero(mock_db: PreCachedDB) -> Result<(), Box<dyn Error>> {
485 let mock_acc_address = Address::from_str("0xb4e16d0168e52d35cacd2c6185b44281ec28c9dc")?;
486 let storage_address = U256::from(1);
487 mock_db
488 .init_account(mock_acc_address, AccountInfo::default(), None, false)
489 .expect("Account init should succeed");
490
491 let storage = mock_db
492 .storage_ref(mock_acc_address, storage_address)
493 .unwrap();
494
495 assert_eq!(storage, U256::ZERO);
496 Ok(())
497 }
498
499 #[rstest]
500 #[should_panic(
501 expected = "called `Result::unwrap()` on an `Err` value: MissingAccount(0xb4e16d0168e52d35cacd2c6185b44281ec28c9dc)"
502 )]
503 fn test_account_storage_missing(mock_db: PreCachedDB) {
504 let mock_acc_address =
505 Address::from_str("0xb4e16d0168e52d35cacd2c6185b44281ec28c9dc").unwrap();
506 let storage_address = U256::from(1);
507
508 mock_db
510 .storage_ref(mock_acc_address, storage_address)
511 .unwrap();
512 }
513
514 fn header(number: u64) -> BlockHeader {
515 BlockHeader { number, ..Default::default() }
516 }
517
518 fn snapshot(address: Address, slots: &[(u64, u64)], balance: u64) -> AccountUpdate {
519 AccountUpdate {
520 address,
521 chain: Chain::Ethereum,
522 slots: slots
523 .iter()
524 .map(|(k, v)| (U256::from(*k), U256::from(*v)))
525 .collect(),
526 balance: Some(U256::from(balance)),
527 code: Some(vec![0x60, 0x00]),
528 change: ChangeType::Creation,
529 }
530 }
531
532 #[rstest]
538 fn test_creation_reapply_overwrites_existing_account(
539 mock_db: PreCachedDB,
540 ) -> Result<(), Box<dyn Error>> {
541 let address = Address::from_str("0xd51a44d3fae010294c616388b506acda1bfaae46")?;
542 let slot = U256::from(1);
543
544 mock_db.update(vec![snapshot(address, &[(1, 100)], 10)], Some(header(1)))?;
545 mock_db.update(vec![snapshot(address, &[(1, 999), (2, 42)], 20)], Some(header(2)))?;
547
548 assert_eq!(
549 mock_db.get_storage(&address, &slot),
550 Some(U256::from(999)),
551 "re-applied snapshot must overwrite a stale slot on an existing account"
552 );
553 assert_eq!(
554 mock_db.get_storage(&address, &U256::from(2)),
555 Some(U256::from(42)),
556 "slots first seen in the re-applied snapshot must be present"
557 );
558 assert_eq!(
559 mock_db
560 .basic_ref(address)
561 .unwrap()
562 .map(|acc| acc.balance),
563 Some(U256::from(20)),
564 "re-applied snapshot must refresh the native balance"
565 );
566 Ok(())
567 }
568
569 #[rstest]
572 fn test_delta_update_refreshes_existing_slot(
573 mock_db: PreCachedDB,
574 ) -> Result<(), Box<dyn Error>> {
575 let address = Address::from_str("0xd51a44d3fae010294c616388b506acda1bfaae46")?;
576 let slot = U256::from(1);
577
578 mock_db.update(vec![snapshot(address, &[(1, 100)], 10)], Some(header(1)))?;
579 let delta = AccountUpdate {
580 address,
581 chain: Chain::Ethereum,
582 slots: HashMap::from([(slot, U256::from(200))]),
583 balance: None,
584 code: None,
585 change: ChangeType::Update,
586 };
587 mock_db.update(vec![delta], Some(header(2)))?;
588
589 assert_eq!(
590 mock_db.get_storage(&address, &slot),
591 Some(U256::from(200)),
592 "delta updates must refresh slots on existing accounts"
593 );
594 Ok(())
595 }
596
597 #[rstest]
598 #[tokio::test]
599 async fn test_update_state(mut mock_db: PreCachedDB) -> Result<(), Box<dyn Error>> {
600 let address = Address::from_str("0xb4e16d0168e52d35cacd2c6185b44281ec28c9dc")?;
601 mock_db
602 .init_account(address, AccountInfo::default(), None, false)
603 .expect("Account init should succeed");
604
605 let mut new_storage = HashMap::default();
606 let new_storage_value_index = U256::from_limbs_slice(&[123]);
607 new_storage.insert(new_storage_value_index, new_storage_value_index);
608 let new_balance = U256::from_limbs_slice(&[500]);
609 let update = StateUpdate { storage: Some(new_storage), balance: Some(new_balance) };
610 let new_block = BlockHeader {
611 number: 1,
612 hash: Bytes::from_str(
613 "0xc6b994ec855fb2b31013c7ae65074406fac46679b5b963469104e0bfeddd66d9",
614 )
615 .unwrap(),
616 timestamp: 123,
617 ..Default::default()
618 };
619 let mut updates = HashMap::default();
620 updates.insert(address, update);
621
622 mock_db
623 .update_state(&updates, new_block)
624 .expect("State update should succeed");
625
626 assert_eq!(
627 mock_db
628 .get_storage(&address, &new_storage_value_index)
629 .unwrap(),
630 new_storage_value_index
631 );
632 let account_info = mock_db
633 .basic_ref(address)
634 .unwrap()
635 .unwrap();
636 assert_eq!(account_info.balance, new_balance);
637 let block = mock_db
638 .inner
639 .read()
640 .unwrap()
641 .block
642 .clone()
643 .expect("block is Some");
644 assert_eq!(block.number, 1);
645
646 Ok(())
647 }
648
649 #[rstest]
650 #[tokio::test]
651 async fn test_block_number_getter(mut mock_db: PreCachedDB) -> Result<(), Box<dyn Error>> {
652 let address = Address::from_str("0xb4e16d0168e52d35cacd2c6185b44281ec28c9dc")?;
653 mock_db
654 .init_account(address, AccountInfo::default(), None, false)
655 .expect("Account init should succeed");
656
657 let new_block = BlockHeader {
658 number: 1,
659 hash: Bytes::from_str(
660 "0xc6b994ec855fb2b31013c7ae65074406fac46679b5b963469104e0bfeddd66d9",
661 )
662 .unwrap(),
663 timestamp: 123,
664 ..Default::default()
665 };
666 let updates = HashMap::default();
667
668 mock_db
669 .update_state(&updates, new_block)
670 .expect("State update should succeed");
671
672 let block_number = mock_db.block_number();
673 assert_eq!(block_number.unwrap().unwrap(), 1);
674
675 Ok(())
676 }
677
678 #[rstest]
679 #[tokio::test]
680 async fn test_update() {
681 let mock_db = PreCachedDB {
682 inner: Arc::new(RwLock::new(PreCachedDBInner {
683 accounts: AccountStorage::new(),
684 block: None,
685 })),
686 };
687
688 let account_update = AccountUpdate::new(
689 Address::from_str("0x7a250d5630B4cF539739dF2C5dAcb4c659F2488D").unwrap(),
690 Chain::Ethereum,
691 HashMap::new(),
692 Some(U256::from(500)),
693 Some(Vec::<u8>::new()),
694 ChangeType::Creation,
695 );
696
697 let new_block = BlockHeader {
698 number: 1,
699 hash: Bytes::from_str(
700 "0xc6b994ec855fb2b31013c7ae65074406fac46679b5b963469104e0bfeddd66d9",
701 )
702 .unwrap(),
703 timestamp: 123,
704 ..Default::default()
705 };
706
707 mock_db
708 .update(vec![account_update], Some(new_block))
709 .unwrap();
710
711 let account_info = mock_db
712 .basic_ref(Address::from_str("0x7a250d5630B4cF539739dF2C5dAcb4c659F2488D").unwrap())
713 .unwrap()
714 .unwrap();
715
716 assert_eq!(
717 account_info,
718 AccountInfo {
719 nonce: 0,
720 balance: U256::from(500),
721 code_hash: B256::from_str(
722 "0xc5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a470"
723 )
724 .unwrap(),
725 code: Some(Bytecode::default()),
726 }
727 );
728
729 assert_eq!(
730 mock_db
731 .inner
732 .read()
733 .unwrap()
734 .block
735 .clone()
736 .expect("block is Some")
737 .number,
738 1
739 );
740 }
741
742 #[rstest]
743 fn test_force_update_accounts_overwrites_existing(mock_db: PreCachedDB) {
744 let address = Address::from_str("0xb4e16d0168e52d35cacd2c6185b44281ec28c9dc").unwrap();
745
746 mock_db
748 .init_account(address, AccountInfo::default(), None, true)
749 .expect("placeholder init should succeed");
750
751 let storage_slot = U256::from(1);
753 let storage_value = U256::from(42);
754 let mut slots = HashMap::new();
755 slots.insert(storage_slot, storage_value);
756 let update = AccountUpdate::new(
757 address,
758 Chain::Ethereum,
759 slots,
760 Some(U256::from(100)),
761 Some(Vec::<u8>::new()),
762 ChangeType::Creation,
763 );
764 mock_db
765 .force_update_accounts(vec![update])
766 .expect("force update should succeed");
767
768 let info = mock_db
769 .basic_ref(address)
770 .unwrap()
771 .unwrap();
772 assert_eq!(info.balance, U256::from(100));
773 assert_eq!(
774 mock_db
775 .get_storage(&address, &storage_slot)
776 .unwrap(),
777 storage_value
778 );
779 }
780
781 #[rstest]
782 fn test_force_update_accounts_non_creation_ignored(mock_db: PreCachedDB) {
783 let address = Address::from_str("0xb4e16d0168e52d35cacd2c6185b44281ec28c9dc").unwrap();
784 let original_balance = U256::from(10);
785
786 mock_db
787 .init_account(
788 address,
789 AccountInfo { balance: original_balance, ..Default::default() },
790 None,
791 true,
792 )
793 .expect("init should succeed");
794
795 let update = AccountUpdate::new(
797 address,
798 Chain::Ethereum,
799 HashMap::new(),
800 Some(U256::from(999)),
801 None,
802 ChangeType::Update,
803 );
804 mock_db
805 .force_update_accounts(vec![update])
806 .expect("force update should succeed");
807
808 let info = mock_db
810 .basic_ref(address)
811 .unwrap()
812 .unwrap();
813 assert_eq!(info.balance, original_balance);
814 }
815
816 #[ignore]
832 #[rstest]
833 fn test_tycho_db_connection() {
834 tracing_subscriber::fmt()
835 .with_env_filter("debug")
836 .init();
837
838 let ambient_contract =
839 Address::from_str("0xaaaaaaaaa24eeeb8d57d431224f73832bc34f688").unwrap();
840
841 let db = PreCachedDB::new().expect("db should initialize");
842
843 let acc_info = db
844 .basic_ref(ambient_contract)
845 .unwrap()
846 .unwrap();
847
848 debug!(?acc_info, "Account info");
849 }
850}