Skip to content

Commit

Permalink
chore: create new types for pc-ic maps
Browse files Browse the repository at this point in the history
  • Loading branch information
DaniPopes committed Jan 9, 2024
1 parent 44a4337 commit f2e371e
Show file tree
Hide file tree
Showing 8 changed files with 88 additions and 89 deletions.
14 changes: 6 additions & 8 deletions crates/cast/bin/cmd/storage.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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;
}
Expand Down
8 changes: 2 additions & 6 deletions crates/common/src/contracts.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
)
Expand Down
4 changes: 2 additions & 2 deletions crates/debugger/src/tui/draw.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
23 changes: 4 additions & 19 deletions crates/debugger/src/tui/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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,
}

Expand All @@ -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()?),
),
))
})
Expand Down
111 changes: 66 additions & 45 deletions crates/evm/core/src/ic.rs
Original file line number Diff line number Diff line change
@@ -1,57 +1,78 @@
use revm::{
interpreter::{opcode, opcode::spec_opcode_gas},
primitives::SpecId,
primitives::{HashMap, SpecId},
};
use std::collections::BTreeMap;

/// 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;

/// 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;
}
i += 1;

Self { inner: map }
}

pc_ic_map
/// Returns the instruction counter for the given program counter.
pub fn get(&self, pc: usize) -> Option<usize> {
self.inner.get(&pc).copied()
}
}

/// 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;
/// 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;
}
i += 1;

Self { inner: map }
}

ic_pc_map
/// Returns the program counter for the given instruction counter.
pub fn get(&self, ic: usize) -> Option<usize> {
self.inner.get(&ic).copied()
}
}
8 changes: 4 additions & 4 deletions crates/evm/coverage/src/anchors.rs
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -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> {
Expand Down Expand Up @@ -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> {
Expand All @@ -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,
Expand Down
8 changes: 4 additions & 4 deletions crates/forge/bin/cmd/coverage.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::{
Expand Down Expand Up @@ -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()),
),
)
})
Expand Down
1 change: 0 additions & 1 deletion crates/forge/bin/cmd/script/executor.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down

0 comments on commit f2e371e

Please sign in to comment.