1use alloy_primitives::{Bytes, U256};
30use revm::bytecode::opcode::OpCode;
31
32#[derive(Debug, Clone, PartialEq, Eq)]
34pub struct DisassemblyInstruction {
35 pub pc: usize,
37 pub opcode: OpCode,
39 pub push_data: Vec<u8>,
42}
43
44impl DisassemblyInstruction {
45 pub fn new(pc: usize, opcode: OpCode) -> Self {
47 Self { pc, opcode, push_data: Vec::new() }
48 }
49
50 pub fn with_push_data(pc: usize, opcode: OpCode, push_data: Vec<u8>) -> Self {
52 Self { pc, opcode, push_data }
53 }
54
55 pub fn is_push(&self) -> bool {
57 let opcode_byte = self.opcode.get();
58 (0x60..=0x7F).contains(&opcode_byte)
59 }
60
61 pub fn push_size(&self) -> usize {
64 if self.is_push() {
65 (self.opcode.get() - 0x60 + 1) as usize
66 } else {
67 0
68 }
69 }
70
71 pub fn instruction_size(&self) -> usize {
73 1 + self.push_size()
74 }
75}
76
77#[derive(Debug, Clone)]
79pub struct DisassemblyResult {
80 pub bytecode: Bytes,
82 pub instructions: Vec<DisassemblyInstruction>,
84}
85
86impl DisassemblyResult {
87 pub fn new(bytecode: Bytes, instructions: Vec<DisassemblyInstruction>) -> Self {
89 Self { bytecode, instructions }
90 }
91
92 pub fn instruction_count(&self) -> usize {
94 self.instructions.len()
95 }
96
97 pub fn get_instruction_at_pc(&self, pc: usize) -> Option<&DisassemblyInstruction> {
99 self.instructions.iter().find(|inst| inst.pc == pc)
100 }
101
102 pub fn get_push_instructions(&self) -> Vec<&DisassemblyInstruction> {
104 self.instructions.iter().filter(|inst| inst.is_push()).collect()
105 }
106
107 pub fn find_instruction_containing_pc(&self, pc: usize) -> Option<&DisassemblyInstruction> {
110 self.instructions.iter().find(|inst| {
111 let start = inst.pc;
112 let end = start + inst.instruction_size();
113 pc >= start && pc < end
114 })
115 }
116}
117
118pub fn disassemble(bytecode: &Bytes) -> DisassemblyResult {
142 let mut instructions = Vec::new();
143 let mut pc = 0;
144
145 while pc < bytecode.len() {
146 let opcode_byte = bytecode[pc];
147
148 let opcode = unsafe { OpCode::new_unchecked(opcode_byte) };
150
151 if (0x60..=0x7F).contains(&opcode_byte) {
153 let push_size = (opcode_byte - 0x60 + 1) as usize;
155 let data_start = pc + 1;
156 let data_end = data_start + push_size;
157
158 let mut push_data = Vec::new();
160 for i in data_start..data_end {
161 if i < bytecode.len() {
162 push_data.push(bytecode[i]);
163 } else {
164 push_data.push(0); }
166 }
167
168 instructions.push(DisassemblyInstruction::with_push_data(pc, opcode, push_data));
169 pc = data_end;
170 } else {
171 instructions.push(DisassemblyInstruction::new(pc, opcode));
173 pc += 1;
174 }
175 }
176
177 DisassemblyResult::new(bytecode.clone(), instructions)
178}
179
180pub fn extract_push_value(instruction: &DisassemblyInstruction) -> Option<U256> {
204 if !instruction.is_push() || instruction.push_data.is_empty() {
205 return None;
206 }
207
208 let mut value = U256::ZERO;
209 for &byte in &instruction.push_data {
210 value = value.wrapping_shl(8).wrapping_add(U256::from(byte));
211 }
212 Some(value)
213}
214
215pub fn format_instruction(instruction: &DisassemblyInstruction, show_pc: bool) -> String {
242 let pc_part = if show_pc { format!("{:04x}: ", instruction.pc) } else { String::new() };
243
244 let opcode_name = if instruction.opcode.is_valid() {
245 instruction.opcode.as_str().to_string()
246 } else {
247 format!("'{:x}'(Unknown Opcode)", instruction.opcode.get())
248 };
249
250 if instruction.is_push() && !instruction.push_data.is_empty() {
251 let hex_data = instruction.push_data.iter().map(|b| format!("{b:02x}")).collect::<String>();
252 format!("{pc_part}{opcode_name} 0x{hex_data}")
253 } else {
254 format!("{pc_part}{opcode_name}")
255 }
256}
257
258#[cfg(test)]
259mod tests {
260 use super::*;
261 use alloy_primitives::Bytes;
262
263 #[test]
264 fn test_disassemble_simple() {
265 let bytecode = Bytes::from(vec![0x80, 0x81, 0x82]); let result = disassemble(&bytecode);
267
268 assert_eq!(result.instructions.len(), 3);
269 assert_eq!(result.instructions[0].pc, 0);
270 assert_eq!(result.instructions[1].pc, 1);
271 assert_eq!(result.instructions[2].pc, 2);
272
273 for inst in &result.instructions {
274 assert!(!inst.is_push());
275 assert!(inst.push_data.is_empty());
276 }
277 }
278
279 #[test]
280 fn test_disassemble_push_instructions() {
281 let bytecode = Bytes::from(vec![
282 0x60, 0x42, 0x61, 0x12, 0x34, 0x80, ]);
286 let result = disassemble(&bytecode);
287
288 assert_eq!(result.instructions.len(), 3);
289
290 assert_eq!(result.instructions[0].pc, 0);
292 assert!(result.instructions[0].is_push());
293 assert_eq!(result.instructions[0].push_data, vec![0x42]);
294 assert_eq!(result.instructions[0].instruction_size(), 2);
295
296 assert_eq!(result.instructions[1].pc, 2);
298 assert!(result.instructions[1].is_push());
299 assert_eq!(result.instructions[1].push_data, vec![0x12, 0x34]);
300 assert_eq!(result.instructions[1].instruction_size(), 3);
301
302 assert_eq!(result.instructions[2].pc, 5);
304 assert!(!result.instructions[2].is_push());
305 assert!(result.instructions[2].push_data.is_empty());
306 assert_eq!(result.instructions[2].instruction_size(), 1);
307 }
308
309 #[test]
310 fn test_extract_push_value() {
311 let push1 = DisassemblyInstruction::with_push_data(
312 0,
313 unsafe { OpCode::new_unchecked(0x60) },
314 vec![0x42],
315 );
316 assert_eq!(extract_push_value(&push1), Some(U256::from(0x42)));
317
318 let push2 = DisassemblyInstruction::with_push_data(
319 0,
320 unsafe { OpCode::new_unchecked(0x61) },
321 vec![0x12, 0x34],
322 );
323 assert_eq!(extract_push_value(&push2), Some(U256::from(0x1234)));
324
325 let push4 = DisassemblyInstruction::with_push_data(
326 0,
327 unsafe { OpCode::new_unchecked(0x63) },
328 vec![0x12, 0x34, 0x56, 0x78],
329 );
330 assert_eq!(extract_push_value(&push4), Some(U256::from(0x12345678)));
331 }
332
333 #[test]
334 fn test_find_instruction_containing_pc() {
335 let bytecode = Bytes::from(vec![
336 0x60, 0x42, 0x61, 0x12, 0x34, 0x80, ]);
340 let result = disassemble(&bytecode);
341
342 assert_eq!(result.find_instruction_containing_pc(0).unwrap().pc, 0);
344 assert_eq!(result.find_instruction_containing_pc(1).unwrap().pc, 0);
345
346 assert_eq!(result.find_instruction_containing_pc(2).unwrap().pc, 2);
348 assert_eq!(result.find_instruction_containing_pc(3).unwrap().pc, 2);
349 assert_eq!(result.find_instruction_containing_pc(4).unwrap().pc, 2);
350
351 assert_eq!(result.find_instruction_containing_pc(5).unwrap().pc, 5);
353
354 assert!(result.find_instruction_containing_pc(6).is_none());
356 }
357
358 #[test]
359 fn test_truncated_push_instruction() {
360 let bytecode = Bytes::from(vec![0x61, 0x12]); let result = disassemble(&bytecode);
362
363 assert_eq!(result.instructions.len(), 1);
364 assert!(result.instructions[0].is_push());
365 assert_eq!(result.instructions[0].push_data, vec![0x12, 0x00]); }
367
368 #[test]
369 fn test_format_instruction() {
370 let push_inst = DisassemblyInstruction::with_push_data(
371 10,
372 unsafe { OpCode::new_unchecked(0x61) },
373 vec![0x12, 0x34],
374 );
375
376 let with_pc = format_instruction(&push_inst, true);
377 assert_eq!(with_pc, "000a: PUSH2 0x1234");
378
379 let without_pc = format_instruction(&push_inst, false);
380 assert_eq!(without_pc, "PUSH2 0x1234");
381
382 let regular_inst = DisassemblyInstruction::new(5, unsafe { OpCode::new_unchecked(0x80) });
383
384 let regular_formatted = format_instruction(®ular_inst, true);
385 assert_eq!(regular_formatted, "0005: DUP1");
386 }
387}