From b648d5ad410a1b1840c24bcd2c27a43f105a401f Mon Sep 17 00:00:00 2001 From: DaniPopes <57450786+DaniPopes@users.noreply.github.com> Date: Tue, 9 Jan 2024 18:08:39 +0100 Subject: [PATCH] chore: new-type IC/PC maps (#6742) * chore: move pc_ic map to a new module * chore: create new types for pc-ic maps --- crates/cast/bin/cmd/storage.rs | 14 ++--- crates/common/src/contracts.rs | 8 +-- crates/debugger/src/tui/draw.rs | 4 +- crates/debugger/src/tui/mod.rs | 23 ++------ crates/evm/core/src/ic.rs | 78 +++++++++++++++++++++++++ crates/evm/core/src/lib.rs | 2 + crates/evm/core/src/utils.rs | 69 +++------------------- crates/evm/coverage/src/anchors.rs | 8 +-- crates/evm/traces/src/lib.rs | 1 + crates/forge/bin/cmd/coverage.rs | 8 +-- crates/forge/bin/cmd/script/executor.rs | 1 - 11 files changed, 112 insertions(+), 104 deletions(-) create mode 100644 crates/evm/core/src/ic.rs diff --git a/crates/cast/bin/cmd/storage.rs b/crates/cast/bin/cmd/storage.rs index b91258c86664a..d574883392ee0 100644 --- a/crates/cast/bin/cmd/storage.rs +++ b/crates/cast/bin/cmd/storage.rs @@ -17,7 +17,9 @@ use foundry_common::{ provider::ethers::RetryProvider, types::{ToAlloy, ToEthers}, }; -use foundry_compilers::{artifacts::StorageLayout, ConfigurableContractArtifact, Project, Solc}; +use foundry_compilers::{ + artifacts::StorageLayout, Artifact, ConfigurableContractArtifact, Project, Solc, +}; use foundry_config::{ figment::{self, value::Dict, Metadata, Profile}, impl_figment_convert_cast, Config, @@ -101,13 +103,9 @@ impl StorageArgs { // Find in artifacts and pretty print add_storage_layout_output(&mut project); let out = ProjectCompiler::new().compile(&project)?; - let match_code = |artifact: &ConfigurableContractArtifact| -> Option { - let bytes = - artifact.deployed_bytecode.as_ref()?.bytecode.as_ref()?.object.as_bytes()?; - Some(bytes == &address_code) - }; - let artifact = - out.artifacts().find(|(_, artifact)| match_code(artifact).unwrap_or_default()); + let artifact = out.artifacts().find(|(_, artifact)| { + artifact.get_deployed_bytecode_bytes().is_some_and(|b| *b == address_code) + }); if let Some((_, artifact)) = artifact { return fetch_and_print_storage(provider, address.clone(), artifact, true).await; } diff --git a/crates/common/src/contracts.rs b/crates/common/src/contracts.rs index be279817603bc..cc3b63f523607 100644 --- a/crates/common/src/contracts.rs +++ b/crates/common/src/contracts.rs @@ -125,13 +125,9 @@ pub fn flatten_contracts( let bytecode = if deployed_code { c.deployed_bytecode.clone().into_bytes() } else { - c.bytecode.clone().object.into_bytes() + c.bytecode.clone().into_bytes() }; - - if let Some(bytecode) = bytecode { - return Some((id.clone(), (c.abi.clone(), bytecode.to_vec()))) - } - None + bytecode.map(|bytecode| (id.clone(), (c.abi.clone(), bytecode.into()))) }) .collect(), ) diff --git a/crates/debugger/src/tui/draw.rs b/crates/debugger/src/tui/draw.rs index 84e56ffe1e120..fe51150714a7b 100644 --- a/crates/debugger/src/tui/draw.rs +++ b/crates/debugger/src/tui/draw.rs @@ -346,8 +346,8 @@ impl DebuggerContext<'_> { let mut source_map = bytecode.source_map()?.ok()?; let pc_ic_map = if is_create { create_map } else { rt_map }; - let ic = pc_ic_map.get(&pc)?; - let source_element = source_map.swap_remove(*ic); + let ic = pc_ic_map.get(pc)?; + let source_element = source_map.swap_remove(ic); (*file_id == source_element.index?).then_some((source_element, source_code)) }) else { diff --git a/crates/debugger/src/tui/mod.rs b/crates/debugger/src/tui/mod.rs index 1b25b8befee68..494c905e236ed 100644 --- a/crates/debugger/src/tui/mod.rs +++ b/crates/debugger/src/tui/mod.rs @@ -8,10 +8,7 @@ use crossterm::{ }; use eyre::Result; use foundry_common::{compile::ContractSources, evm::Breakpoints}; -use foundry_evm_core::{ - debug::DebugNodeFlat, - utils::{build_pc_ic_map, PCICMap}, -}; +use foundry_evm_core::{debug::DebugNodeFlat, utils::PcIcMap}; use ratatui::{ backend::{Backend, CrosstermBackend}, Terminal, @@ -50,7 +47,7 @@ pub struct Debugger { /// Source map of contract sources contracts_sources: ContractSources, /// A mapping of source -> (PC -> IC map for deploy code, PC -> IC map for runtime code) - pc_ic_maps: BTreeMap, + pc_ic_maps: BTreeMap, breakpoints: Breakpoints, } @@ -76,20 +73,8 @@ impl Debugger { Some(( contract_name.clone(), ( - build_pc_ic_map( - SpecId::LATEST, - contract.bytecode.object.as_bytes()?.as_ref(), - ), - build_pc_ic_map( - SpecId::LATEST, - contract - .deployed_bytecode - .bytecode - .as_ref()? - .object - .as_bytes()? - .as_ref(), - ), + PcIcMap::new(SpecId::LATEST, contract.bytecode.bytes()?), + PcIcMap::new(SpecId::LATEST, contract.deployed_bytecode.bytes()?), ), )) }) diff --git a/crates/evm/core/src/ic.rs b/crates/evm/core/src/ic.rs new file mode 100644 index 0000000000000..479a50b0a528d --- /dev/null +++ b/crates/evm/core/src/ic.rs @@ -0,0 +1,78 @@ +use revm::{ + interpreter::{opcode, opcode::spec_opcode_gas}, + primitives::{HashMap, SpecId}, +}; + +/// Maps from program counter to instruction counter. +/// +/// Inverse of [`IcPcMap`]. +pub struct PcIcMap { + pub inner: HashMap, +} + +impl PcIcMap { + /// Creates a new `PcIcMap` for the given code. + pub fn new(spec: SpecId, code: &[u8]) -> Self { + let opcode_infos = spec_opcode_gas(spec); + let mut map = HashMap::new(); + + let mut i = 0; + let mut cumulative_push_size = 0; + while i < code.len() { + let op = code[i]; + map.insert(i, i - cumulative_push_size); + if opcode_infos[op as usize].is_push() { + // Skip the push bytes. + // + // For more context on the math, see: https://github.com/bluealloy/revm/blob/007b8807b5ad7705d3cacce4d92b89d880a83301/crates/revm/src/interpreter/contract.rs#L114-L115 + i += (op - opcode::PUSH1 + 1) as usize; + cumulative_push_size += (op - opcode::PUSH1 + 1) as usize; + } + i += 1; + } + + Self { inner: map } + } + + /// Returns the instruction counter for the given program counter. + pub fn get(&self, pc: usize) -> Option { + self.inner.get(&pc).copied() + } +} + +/// Map from instruction counter to program counter. +/// +/// Inverse of [`PcIcMap`]. +pub struct IcPcMap { + pub inner: HashMap, +} + +impl IcPcMap { + /// Creates a new `IcPcMap` for the given code. + pub fn new(spec: SpecId, code: &[u8]) -> Self { + let opcode_infos = spec_opcode_gas(spec); + let mut map = HashMap::new(); + + let mut i = 0; + let mut cumulative_push_size = 0; + while i < code.len() { + let op = code[i]; + map.insert(i - cumulative_push_size, i); + if opcode_infos[op as usize].is_push() { + // Skip the push bytes. + // + // For more context on the math, see: https://github.com/bluealloy/revm/blob/007b8807b5ad7705d3cacce4d92b89d880a83301/crates/revm/src/interpreter/contract.rs#L114-L115 + i += (op - opcode::PUSH1 + 1) as usize; + cumulative_push_size += (op - opcode::PUSH1 + 1) as usize; + } + i += 1; + } + + Self { inner: map } + } + + /// Returns the program counter for the given instruction counter. + pub fn get(&self, ic: usize) -> Option { + self.inner.get(&ic).copied() + } +} diff --git a/crates/evm/core/src/lib.rs b/crates/evm/core/src/lib.rs index 9864ded06f0b0..caafa8e0f4d23 100644 --- a/crates/evm/core/src/lib.rs +++ b/crates/evm/core/src/lib.rs @@ -7,6 +7,8 @@ #[macro_use] extern crate tracing; +mod ic; + pub mod abi; pub mod backend; pub mod constants; diff --git a/crates/evm/core/src/utils.rs b/crates/evm/core/src/utils.rs index cbce9fc9984da..55ec9a15bdf7d 100644 --- a/crates/evm/core/src/utils.rs +++ b/crates/evm/core/src/utils.rs @@ -1,18 +1,19 @@ -use alloy_json_abi::{Function, JsonAbi as Abi}; -use alloy_primitives::{Address, FixedBytes, B256}; +use alloy_json_abi::{Function, JsonAbi}; +use alloy_primitives::FixedBytes; use alloy_rpc_types::{Block, Transaction}; use ethers::types::{ActionType, CallType, Chain, H256, U256}; use eyre::ContextCompat; use foundry_common::types::ToAlloy; -pub use foundry_compilers::utils::RuntimeOrHandle; -pub use revm::primitives::State as StateChangeset; - use revm::{ - interpreter::{opcode, opcode::spec_opcode_gas, CallScheme, CreateInputs, InstructionResult}, + interpreter::{CallScheme, InstructionResult}, primitives::{CreateScheme, Eval, Halt, SpecId, TransactTo}, }; use serde::{Deserialize, Serialize}; -use std::collections::BTreeMap; + +pub use foundry_compilers::utils::RuntimeOrHandle; +pub use revm::primitives::State as StateChangeset; + +pub use crate::ic::*; // TODO(onbjerg): Remove this and use `CallKind` from the tracer. #[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)] @@ -177,63 +178,11 @@ pub fn apply_chain_and_block_specific_env_changes(env: &mut revm::primitives::En } } -/// A map of program counters to instruction counters. -pub type PCICMap = BTreeMap; - -/// Builds a mapping from program counters to instruction counters. -pub fn build_pc_ic_map(spec: SpecId, code: &[u8]) -> PCICMap { - let opcode_infos = spec_opcode_gas(spec); - let mut pc_ic_map: PCICMap = BTreeMap::new(); - - let mut i = 0; - let mut cumulative_push_size = 0; - while i < code.len() { - let op = code[i]; - pc_ic_map.insert(i, i - cumulative_push_size); - if opcode_infos[op as usize].is_push() { - // Skip the push bytes. - // - // For more context on the math, see: https://github.com/bluealloy/revm/blob/007b8807b5ad7705d3cacce4d92b89d880a83301/crates/revm/src/interpreter/contract.rs#L114-L115 - i += (op - opcode::PUSH1 + 1) as usize; - cumulative_push_size += (op - opcode::PUSH1 + 1) as usize; - } - i += 1; - } - - pc_ic_map -} - -/// A map of instruction counters to program counters. -pub type ICPCMap = BTreeMap; - -/// Builds a mapping from instruction counters to program counters. -pub fn build_ic_pc_map(spec: SpecId, code: &[u8]) -> ICPCMap { - let opcode_infos = spec_opcode_gas(spec); - let mut ic_pc_map: ICPCMap = ICPCMap::new(); - - let mut i = 0; - let mut cumulative_push_size = 0; - while i < code.len() { - let op = code[i]; - ic_pc_map.insert(i - cumulative_push_size, i); - if opcode_infos[op as usize].is_push() { - // Skip the push bytes. - // - // For more context on the math, see: https://github.com/bluealloy/revm/blob/007b8807b5ad7705d3cacce4d92b89d880a83301/crates/revm/src/interpreter/contract.rs#L114-L115 - i += (op - opcode::PUSH1 + 1) as usize; - cumulative_push_size += (op - opcode::PUSH1 + 1) as usize; - } - i += 1; - } - - ic_pc_map -} - /// Given an ABI and selector, it tries to find the respective function. pub fn get_function( contract_name: &str, selector: &FixedBytes<4>, - abi: &Abi, + abi: &JsonAbi, ) -> eyre::Result { abi.functions() .find(|func| func.selector().as_slice() == selector.as_slice()) diff --git a/crates/evm/coverage/src/anchors.rs b/crates/evm/coverage/src/anchors.rs index 0342b38e9527e..237d941d1a3d0 100644 --- a/crates/evm/coverage/src/anchors.rs +++ b/crates/evm/coverage/src/anchors.rs @@ -1,7 +1,7 @@ use super::{CoverageItem, CoverageItemKind, ItemAnchor, SourceLocation}; use alloy_primitives::Bytes; use foundry_compilers::sourcemap::{SourceElement, SourceMap}; -use foundry_evm_core::utils::ICPCMap; +use foundry_evm_core::utils::IcPcMap; use revm::{ interpreter::opcode::{self, spec_opcode_gas}, primitives::SpecId, @@ -11,7 +11,7 @@ use revm::{ pub fn find_anchors( bytecode: &Bytes, source_map: &SourceMap, - ic_pc_map: &ICPCMap, + ic_pc_map: &IcPcMap, item_ids: &[usize], items: &[CoverageItem], ) -> Vec { @@ -49,7 +49,7 @@ pub fn find_anchors( /// Find an anchor representing the first opcode within the given source range. pub fn find_anchor_simple( source_map: &SourceMap, - ic_pc_map: &ICPCMap, + ic_pc_map: &IcPcMap, item_id: usize, loc: &SourceLocation, ) -> eyre::Result { @@ -62,7 +62,7 @@ pub fn find_anchor_simple( })?; Ok(ItemAnchor { - instruction: *ic_pc_map.get(&instruction).ok_or_else(|| { + instruction: ic_pc_map.get(instruction).ok_or_else(|| { eyre::eyre!("We found an anchor, but we cant translate it to a program counter") })?, item_id, diff --git a/crates/evm/traces/src/lib.rs b/crates/evm/traces/src/lib.rs index 3d5d09635807b..82186f3b7b643 100644 --- a/crates/evm/traces/src/lib.rs +++ b/crates/evm/traces/src/lib.rs @@ -6,6 +6,7 @@ #[macro_use] extern crate tracing; + use alloy_primitives::{Log, U256}; use foundry_common::contracts::{ContractsByAddress, ContractsByArtifact}; use foundry_evm_core::constants::CHEATCODE_ADDRESS; diff --git a/crates/forge/bin/cmd/coverage.rs b/crates/forge/bin/cmd/coverage.rs index db550b4b07d2a..771bd8bb09066 100644 --- a/crates/forge/bin/cmd/coverage.rs +++ b/crates/forge/bin/cmd/coverage.rs @@ -11,7 +11,7 @@ use forge::{ opts::EvmOpts, result::SuiteResult, revm::primitives::SpecId, - utils::{build_ic_pc_map, ICPCMap}, + utils::IcPcMap, MultiContractRunnerBuilder, TestOptions, }; use foundry_cli::{ @@ -233,15 +233,15 @@ impl CoverageArgs { // Since our coverage inspector collects hit data using program counters, the anchors also // need to be based on program counters. // TODO: Index by contract ID - let ic_pc_maps: HashMap = bytecodes + let ic_pc_maps: HashMap = bytecodes .iter() .map(|(id, bytecodes)| { // TODO: Creation bytecode as well ( id.clone(), ( - build_ic_pc_map(SpecId::LATEST, bytecodes.0.as_ref()), - build_ic_pc_map(SpecId::LATEST, bytecodes.1.as_ref()), + IcPcMap::new(SpecId::LATEST, bytecodes.0.as_ref()), + IcPcMap::new(SpecId::LATEST, bytecodes.1.as_ref()), ), ) }) diff --git a/crates/forge/bin/cmd/script/executor.rs b/crates/forge/bin/cmd/script/executor.rs index cee58db6852e6..3d1eb7a887bec 100644 --- a/crates/forge/bin/cmd/script/executor.rs +++ b/crates/forge/bin/cmd/script/executor.rs @@ -44,7 +44,6 @@ impl ScriptArgs { let abi = abi.ok_or_else(|| eyre::eyre!("no ABI found for contract"))?; let bytecode = bytecode .ok_or_else(|| eyre::eyre!("no bytecode found for contract"))? - .object .into_bytes() .ok_or_else(|| { eyre::eyre!("expected fully linked bytecode, found unlinked bytecode")