refactor: encapsulate OpcodeId arithmetics (privacy-scaling-explorati…

…ons#824)

* refactor: encapsulate OpcodeId arithmetics

Closes privacy-scaling-explorations#822

bus-mapping/src/evm/opcodes/logs.rs

@@ -3,7 +3,6 @@ use crate::circuit_input_builder::{CircuitInputStateRef, ExecState, ExecStep};
 use crate::circuit_input_builder::{CopyDataType, CopyEvent, NumberOrHash};
 use crate::operation::{CallContextField, TxLogField};
 use crate::Error;
-use eth_types::evm_types::OpcodeId;
 use eth_types::Word;
 use eth_types::{GethExecStep, ToWord};
 
@@ -100,7 +99,7 @@ fn gen_log_step(
 
     // generates topic operation dynamically
     let topic_count = match exec_step.exec_state {
-        ExecState::Op(op_id) => (op_id.as_u8() - OpcodeId::LOG0.as_u8()) as usize,
+        ExecState::Op(op_id) => op_id.postfix().expect("opcode with postfix") as usize,
         _ => panic!("currently only handle successful log state"),
     };
 

eth-types/src/bytecode.rs

@@ -95,20 +95,20 @@ impl Bytecode {
     }
 
     /// Push
-    pub fn push(&mut self, n: usize, value: Word) -> &mut Self {
+    pub fn push(&mut self, n: u8, value: Word) -> &mut Self {
         debug_assert!((1..=32).contains(&n), "invalid push");
 
         // Write the op code
-        self.write_op_internal(OpcodeId::PUSH1.as_u8() + ((n - 1) as u8));
+        self.write_op((OpcodeId::push_n(n)).expect("valid push size"));
 
         let mut bytes = [0u8; 32];
         value.to_little_endian(&mut bytes);
         // Write the bytes MSB to LSB
         for i in 0..n {
-            self.write(bytes[n - 1 - i], false);
+            self.write(bytes[(n - 1 - i) as usize], false);
         }
         // Check if the full value could be pushed
-        for byte in bytes.iter().skip(n) {
+        for byte in bytes.iter().skip(n as usize) {
             debug_assert!(*byte == 0u8, "value too big for PUSH{}: {}", n, value);
         }
         self
@@ -217,17 +217,15 @@ pub enum OpcodeWithData {
     /// A non-push opcode
     Opcode(OpcodeId),
     /// A push opcode
-    Push(usize, Word),
+    Push(u8, Word),
 }
 
 impl OpcodeWithData {
     /// get the opcode
     pub fn opcode(&self) -> OpcodeId {
         match self {
             OpcodeWithData::Opcode(op) => *op,
-            OpcodeWithData::Push(n, _) => {
-                OpcodeId::try_from(OpcodeId::PUSH1.as_u8() + (*n as u8) - 1).unwrap()
-            }
+            OpcodeWithData::Push(n, _) => OpcodeId::push_n(*n).expect("valid push size"),
         }
     }
 }
@@ -240,7 +238,7 @@ impl FromStr for OpcodeWithData {
         let err = || Error::InvalidAsmError(op.to_string());
         if let Some(push) = op.strip_prefix("PUSH") {
             let n_value: Vec<_> = push.splitn(3, ['(', ')']).collect();
-            let n = n_value[0].parse::<usize>().map_err(|_| err())?;
+            let n = n_value[0].parse::<u8>().map_err(|_| err())?;
             if n < 1 || n > 32 {
                 return Err(err());
             }
@@ -276,12 +274,12 @@ impl<'a> Iterator for BytecodeIterator<'a> {
         self.0.next().map(|byte| {
             if let Ok(op) = OpcodeId::try_from(byte.value) {
                 if op.is_push() {
-                    let n = op.as_u8() - OpcodeId::PUSH1.as_u8() + 1;
+                    let n = op.postfix().expect("opcode with postfix");
                     let mut value = vec![0u8; n as usize];
                     for value_byte in value.iter_mut() {
                         *value_byte = self.0.next().unwrap().value;
                     }
-                    OpcodeWithData::Push(n as usize, Word::from(value.as_slice()))
+                    OpcodeWithData::Push(n, Word::from(value.as_slice()))
                 } else {
                     OpcodeWithData::Opcode(op)
                 }
@@ -301,7 +299,7 @@ impl From<Vec<u8>> for Bytecode {
             if let Ok(op) = OpcodeId::try_from(*byte) {
                 code.write_op(op);
                 if op.is_push() {
-                    let n = (op.as_u8() - OpcodeId::PUSH1.as_u8() + 1) as usize;
+                    let n = op.postfix().expect("opcode with postfix");
                     for _ in 0..n {
                         match input_iter.next() {
                             Some(v) => {
@@ -342,10 +340,8 @@ macro_rules! bytecode_internal {
     // PUSHX op codes
     ($code:ident, $x:ident ($v:expr) $($rest:tt)*) => {{
         debug_assert!($crate::evm_types::OpcodeId::$x.is_push(), "invalid push");
-        let n = $crate::evm_types::OpcodeId::$x.as_u8()
-            - $crate::evm_types::OpcodeId::PUSH1.as_u8()
-            + 1;
-        $code.push(n as usize, $v.into());
+        let n = $crate::evm_types::OpcodeId::$x.postfix().expect("opcode with postfix");
+        $code.push(n, $v.into());
         $crate::bytecode_internal!($code, $($rest)*);
     }};
     // Default opcode without any inputs

eth-types/src/evm_types/opcode_ids.rs

@@ -675,6 +675,40 @@ impl OpcodeId {
                 | OpcodeId::EXTCODECOPY
         )
     }
+
+    /// Returns PUSHn opcode from parameter n.
+    pub fn push_n(n: u8) -> Result<Self, Error> {
+        if (1..=32).contains(&n) {
+            OpcodeId::try_from(OpcodeId::PUSH1.as_u8() + n - 1)
+        } else {
+            Err(Error::InvalidOpConversion)
+        }
+    }
+
+    /// If operation has postfix returns it, otherwise None.
+    pub fn postfix(&self) -> Option<u8> {
+        if self.is_push() {
+            Some(self.as_u8() - OpcodeId::PUSH1.as_u8() + 1)
+        } else if self.is_dup() {
+            Some(self.as_u8() - OpcodeId::DUP1.as_u8() + 1)
+        } else if self.is_swap() {
+            Some(self.as_u8() - OpcodeId::SWAP1.as_u8() + 1)
+        } else if self.is_log() {
+            Some(self.as_u8() - OpcodeId::LOG0.as_u8())
+        } else {
+            None
+        }
+    }
+
+    /// Returns number of bytes used by immediate data. This is > 0 only for
+    /// push opcodes.
+    pub fn data_len(&self) -> usize {
+        if self.is_push() {
+            (self.as_u8() - OpcodeId::PUSH1.as_u8() + 1) as usize
+        } else {
+            0
+        }
+    }
 }
 
 impl TryFrom<u8> for OpcodeId {
@@ -1013,3 +1047,38 @@ impl fmt::Display for OpcodeId {
         write!(f, "{:?}", self)
     }
 }
+
+#[cfg(test)]
+mod opcode_ids_tests {
+    use super::*;
+
+    #[test]
+    fn push_n() {
+        assert!(matches!(OpcodeId::push_n(1), Ok(OpcodeId::PUSH1)));
+        assert!(matches!(OpcodeId::push_n(10), Ok(OpcodeId::PUSH10)));
+        assert!(matches!(
+            OpcodeId::push_n(100),
+            Err(Error::InvalidOpConversion)
+        ));
+        assert!(matches!(
+            OpcodeId::push_n(0),
+            Err(Error::InvalidOpConversion)
+        ));
+    }
+
+    #[test]
+    fn postfix() {
+        assert_eq!(OpcodeId::PUSH1.postfix(), Some(1));
+        assert_eq!(OpcodeId::PUSH10.postfix(), Some(10));
+        assert_eq!(OpcodeId::LOG2.postfix(), Some(2));
+        assert_eq!(OpcodeId::CALLCODE.postfix(), None);
+    }
+
+    #[test]
+    fn data_len() {
+        assert_eq!(OpcodeId::PUSH1.data_len(), 1);
+        assert_eq!(OpcodeId::PUSH10.data_len(), 10);
+        assert_eq!(OpcodeId::LOG2.data_len(), 0);
+        assert_eq!(OpcodeId::CALLCODE.data_len(), 0);
+    }
+}

zkevm-circuits/src/bytecode_circuit/bytecode_unroller.rs

@@ -657,7 +657,10 @@ pub fn unroll<F: Field>(bytes: Vec<u8>) -> UnrolledBytecode<F> {
 }
 
 fn is_push(byte: u8) -> bool {
-    OpcodeId::PUSH1.as_u8() <= byte && byte <= OpcodeId::PUSH32.as_u8()
+    match OpcodeId::try_from(byte) {
+        Ok(op) => op.is_push(),
+        Err(_) => false,
+    }
 }
 
 fn get_push_size(byte: u8) -> u64 {
@@ -720,7 +723,10 @@ mod tests {
         // Now add the different push ops
         for n in 1..=32 {
             let data_byte = OpcodeId::PUSH32.as_u8();
-            bytecode.push(n, Word::from_little_endian(&vec![data_byte; n][..]));
+            bytecode.push(
+                n,
+                Word::from_little_endian(&vec![data_byte; n as usize][..]),
+            );
             rows.push(BytecodeRow {
                 code_hash: Word::zero(),
                 tag: Fr::from(BytecodeFieldTag::Byte as u64),

zkevm-circuits/src/evm_circuit/execution/dup.rs

@@ -88,7 +88,7 @@ mod test {
     use mock::TestContext;
 
     fn test_ok(opcode: OpcodeId, value: Word) {
-        let n = (opcode.as_u8() - OpcodeId::DUP1.as_u8() + 1) as usize;
+        let n = opcode.postfix().expect("opcode with postfix");
         let mut bytecode = bytecode! {
             PUSH32(value)
         };

zkevm-circuits/src/evm_circuit/execution/logs.rs

@@ -208,7 +208,7 @@ impl<F: Field> ExecutionGadget<F> for LogGadget<F> {
             .assign(region, offset, step.memory_word_size(), [memory_address])?;
 
         let opcode = step.opcode.unwrap();
-        let topic_count = (opcode.as_u8() - OpcodeId::LOG0.as_u8()) as usize;
+        let topic_count = opcode.postfix().expect("opcode with postfix") as usize;
         assert!(topic_count <= 4);
 
         let is_persistent = call.is_persistent as u64;

zkevm-circuits/src/evm_circuit/execution/push.rs

@@ -129,12 +129,12 @@ impl<F: Field> ExecutionGadget<F> for PushGadget<F> {
         self.value
             .assign(region, offset, Some(value.to_le_bytes()))?;
 
-        let num_additional_pushed = (opcode.as_u8() - OpcodeId::PUSH1.as_u8()) as usize;
+        let num_additional_pushed = opcode.postfix().expect("opcode with postfix") - 1;
         for (idx, selector) in self.selectors.iter().enumerate() {
             selector.assign(
                 region,
                 offset,
-                Value::known(F::from((idx < num_additional_pushed) as u64)),
+                Value::known(F::from((idx < num_additional_pushed as usize) as u64)),
             )?;
         }
 
@@ -150,7 +150,7 @@ mod test {
     use mock::TestContext;
 
     fn test_ok(opcode: OpcodeId, bytes: &[u8]) {
-        assert!(bytes.len() as u8 == opcode.as_u8() - OpcodeId::PUSH1.as_u8() + 1,);
+        assert!(bytes.len() == opcode.data_len());
 
         let mut bytecode = bytecode! {
             .write_op(opcode)
@@ -239,7 +239,7 @@ mod test {
     }
 
     fn test_stack_overflow(opcode: OpcodeId, bytes: &[u8]) {
-        assert!(bytes.len() as u8 == opcode.as_u8() - OpcodeId::PUSH1.as_u8() + 1,);
+        assert!(bytes.len() == opcode.data_len());
 
         let mut bytecode = bytecode! {
             .write_op(opcode)

zkevm-circuits/src/evm_circuit/execution/swap.rs

@@ -96,7 +96,7 @@ mod test {
     use mock::TestContext;
 
     fn test_ok(opcode: OpcodeId, lhs: Word, rhs: Word) {
-        let n = (opcode.as_u8() - OpcodeId::SWAP1.as_u8() + 1) as usize;
+        let n = opcode.postfix().expect("opcode with postfix");
 
         let mut bytecode = bytecode! {
             PUSH32(lhs)

zkevm-circuits/src/witness/bytecode.rs

@@ -44,13 +44,17 @@ impl Bytecode {
 
         let mut push_data_left = 0;
         for (idx, byte) in self.bytes.iter().enumerate() {
-            let mut is_code = true;
-            if push_data_left > 0 {
-                is_code = false;
+            let is_code = push_data_left == 0;
+
+            if is_code {
+                let op = OpcodeId::try_from(*byte).expect("byte with valid opcode");
+
+                // push_data_left will be > 0 only if it is a push opcode
+                push_data_left = op.data_len();
+            } else {
                 push_data_left -= 1;
-            } else if (OpcodeId::PUSH1.as_u8()..=OpcodeId::PUSH32.as_u8()).contains(byte) {
-                push_data_left = *byte as usize - (OpcodeId::PUSH1.as_u8() - 1) as usize;
             }
+
             rows.push([
                 hash,
                 Value::known(F::from(BytecodeFieldTag::Byte as u64)),