#[cfg(not(feature = "std"))]
use alloc as std;
use core::{
cell::RefCell,
convert::Infallible,
fmt::{Debug, Display},
};
use std::{rc::Rc, vec::Vec};
use alloy_primitives::{Address, BlockNumber, Bytes, B256, U256};
use revm::primitives::HashMap;
use crate::{BucketId, ExternalEnvFactory, ExternalEnvTypes, ExternalEnvs, OracleEnv, SaltEnv};
pub trait BucketHasher: Debug + Clone + Unpin + 'static {
fn bucket_id(key: &[u8]) -> BucketId;
}
#[derive(Debug, Clone, Copy)]
pub struct SimpleBucketHasher;
impl BucketHasher for SimpleBucketHasher {
fn bucket_id(key: &[u8]) -> BucketId {
let mut hash: u64 = 0xcbf29ce484222325;
for &byte in key {
hash ^= byte as u64;
hash = hash.wrapping_mul(0x100000001b3);
}
const NUM_BUCKETS: u64 = 1 << 24; const NUM_META_BUCKETS: u64 = NUM_BUCKETS / 256; const NUM_KV_BUCKETS: u64 = NUM_BUCKETS - NUM_META_BUCKETS;
(hash % NUM_KV_BUCKETS + NUM_META_BUCKETS) as BucketId
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RecordedHint {
pub from: Address,
pub topic: B256,
pub data: Bytes,
}
#[derive(derive_more::Debug, Clone)]
pub struct TestExternalEnvs<Error = Infallible, Hasher = SimpleBucketHasher> {
#[debug(ignore)]
_phantom: core::marker::PhantomData<(Error, Hasher)>,
oracle_storage: Rc<RefCell<HashMap<U256, U256>>>,
bucket_capacity: Rc<RefCell<HashMap<BucketId, u64>>>,
recorded_hints: Rc<RefCell<Vec<RecordedHint>>>,
}
impl Default for TestExternalEnvs {
fn default() -> Self {
Self::new()
}
}
impl<H: BucketHasher> From<TestExternalEnvs<Infallible, H>>
for ExternalEnvs<TestExternalEnvs<Infallible, H>>
{
fn from(value: TestExternalEnvs<Infallible, H>) -> Self {
Self { salt_env: value.clone(), oracle_env: value }
}
}
impl<'a, H: BucketHasher> From<&'a TestExternalEnvs<Infallible, H>>
for ExternalEnvs<&'a TestExternalEnvs<Infallible, H>>
{
fn from(value: &'a TestExternalEnvs<Infallible, H>) -> Self {
ExternalEnvs { salt_env: value.clone(), oracle_env: value.clone() }
}
}
impl<Error: Unpin + Clone + Display + 'static, Hasher: BucketHasher>
TestExternalEnvs<Error, Hasher>
{
pub fn new() -> Self {
Self {
_phantom: core::marker::PhantomData,
oracle_storage: Rc::new(RefCell::new(HashMap::default())),
bucket_capacity: Rc::new(RefCell::new(HashMap::default())),
recorded_hints: Rc::new(RefCell::new(Vec::new())),
}
}
pub fn recorded_hints(&self) -> Vec<RecordedHint> {
self.recorded_hints.borrow().clone()
}
pub fn clear_recorded_hints(&self) {
self.recorded_hints.borrow_mut().clear();
}
pub fn with_bucket_capacity(self, bucket_id: BucketId, capacity: u64) -> Self {
self.bucket_capacity.borrow_mut().insert(bucket_id, capacity);
self
}
pub fn clear_bucket_capacity(&self) {
self.bucket_capacity.borrow_mut().clear();
}
pub fn with_oracle_storage(self, slot: U256, value: U256) -> Self {
self.oracle_storage.borrow_mut().insert(slot, value);
self
}
pub fn clear_oracle_storage(&self) {
self.oracle_storage.borrow_mut().clear();
}
}
impl<Error: Unpin + Clone + Display, Hasher: BucketHasher> ExternalEnvFactory
for TestExternalEnvs<Error, Hasher>
{
type EnvTypes = Self;
fn external_envs(&self, _block: BlockNumber) -> ExternalEnvs<Self::EnvTypes> {
ExternalEnvs { salt_env: self.clone(), oracle_env: self.clone() }
}
}
impl<Error: Unpin + Display, Hasher: BucketHasher> ExternalEnvTypes
for TestExternalEnvs<Error, Hasher>
{
type SaltEnv = Self;
type OracleEnv = Self;
}
const SLOT_KEY_LEN: usize = B256::len_bytes();
const PLAIN_ACCOUNT_KEY_LEN: usize = Address::len_bytes();
const PLAIN_STORAGE_KEY_LEN: usize = PLAIN_ACCOUNT_KEY_LEN + SLOT_KEY_LEN;
impl<Error: Unpin + Display, Hasher: BucketHasher> SaltEnv for TestExternalEnvs<Error, Hasher> {
type Error = Error;
fn get_bucket_capacity(&self, bucket_id: BucketId) -> Result<u64, Self::Error> {
Ok(self
.bucket_capacity
.borrow()
.get(&bucket_id)
.copied()
.unwrap_or(crate::MIN_BUCKET_SIZE as u64))
}
fn bucket_id_for_account(account: Address) -> BucketId {
Hasher::bucket_id(account.as_slice())
}
fn bucket_id_for_slot(address: Address, key: U256) -> BucketId {
Hasher::bucket_id(
address.concat_const::<SLOT_KEY_LEN, PLAIN_STORAGE_KEY_LEN>(key.into()).as_slice(),
)
}
}
impl<Error: Unpin + Display, Hasher: BucketHasher> OracleEnv for TestExternalEnvs<Error, Hasher> {
fn get_oracle_storage(&self, slot: U256) -> Option<U256> {
self.oracle_storage.borrow().get(&slot).copied()
}
fn on_hint(&self, from: Address, topic: B256, data: Bytes) {
self.recorded_hints.borrow_mut().push(RecordedHint { from, topic, data });
}
}