foundry_evm_core/
ic.rs

1use alloy_primitives::map::rustc_hash::FxHashMap;
2use eyre::Result;
3use revm::bytecode::opcode::{OpCode, PUSH0, PUSH1, PUSH32};
4use serde::Serialize;
5
6/// Maps from program counter to instruction counter.
7///
8/// Inverse of [`IcPcMap`].
9#[derive(Debug, Clone, Serialize)]
10#[serde(transparent)]
11pub struct PcIcMap {
12    pub inner: FxHashMap<u32, u32>,
13}
14
15impl PcIcMap {
16    /// Creates a new `PcIcMap` for the given code.
17    pub fn new(code: &[u8]) -> Self {
18        Self { inner: make_map::<true>(code) }
19    }
20
21    /// Returns the length of the map.
22    pub fn len(&self) -> usize {
23        self.inner.len()
24    }
25
26    /// Returns `true` if the map is empty.
27    pub fn is_empty(&self) -> bool {
28        self.inner.is_empty()
29    }
30
31    /// Returns the instruction counter for the given program counter.
32    pub fn get(&self, pc: u32) -> Option<u32> {
33        self.inner.get(&pc).copied()
34    }
35}
36
37/// Map from instruction counter to program counter.
38///
39/// Inverse of [`PcIcMap`].
40pub struct IcPcMap {
41    pub inner: FxHashMap<u32, u32>,
42}
43
44impl IcPcMap {
45    /// Creates a new `IcPcMap` for the given code.
46    pub fn new(code: &[u8]) -> Self {
47        Self { inner: make_map::<false>(code) }
48    }
49
50    /// Returns the length of the map.
51    pub fn len(&self) -> usize {
52        self.inner.len()
53    }
54
55    /// Returns `true` if the map is empty.
56    pub fn is_empty(&self) -> bool {
57        self.inner.is_empty()
58    }
59
60    /// Returns the program counter for the given instruction counter.
61    pub fn get(&self, ic: u32) -> Option<u32> {
62        self.inner.get(&ic).copied()
63    }
64}
65
66fn make_map<const PC_FIRST: bool>(code: &[u8]) -> FxHashMap<u32, u32> {
67    assert!(code.len() <= u32::MAX as usize, "bytecode is too big");
68
69    let mut map = FxHashMap::with_capacity_and_hasher(code.len(), Default::default());
70
71    let mut pc = 0usize;
72    let mut cumulative_push_size = 0usize;
73    while pc < code.len() {
74        let ic = pc - cumulative_push_size;
75        if PC_FIRST {
76            map.insert(pc as u32, ic as u32);
77        } else {
78            map.insert(ic as u32, pc as u32);
79        }
80
81        if (PUSH1..=PUSH32).contains(&code[pc]) {
82            // Skip the push bytes.
83            let push_size = (code[pc] - PUSH0) as usize;
84            pc += push_size;
85            cumulative_push_size += push_size;
86        }
87
88        pc += 1;
89    }
90
91    map.shrink_to_fit();
92
93    map
94}
95
96/// Represents a single instruction consisting of the opcode and its immediate data.
97pub struct Instruction {
98    /// OpCode, if it could be decoded.
99    pub op: Option<OpCode>,
100    /// Immediate data following the opcode.
101    pub immediate: Box<[u8]>,
102    /// Program counter of the opcode.
103    pub pc: u32,
104}
105
106/// Decodes raw opcode bytes into [`Instruction`]s.
107pub fn decode_instructions(code: &[u8]) -> Result<Vec<Instruction>> {
108    assert!(code.len() <= u32::MAX as usize, "bytecode is too big");
109
110    let mut pc = 0usize;
111    let mut steps = Vec::new();
112
113    while pc < code.len() {
114        let op = OpCode::new(code[pc]);
115        let next_pc = pc + 1;
116        let immediate_size = op.map(|op| op.info().immediate_size()).unwrap_or(0) as usize;
117        let is_normal_push = op.map(|op| op.is_push()).unwrap_or(false);
118
119        if !is_normal_push && next_pc + immediate_size > code.len() {
120            eyre::bail!("incomplete sequence of bytecode");
121        }
122
123        // Ensure immediate is padded if needed.
124        let immediate_end = (next_pc + immediate_size).min(code.len());
125        let mut immediate = vec![0u8; immediate_size];
126        let immediate_part = &code[next_pc..immediate_end];
127        immediate[..immediate_part.len()].copy_from_slice(immediate_part);
128
129        steps.push(Instruction { op, pc: pc as u32, immediate: immediate.into_boxed_slice() });
130
131        pc = next_pc + immediate_size;
132    }
133
134    Ok(steps)
135}
136
137#[cfg(test)]
138pub mod tests {
139    use super::*;
140
141    #[test]
142    fn decode_push2_and_stop() -> Result<()> {
143        // 0x61 0xAA 0xBB = PUSH2 0xAABB
144        // 0x00           = STOP
145        let code = vec![0x61, 0xAA, 0xBB, 0x00];
146        let insns = decode_instructions(&code)?;
147
148        // PUSH2 then STOP
149        assert_eq!(insns.len(), 2);
150
151        // PUSH2 at pc = 0
152        let i0 = &insns[0];
153        assert_eq!(i0.pc, 0);
154        assert_eq!(i0.op, Some(OpCode::PUSH2));
155        assert_eq!(i0.immediate.as_ref(), &[0xAA, 0xBB]);
156
157        // STOP at pc = 3
158        let i1 = &insns[1];
159        assert_eq!(i1.pc, 3);
160        assert_eq!(i1.op, Some(OpCode::STOP));
161        assert!(i1.immediate.is_empty());
162
163        Ok(())
164    }
165
166    #[test]
167    fn decode_arithmetic_ops() -> Result<()> {
168        // 0x01 = ADD, 0x02 = MUL, 0x03 = SUB, 0x04 = DIV
169        let code = vec![0x01, 0x02, 0x03, 0x04];
170        let insns = decode_instructions(&code)?;
171
172        assert_eq!(insns.len(), 4);
173
174        let expected = [(0, OpCode::ADD), (1, OpCode::MUL), (2, OpCode::SUB), (3, OpCode::DIV)];
175        for ((pc, want_op), insn) in expected.iter().zip(insns.iter()) {
176            assert_eq!(insn.pc, *pc);
177            assert_eq!(insn.op, Some(*want_op));
178            assert!(insn.immediate.is_empty());
179        }
180
181        Ok(())
182    }
183}