Add heap size optimization

src/execution.rs

@@ -26,6 +26,7 @@ pub struct Stack {
 #[derive(Debug, Clone, PartialEq)]
 pub struct Heap {
     heap: Vec<u8>,
+    size: u32,
 }
 
 #[derive(Debug, Clone, PartialEq)]
@@ -448,34 +449,47 @@ impl Stack {
 
 impl Default for Heap {
     fn default() -> Self {
-        Self::new(vec![])
+        Self::new(vec![], 0)
     }
 }
 
 impl Heap {
-    pub fn new(values: Vec<u8>) -> Self {
-        Self { heap: values }
+    pub fn new(values: Vec<u8>, size: u32) -> Self {
+        Self { heap: values, size }
     }
-    // Returns how many ergs the expand costs
-    pub fn expand_memory(&mut self, address: u32) -> u32 {
+    fn expand_memory_inner(&mut self, address: u32) {
         if address >= self.heap.len() as u32 {
-            let old_size = self.heap.len() as u32;
             self.heap.resize(address as usize, 0);
-            return MEMORY_GROWTH_ERGS_PER_BYTE * (address - old_size);
+        }
+    }
+
+    // Returns how many ergs the expand costs
+    pub fn expand_memory(&mut self, address: u32) -> u32 {
+        if address >= self.size {
+            let old_size = self.size;
+            self.size = address;
+            return MEMORY_GROWTH_ERGS_PER_BYTE * (self.size - old_size);
         }
         0
     }
 
     pub fn store(&mut self, address: u32, value: U256) {
         let start = address as usize;
         let end = start + 32;
+        self.expand_memory_inner(end as u32);
         value.to_big_endian(&mut self.heap[start..end]);
     }
 
     pub fn read(&self, address: u32) -> U256 {
         let start = address as usize;
-        let end = start + 32;
-        U256::from_big_endian(&self.heap[start..end])
+        let mut end = start + 32;
+        if start >= self.heap.len() {
+            return U256::zero();
+        }
+        end = end.min(self.heap.len());
+        let mut result = self.heap[start..end].to_vec();
+        result.resize(32, 0);
+        U256::from_big_endian(&result)
     }
 
     pub fn expanded_read(&mut self, address: u32) -> (U256, u32) {
@@ -485,13 +499,19 @@ impl Heap {
     }
 
     pub fn read_byte(&self, address: u32) -> u8 {
+        if address >= self.heap.len() as u32 {
+            return 0;
+        }
         self.heap[address as usize]
     }
 
     pub fn read_from_pointer(&self, pointer: &FatPointer) -> U256 {
         let mut result = U256::zero();
         for i in 0..32 {
             let addr = pointer.start + pointer.offset + (31 - i);
+            if addr >= self.heap.len() as u32 {
+                continue;
+            }
             result |= U256::from(self.heap[addr as usize]) << (i * 8);
         }
         result
@@ -504,10 +524,10 @@ impl Heap {
     }
 
     pub fn len(&self) -> usize {
-        self.heap.len()
+        self.size as usize
     }
 
     pub fn is_empty(&self) -> bool {
-        self.heap.is_empty()
+        self.size == 0
     }
 }

src/heaps.rs

@@ -11,9 +11,10 @@ impl Heaps {
     pub fn new(calldata: Vec<u8>) -> Self {
         // The first heap can never be used because heap zero
         // means the current heap in precompile calls
+        let size = calldata.len() as u32;
         let heaps = vec![
             Heap::default(),
-            Heap::new(calldata),
+            Heap::new(calldata, size),
             Heap::default(),
             Heap::default(),
         ];
@@ -23,15 +24,13 @@ impl Heaps {
 
     pub fn allocate(&mut self) -> u32 {
         let id = self.heaps.len() as u32;
-        self.heaps
-            .push(Heap::new(vec![0; NEW_FRAME_MEMORY_STIPEND as usize]));
+        self.heaps.push(Heap::new(vec![], NEW_FRAME_MEMORY_STIPEND));
         id
     }
 
     pub fn allocate_copy(&mut self) -> u32 {
-        let id = self.heaps.len() as u32;
-        self.heaps
-            .push(Heap::new(vec![0; NEW_FRAME_MEMORY_STIPEND as usize]));
+        let id: u32 = self.heaps.len() as u32;
+        self.heaps.push(Heap::new(vec![], NEW_FRAME_MEMORY_STIPEND));
         id
     }