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    use std::{vec, vec::Vec};
106
107    #[test]
108    fn test_simple_bytecode_iteration() {
109        // Create a simple bytecode: PUSH1 0x01 PUSH1 0x02 ADD STOP
110        let bytecode = Bytecode::new_legacy(Bytes::from_static(&[
111            opcode::PUSH1,
112            0x01,
113            opcode::PUSH1,
114            0x02,
115            opcode::ADD,
116            opcode::STOP,
117        ]));
118        let opcodes: Vec<u8> = bytecode.iter_opcodes().collect();
119        assert_eq!(
120            opcodes,
121            vec![opcode::PUSH1, opcode::PUSH1, opcode::ADD, opcode::STOP]
122        );
123    }
124
125    #[test]
126    fn test_bytecode_with_various_push_sizes() {
127        let bytecode = Bytecode::new_legacy(Bytes::from_static(&[
128            opcode::PUSH1,
129            0x01,
130            opcode::PUSH2,
131            0x02,
132            0x03,
133            opcode::PUSH3,
134            0x04,
135            0x05,
136            0x06,
137            opcode::STOP,
138        ]));
139
140        let opcodes: Vec<u8> = bytecode.iter_opcodes().collect();
141
142        // We should only see the opcodes, not the immediates
143        assert_eq!(
144            opcodes,
145            vec![opcode::PUSH1, opcode::PUSH2, opcode::PUSH3, opcode::STOP]
146        );
147    }
148
149    #[test]
150    fn test_bytecode_skips_immediates() {
151        let bytecode = Bytecode::new_legacy(Bytes::from_static(&[
152            opcode::PUSH1,
153            0x01,
154            opcode::PUSH2,
155            0x02,
156            0x03,
157            opcode::ADD,
158            opcode::PUSH3,
159            0x04,
160            0x05,
161            0x06,
162            opcode::PUSH32,
163            0x10,
164            0x11,
165            0x12,
166            0x13,
167            0x14,
168            0x15,
169            0x16,
170            0x17,
171            0x18,
172            0x19,
173            0x1a,
174            0x1b,
175            0x1c,
176            0x1d,
177            0x1e,
178            0x1f,
179            0x20,
180            0x21,
181            0x22,
182            0x23,
183            0x24,
184            0x25,
185            0x26,
186            0x27,
187            0x28,
188            0x29,
189            0x2a,
190            0x2b,
191            0x2c,
192            0x2d,
193            0x2e,
194            0x2f,
195            opcode::MUL,
196            opcode::STOP,
197        ]));
198
199        let opcodes: Vec<u8> = bytecode.iter_opcodes().collect();
200        assert_eq!(
201            opcodes,
202            vec![
203                opcode::PUSH1,
204                opcode::PUSH2,
205                opcode::ADD,
206                opcode::PUSH3,
207                opcode::PUSH32,
208                opcode::MUL,
209                opcode::STOP,
210            ]
211        );
212    }
213
214    #[test]
215    fn test_position_tracking() {
216        let bytecode = Bytecode::new_legacy(Bytes::from_static(&[
217            opcode::PUSH1,
218            0x01,
219            opcode::PUSH1,
220            0x02,
221            opcode::ADD,
222            opcode::STOP,
223        ]));
224
225        let mut iter = bytecode.iter_opcodes();
226
227        assert_eq!(iter.position(), 0);
228        assert_eq!(iter.next(), Some(opcode::PUSH1));
229        assert_eq!(iter.position(), 2);
230
231        assert_eq!(iter.next(), Some(opcode::PUSH1));
232        assert_eq!(iter.position(), 4);
233
234        assert_eq!(iter.next(), Some(opcode::ADD));
235        assert_eq!(iter.position(), 5);
236
237        assert_eq!(iter.next(), Some(opcode::STOP));
238        assert_eq!(iter.position(), 6);
239
240        assert_eq!(iter.next(), None);
241        assert_eq!(iter.position(), 6);
242    }
243
244    #[test]
245    fn test_empty_bytecode() {
246        let bytecode = Bytecode::new_legacy(Bytes::from_static(&[opcode::STOP]));
247        let opcodes: Vec<u8> = bytecode.iter_opcodes().collect();
248        assert_eq!(opcodes, vec![opcode::STOP]);
249    }
250
251    #[test]
252    fn test_truncated_push_does_not_iterate_padding() {
253        let bytecode = Bytecode::new_legacy(Bytes::from_static(&[opcode::PUSH1]));
254        let opcodes: Vec<u8> = bytecode.iter_opcodes().collect();
255        assert_eq!(opcodes, vec![opcode::PUSH1]);
256    }
257
258    #[test]
259    fn test_position_after_truncated_push() {
260        let bytecode = Bytecode::new_legacy(Bytes::from_static(&[opcode::PUSH2, 0x01]));
261        let mut iter = bytecode.iter_opcodes();
262
263        assert_eq!(iter.next(), Some(opcode::PUSH2));
264        assert_eq!(iter.position(), 2);
265        assert_eq!(iter.next(), None);
266    }
267}
revm_bytecode/iter.rs

revm_bytecode/
iter.rs
