Skip to content

Commit

Permalink
add create and select blocks
Browse files Browse the repository at this point in the history 
Signed-off-by: Danil <[email protected]>
  • Loading branch information
@Deniallugo
Deniallugo committed Dec 18, 2023
1 parent c0b3a72 commit 6c87893
Show file tree
Hide file tree
Showing 2 changed files with 124 additions and 6 deletions.
106 changes: 100 additions & 6 deletions crates/era-cheatcodes/src/cheatcodes.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ use multivm::{
},
};
use revm::{
primitives::{BlockEnv, CfgEnv, Env, SpecId},
primitives::{BlockEnv, CfgEnv, Env, SpecId, U256 as rU256},
JournaledState,
};
use serde::Serialize;
Expand Down Expand Up @@ -118,7 +118,9 @@ impl LogEntry {
enum FinishCycleOneTimeActions {
StorageWrite { key: StorageKey, read_value: H256, write_value: H256 },
StoreFactoryDep { hash: U256, bytecode: Vec<U256> },
CreateFork { url_or_alias: String, block_number: u64 },
CreateSelectFork { url_or_alias: String, block_number: Option<u64> },
CreateFork { url_or_alias: String, block_number: Option<u64> },
SelectFork { fork_id: U256 },
}

#[derive(Debug, Default, Clone)]
Expand Down Expand Up @@ -251,7 +253,7 @@ impl<S: DatabaseExt + Send, H: HistoryMode> VmTracer<EraDb<S>, H> for CheatcodeT
FinishCycleOneTimeActions::StoreFactoryDep { hash, bytecode } => state
.decommittment_processor
.populate(vec![(hash, bytecode)], Timestamp(state.local_state.timestamp)),
FinishCycleOneTimeActions::CreateFork { url_or_alias, block_number } => {
FinishCycleOneTimeActions::CreateSelectFork { url_or_alias, block_number } => {
let modified_storage: HashMap<StorageKey, H256> = storage
.borrow_mut()
.modified_storage_keys()
Expand Down Expand Up @@ -284,7 +286,7 @@ impl<S: DatabaseExt + Send, H: HistoryMode> VmTracer<EraDb<S>, H> for CheatcodeT
create_fork_request(
era_env,
self.config.clone(),
Some(block_number),
block_number,
&url_or_alias,
),
&mut env,
Expand All @@ -299,6 +301,65 @@ impl<S: DatabaseExt + Send, H: HistoryMode> VmTracer<EraDb<S>, H> for CheatcodeT

self.return_data = Some(vec![fork_id.unwrap().to_u256()]);
}
FinishCycleOneTimeActions::CreateFork { url_or_alias, block_number } => {
let handle: &ForkStorage<RevmDatabaseForEra<S>> =
&storage.borrow_mut().storage_handle;
let era_db =
handle.inner.write().unwrap().fork.as_ref().unwrap().fork_source.clone();

let mut db = era_db.db.lock().unwrap();
let era_env = self.env.get().unwrap();
let fork_id = db.create_fork(create_fork_request(
era_env,
self.config.clone(),
block_number,
&url_or_alias,
));
self.return_data = Some(vec![fork_id.unwrap().to_u256()]);
}
FinishCycleOneTimeActions::SelectFork { fork_id } => {
let modified_storage: HashMap<StorageKey, H256> = storage
.borrow_mut()
.modified_storage_keys()
.clone()
.into_iter()
.filter(|(key, _)| key.address() != &zksync_types::SYSTEM_CONTEXT_ADDRESS)
.collect();
{
storage.borrow_mut().clean_cache();
let handle: &ForkStorage<RevmDatabaseForEra<S>> =
&storage.borrow_mut().storage_handle;
let mut fork_storage = handle.inner.write().unwrap();
fork_storage.value_read_cache.clear();
let era_db = fork_storage.fork.as_ref().unwrap().fork_source.clone();
let bytecodes = bootloader_state
.get_last_tx_compressed_bytecodes()
.iter()
.map(|b| bytecode_to_factory_dep(b.original.clone()))
.collect();

let mut journaled_state = JournaledState::new(SpecId::LATEST, vec![]);
let state =
storage_to_state(modified_storage.clone(), bytecodes, era_db.clone());
*journaled_state.state() = state;

let mut db = era_db.db.lock().unwrap();
let era_env = self.env.get().unwrap();
let mut env = into_revm_env(era_env);
db.select_fork(
rU256::from(fork_id.as_u128()),
&mut env,
&mut journaled_state,
)
.unwrap();
drop(db);
let mut db_env = era_db.env.lock().unwrap();
*db_env = env;
}
storage.borrow_mut().modified_storage_keys = modified_storage;

self.return_data = Some(vec![fork_id]);
}
}
}

