Return error instead of panic in memory operation (#928)

evm/src/witness/errors.rs

@@ -1,3 +1,5 @@
+use ethereum_types::U256;
+
 #[allow(dead_code)]
 #[derive(Debug)]
 pub enum ProgramError {
@@ -8,4 +10,13 @@ pub enum ProgramError {
     InvalidJumpiDestination,
     StackOverflow,
     KernelPanic,
+    MemoryError(MemoryError),
+}
+
+#[allow(clippy::enum_variant_names)]
+#[derive(Debug)]
+pub enum MemoryError {
+    ContextTooLarge { context: U256 },
+    SegmentTooLarge { segment: U256 },
+    VirtTooLarge { virt: U256 },
 }

evm/src/witness/memory.rs

@@ -12,6 +12,9 @@ use MemoryChannel::{Code, GeneralPurpose};
 
 use crate::cpu::kernel::constants::global_metadata::GlobalMetadata;
 use crate::memory::segments::Segment;
+use crate::witness::errors::MemoryError::{ContextTooLarge, SegmentTooLarge, VirtTooLarge};
+use crate::witness::errors::ProgramError;
+use crate::witness::errors::ProgramError::MemoryError;
 
 impl MemoryChannel {
     pub fn index(&self) -> usize {
@@ -41,19 +44,25 @@ impl MemoryAddress {
         }
     }
 
-    pub(crate) fn new_u256s(context: U256, segment: U256, virt: U256) -> Self {
-        assert!(context.bits() <= 32, "context too large: {}", context);
-        assert!(
-            segment < Segment::COUNT.into(),
-            "segment too large: {}",
-            segment
-        );
-        assert!(virt.bits() <= 32, "virt too large: {}", virt);
-        Self {
+    pub(crate) fn new_u256s(
+        context: U256,
+        segment: U256,
+        virt: U256,
+    ) -> Result<Self, ProgramError> {
+        if context.bits() > 32 {
+            return Err(MemoryError(ContextTooLarge { context }));
+        }
+        if segment >= Segment::COUNT.into() {
+            return Err(MemoryError(SegmentTooLarge { segment }));
+        }
+        if virt.bits() > 32 {
+            return Err(MemoryError(VirtTooLarge { virt }));
+        }
+        Ok(Self {
             context: context.as_usize(),
             segment: segment.as_usize(),
             virt: virt.as_usize(),
-        }
+        })
     }
 
     pub(crate) fn increment(&mut self) {

evm/src/witness/operation.rs

@@ -111,7 +111,7 @@ pub(crate) fn generate_keccak_general<F: Field>(
         stack_pop_with_log_and_fill::<4, _>(state, &mut row)?;
     let len = len.as_usize();
 
-    let base_address = MemoryAddress::new_u256s(context, segment, base_virt);
+    let base_address = MemoryAddress::new_u256s(context, segment, base_virt)?;
     let input = (0..len)
         .map(|i| {
             let address = MemoryAddress {
@@ -608,7 +608,7 @@ pub(crate) fn generate_mload_general<F: Field>(
 
     let (val, log_read) = mem_read_gp_with_log_and_fill(
         3,
-        MemoryAddress::new_u256s(context, segment, virt),
+        MemoryAddress::new_u256s(context, segment, virt)?,
         state,
         &mut row,
     );