feat: adds zk cheatcode for retrieving bytecode hash

crates/cheatcodes/spec/src/vm.rs

@@ -465,6 +465,10 @@ interface Vm {
     #[cheatcode(group = Evm, safety = Unsafe, status = Experimental)]
     function cool(address target) external;
 
+    /// Gets the deployed bytecodeHash from state. Takes in the address of the contract.
+    #[cheatcode(group = Evm, safety = Safe)]
+    function getRawCodeHash(address target) external view returns (bytes memory byteCodeHash);
+
     // -------- Call Manipulation --------
     // --- Mocks ---
 

crates/cheatcodes/src/evm.rs

@@ -598,6 +598,18 @@ impl Cheatcode for setBlockhashCall {
     }
 }
 
+impl Cheatcode for getRawCodeHashCall {
+    fn apply_stateful<DB: DatabaseExt>(&self, ccx: &mut CheatsCtxt<DB>) -> Result {
+        let Self { target } = self;
+        if ccx.state.use_zk_vm {
+            let code = foundry_zksync_core::cheatcodes::get_raw_code_hash(*target, ccx.ecx);
+            return Ok(code.abi_encode());
+        }
+
+        Ok(Bytecode::new().bytes().abi_encode())
+    }
+}
+
 pub(super) fn get_nonce<DB: DatabaseExt>(ccx: &mut CheatsCtxt<DB>, address: &Address) -> Result {
     let (account, _) = ccx.ecx.journaled_state.load_account(*address, &mut ccx.ecx.db)?;
     Ok(account.info.nonce.abi_encode())

crates/forge/tests/it/zk/cheats.rs

@@ -121,3 +121,11 @@ async fn test_zk_cheatcodes_in_zkvm() {
 
     TestConfig::with_filter(runner, filter).evm_spec(SpecId::SHANGHAI).run().await;
 }
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_zk_contract_get_raw_code_hash() {
+    let runner = TEST_DATA_DEFAULT.runner_zksync();
+    let filter = Filter::new("testZkGetRawCodeHash", "ZkCheatcodesGetRawCodeHashTest", ".*");
+
+    TestConfig::with_filter(runner, filter).evm_spec(SpecId::SHANGHAI).run().await;
+}

crates/zksync/core/src/cheatcodes.rs

@@ -199,6 +199,32 @@ where
     }
 }
 
+    /// Retrieves the bytecode hash for a given address.
+    pub fn get_raw_code_hash<DB>(address: Address, ecx: &mut InnerEvmContext<DB>) -> rU256
+    where
+        DB: Database,
+        <DB as Database>::Error: Debug,
+    {
+    info!(?address, "cheatcode getRawCodeHash");
+
+    // Load the account code storage system address
+    let account_code_addr = ACCOUNT_CODE_STORAGE_ADDRESS.to_address();
+    ecx.load_account(account_code_addr).expect("account 'ACCOUNT_CODE_STORAGE_ADDRESS' could not be loaded");
+
+    let zk_address = address.to_h160();
+    let account_key = zk_address.to_h256().to_ru256();
+
+    // Load the bytecode hash from the system storage
+    let (bytecode_hash, _) = ecx.sload(account_code_addr, account_key).unwrap_or_default();
+
+    if bytecode_hash.is_zero() {
+        info!(?address, "no bytecode found for address");
+        return rU256::ZERO;
+    }
+
+    bytecode_hash
+}
+
 #[cfg(test)]
 mod tests {
     use revm::db::EmptyDB;