Expand Down Expand Up @@ -748,15 +809,48 @@ impl CheatcodeTracer {
&mut storage,
);
}
createSelectFork_1(createSelectFork_1Call { urlOrAlias, blockNumber }) => {
createSelectFork_0(createSelectFork_0Call { urlOrAlias }) => {
tracing::info!("👷 Creating and selecting fork {}", urlOrAlias,);

self.one_time_actions.push(FinishCycleOneTimeActions::CreateSelectFork {
url_or_alias: urlOrAlias,
block_number: None,
});
}
createSelectFork_1(createSelectFork_1Call { urlOrAlias, blockNumber }) => {
let block_number = blockNumber.to_u256().as_u64();
tracing::info!(
"👷 Creating and selecting fork {} for block number {}",
urlOrAlias,
block_number
);
self.one_time_actions.push(FinishCycleOneTimeActions::CreateSelectFork {
url_or_alias: urlOrAlias,
block_number: Some(block_number),
});
}
createFork_0(createFork_0Call { urlOrAlias }) => {
tracing::info!("👷 Creating fork {}", urlOrAlias,);

self.one_time_actions.push(FinishCycleOneTimeActions::CreateFork {
url_or_alias: urlOrAlias,
block_number: None,
});
}
createFork_1(createFork_1Call { urlOrAlias, blockNumber }) => {
let block_number = blockNumber.to_u256().as_u64();
tracing::info!("👷 Creating fork {} for block number {}", urlOrAlias, block_number);
self.one_time_actions.push(FinishCycleOneTimeActions::CreateFork {
url_or_alias: urlOrAlias,
block_number,
block_number: Some(block_number),
});
}
selectFork(selectForkCall { forkId }) => {
tracing::info!("👷 Selecting fork {}", forkId);

self.one_time_actions
.push(FinishCycleOneTimeActions::SelectFork { fork_id: forkId.to_u256() });
}
writeFile(writeFileCall { path, data }) => {
tracing::info!("👷 Writing data to file in path {}", path);
if fs::write(path, data).is_err() {
Expand Down
24 changes: 24 additions & 0 deletions crates/era-cheatcodes/tests/src/cheatcodes/Fork.t.sol
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,4 +35,28 @@ contract ForkTest is Test {
require(block.number == FORK_BLOCK + 1, "ENV for blocks is not set correctly");

}

function testCreateSelectFork() public{
(bool success, bytes memory data) = Constants.CHEATCODE_ADDRESS.call(
abi.encodeWithSignature("createFork(string,uint256)", "mainnet", FORK_BLOCK + 100)
);
require(success, "fork failed");

uint256 forkId = uint256(bytes32(data));
(bool success1, bytes memory data1) = Constants.CHEATCODE_ADDRESS.call(
abi.encodeWithSignature("selectFork(uint256)", forkId)
);
require(success1, "select fork failed");


/// After createSelect fork the decimals should exist
(bool success2, bytes memory data2) = TOKEN_ADDRESS.call(
abi.encodeWithSignature("decimals()")
);
uint256 decimals_after = uint256(bytes32(data2));
console.log("decimals_after", decimals_after);
console.log("block ", block.number);
require(decimals_after == TOKEN_DECIMALS, "Contract dosent exists in fork");
require(block.number == FORK_BLOCK + 100, "ENV for blocks is not set correctly");
}
}

0 comments on commit 6c87893

Please sign in to comment.