chore: new-type IC/PC maps (foundry-rs#6742)

* chore: move pc_ic map to a new module

* chore: create new types for pc-ic maps

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<bool> {
-                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;
             }

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(),
     )

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 {

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<String, (PCICMap, PCICMap)>,
+    pc_ic_maps: BTreeMap<String, (PcIcMap, PcIcMap)>,
     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()?),
                         ),
                     ))
                 })

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<usize, usize>,
+}
+
+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<usize> {
+        self.inner.get(&pc).copied()
+    }
+}
+
+/// Map from instruction counter to program counter.
+///
+/// Inverse of [`PcIcMap`].
+pub struct IcPcMap {
+    pub inner: HashMap<usize, usize>,
+}
+
+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<usize> {
+        self.inner.get(&ic).copied()
+    }
+}

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;

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<usize, usize>;
-
-/// 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<usize, usize>;
-
-/// 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<Function> {
     abi.functions()
         .find(|func| func.selector().as_slice() == selector.as_slice())

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<ItemAnchor> {
@@ -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<ItemAnchor> {
@@ -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,

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;

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<ContractId, (ICPCMap, ICPCMap)> = bytecodes
+        let ic_pc_maps: HashMap<ContractId, (IcPcMap, IcPcMap)> = 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()),
                     ),
                 )
             })

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")