revm_bytecode/
iter.rs

1use crate::{opcode, Bytecode, OpCode};
2
3/// Iterator over opcodes in a bytecode, skipping immediates.
4///
5/// This allows you to iterate through the actual opcodes in the bytecode,
6/// without dealing with the immediate values that follow instructions.
7#[derive(Debug, Clone)]
8pub struct BytecodeIterator<'a> {
9    /// Iterator over the bytecode bytes.
10    bytes: core::slice::Iter<'a, u8>,
11    /// Start pointer of the bytecode. Only used to calculate [`position`](Self::position).
12    start: *const u8,
13}
14
15impl<'a> BytecodeIterator<'a> {
16    /// Creates a new iterator from a bytecode reference.
17    #[inline]
18    pub fn new(bytecode: &'a Bytecode) -> Self {
19        let bytes = if bytecode.is_legacy() {
20            bytecode.original_byte_slice()
21        } else {
22            &[]
23        };
24        Self {
25            bytes: bytes.iter(),
26            start: bytes.as_ptr(),
27        }
28    }
29
30    /// Skips to the next opcode, taking into account PUSH instructions.
31    pub fn skip_to_next_opcode(&mut self) {
32        self.next();
33    }
34
35    /// Returns the remaining bytes in the bytecode as a slice.
36    #[inline]
37    pub fn as_slice(&self) -> &[u8] {
38        self.bytes.as_slice()
39    }
40
41    /// Returns the current position in the bytecode.
42    #[inline]
43    pub fn position(&self) -> usize {
44        // SAFETY: `start` always points to the start of the bytecode.
45        unsafe {
46            self.bytes
47                .as_slice()
48                .as_ptr()
49                .offset_from_unsigned(self.start)
50        }
51    }
52
53    #[inline]
54    fn skip_immediate(&mut self, opcode: u8) {
55        // Get base immediate size from opcode info
56        let immediate_size = opcode::OPCODE_INFO[opcode as usize]
57            .map(|info| info.immediate_size() as usize)
58            .unwrap_or_default();
59
60        // Advance the iterator by the immediate size
61        if immediate_size > 0 {
62            let remaining = self.bytes.as_slice();
63            self.bytes = remaining[immediate_size.min(remaining.len())..].iter();
64        }
65    }
66
67    /// Returns the current opcode without advancing the iterator.
68    #[inline]
69    pub fn peek(&self) -> Option<u8> {
70        self.bytes.as_slice().first().copied()
71    }
72
73    /// Returns the current opcode wrapped in OpCode without advancing the iterator.
74    #[inline]
75    pub fn peek_opcode(&self) -> Option<OpCode> {
76        self.peek().and_then(OpCode::new)
77    }
78}
79
80impl Iterator for BytecodeIterator<'_> {
81    type Item = u8;
82
83    #[inline]
84    fn next(&mut self) -> Option<Self::Item> {
85        self.bytes
86            .next()
87            .copied()
88            .inspect(|&current| self.skip_immediate(current))
89    }
90
91    #[inline]
92    fn size_hint(&self) -> (usize, Option<usize>) {
93        // Lower bound is 0 if empty, 1 if not empty as it depends on the bytes.
94        let byte_len = self.bytes.len();
95        (byte_len.min(1), Some(byte_len))
96    }
97}
98
99impl core::iter::FusedIterator for BytecodeIterator<'_> {}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104    use primitives::Bytes;
105
106    #[test]
107    fn test_simple_bytecode_iteration() {
108        // Create a simple bytecode: PUSH1 0x01 PUSH1 0x02 ADD STOP
109        let bytecode = Bytecode::new_legacy(Bytes::from_static(&[
110            opcode::PUSH1,
111            0x01,
112            opcode::PUSH1,
113            0x02,
114            opcode::ADD,
115            opcode::STOP,
116        ]));
117        let opcodes: Vec<u8> = bytecode.iter_opcodes().collect();
118        assert_eq!(
119            opcodes,
120            vec![opcode::PUSH1, opcode::PUSH1, opcode::ADD, opcode::STOP]
121        );
122    }
123
124    #[test]
125    fn test_bytecode_with_various_push_sizes() {
126        let bytecode = Bytecode::new_legacy(Bytes::from_static(&[
127            opcode::PUSH1,
128            0x01,
129            opcode::PUSH2,
130            0x02,
131            0x03,
132            opcode::PUSH3,
133            0x04,
134            0x05,
135            0x06,
136            opcode::STOP,
137        ]));
138
139        let opcodes: Vec<u8> = bytecode.iter_opcodes().collect();
140
141        // We should only see the opcodes, not the immediates
142        assert_eq!(
143            opcodes,
144            vec![opcode::PUSH1, opcode::PUSH2, opcode::PUSH3, opcode::STOP]
145        );
146    }
147
148    #[test]
149    fn test_bytecode_skips_immediates() {
150        let bytecode = Bytecode::new_legacy(Bytes::from_static(&[
151            opcode::PUSH1,
152            0x01,
153            opcode::PUSH2,
154            0x02,
155            0x03,
156            opcode::ADD,
157            opcode::PUSH3,
158            0x04,
159            0x05,
160            0x06,
161            opcode::PUSH32,
162            0x10,
163            0x11,
164            0x12,
165            0x13,
166            0x14,
167            0x15,
168            0x16,
169            0x17,
170            0x18,
171            0x19,
172            0x1a,
173            0x1b,
174            0x1c,
175            0x1d,
176            0x1e,
177            0x1f,
178            0x20,
179            0x21,
180            0x22,
181            0x23,
182            0x24,
183            0x25,
184            0x26,
185            0x27,
186            0x28,
187            0x29,
188            0x2a,
189            0x2b,
190            0x2c,
191            0x2d,
192            0x2e,
193            0x2f,
194            opcode::MUL,
195            opcode::STOP,
196        ]));
197
198        let opcodes: Vec<u8> = bytecode.iter_opcodes().collect();
199        assert_eq!(
200            opcodes,
201            vec![
202                opcode::PUSH1,
203                opcode::PUSH2,
204                opcode::ADD,
205                opcode::PUSH3,
206                opcode::PUSH32,
207                opcode::MUL,
208                opcode::STOP,
209            ]
210        );
211    }
212
213    #[test]
214    fn test_position_tracking() {
215        let bytecode = Bytecode::new_legacy(Bytes::from_static(&[
216            opcode::PUSH1,
217            0x01,
218            opcode::PUSH1,
219            0x02,
220            opcode::ADD,
221            opcode::STOP,
222        ]));
223
224        let mut iter = bytecode.iter_opcodes();
225
226        assert_eq!(iter.position(), 0);
227        assert_eq!(iter.next(), Some(opcode::PUSH1));
228        assert_eq!(iter.position(), 2);
229
230        assert_eq!(iter.next(), Some(opcode::PUSH1));
231        assert_eq!(iter.position(), 4);
232
233        assert_eq!(iter.next(), Some(opcode::ADD));
234        assert_eq!(iter.position(), 5);
235
236        assert_eq!(iter.next(), Some(opcode::STOP));
237        assert_eq!(iter.position(), 6);
238
239        assert_eq!(iter.next(), None);
240        assert_eq!(iter.position(), 6);
241    }
242
243    #[test]
244    fn test_empty_bytecode() {
245        let bytecode = Bytecode::new_legacy(Bytes::from_static(&[opcode::STOP]));
246        let opcodes: Vec<u8> = bytecode.iter_opcodes().collect();
247        assert_eq!(opcodes, vec![opcode::STOP]);
248    }
249
250    #[test]
251    fn test_truncated_push_does_not_iterate_padding() {
252        let bytecode = Bytecode::new_legacy(Bytes::from_static(&[opcode::PUSH1]));
253        let opcodes: Vec<u8> = bytecode.iter_opcodes().collect();
254        assert_eq!(opcodes, vec![opcode::PUSH1]);
255    }
256
257    #[test]
258    fn test_position_after_truncated_push() {
259        let bytecode = Bytecode::new_legacy(Bytes::from_static(&[opcode::PUSH2, 0x01]));
260        let mut iter = bytecode.iter_opcodes();
261
262        assert_eq!(iter.next(), Some(opcode::PUSH2));
263        assert_eq!(iter.position(), 2);
264        assert_eq!(iter.next(), None);
265    }
266}
revm_bytecode/iter.rs

revm_bytecode/
iter.rs
