From 09f0670a7730f706bf802e9e9d5e36b7e96d17c6 Mon Sep 17 00:00:00 2001
From: DaniPopes <57450786+DaniPopes@users.noreply.github.com>
Date: Wed, 13 Mar 2024 18:37:42 +0100
Subject: [PATCH] chore: factor out common ic-pc code

---
 crates/evm/core/src/ic.rs | 67 +++++++++++++++++----------------------
 1 file changed, 29 insertions(+), 38 deletions(-)

diff --git a/crates/evm/core/src/ic.rs b/crates/evm/core/src/ic.rs
index 7bb2179c5e9a..9f679fc36e50 100644
--- a/crates/evm/core/src/ic.rs
+++ b/crates/evm/core/src/ic.rs
@@ -14,25 +14,7 @@ pub struct PcIcMap {
 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 = FxHashMap::default();
-
-        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;
-        }
-
-        Self { inner: map }
+        Self { inner: make_map::<true>(spec, code) }
     }
 
     /// Returns the instruction counter for the given program counter.
@@ -51,25 +33,7 @@ pub struct IcPcMap {
 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 = FxHashMap::default();
-
-        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;
-        }
-
-        Self { inner: map }
+        Self { inner: make_map::<false>(spec, code) }
     }
 
     /// Returns the program counter for the given instruction counter.
@@ -77,3 +41,30 @@ impl IcPcMap {
         self.inner.get(&ic).copied()
     }
 }
+
+fn make_map<const PC_FIRST: bool>(spec: SpecId, code: &[u8]) -> FxHashMap<usize, usize> {
+    let opcode_infos = spec_opcode_gas(spec);
+    let mut map = FxHashMap::default();
+
+    let mut pc = 0;
+    let mut cumulative_push_size = 0;
+    while pc < code.len() {
+        let ic = pc - cumulative_push_size;
+        if PC_FIRST {
+            map.insert(pc, ic);
+        } else {
+            map.insert(ic, pc);
+        }
+
+        let op = code[pc];
+        if opcode_infos[op as usize].is_push() {
+            // Skip the push bytes.
+            let push_size = (op - opcode::PUSH0) as usize;
+            pc += push_size;
+            cumulative_push_size += push_size;
+        }
+
+        pc += 1;
+    }
+    map
+}