From 2a213dc0841d0d029a3366d59ab728be1906c2cb Mon Sep 17 00:00:00 2001
From: Gianbelinche <39842759+gianbelinche@users.noreply.github.com>
Date: Wed, 4 Sep 2024 13:06:19 -0300
Subject: [PATCH] Add heap size optimization

---
 src/execution.rs | 42 +++++++++++++++++++++++++++++++-----------
 src/heaps.rs     | 11 +++++------
 2 files changed, 36 insertions(+), 17 deletions(-)

diff --git a/src/execution.rs b/src/execution.rs
index 63acca84..dd9fe01d 100644
--- a/src/execution.rs
+++ b/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,20 +449,26 @@ 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
     }
@@ -469,13 +476,20 @@ impl Heap {
     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,6 +499,9 @@ impl Heap {
     }
 
     pub fn read_byte(&self, address: u32) -> u8 {
+        if address >= self.heap.len() as u32 {
+            return 0;
+        }
         self.heap[address as usize]
     }
 
@@ -492,6 +509,9 @@ impl Heap {
         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
     }
 }
diff --git a/src/heaps.rs b/src/heaps.rs
index e02f4747..8fadfddf 100644
--- a/src/heaps.rs
+++ b/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
     }