lambdaclass / era_vm

EraVM implementation
MIT License
20 stars 3 forks source link

Encode existence of current context in `Execution` type #222

Open Oppen opened 2 months ago

Oppen commented 2 months ago

At every point of execution we expect a current execution context to exist. This causes us to always need to check the call stack for emptiness. Instead, we can encode that by using an unconditional field of type Context, separate from the stack, and refresh it and push the old value on call, as well as overwrite it with the top on returns. The initial change would look sort of like this:

diff --git a/src/execution.rs b/src/execution.rs
index 63acca8..648ec24 100644
--- a/src/execution.rs
+++ b/src/execution.rs
@@ -41,6 +41,7 @@ pub struct Execution {
     /// Equal flag
     pub flag_eq: bool,
     pub running_contexts: Vec<Context>,
+    pub current_context: Context,
     pub program: Vec<U256>,
     pub tx_number: u64,
     pub heaps: Heaps,
@@ -163,7 +164,7 @@ impl Execution {
         is_static: bool,
         stipend: u32,
     ) -> Result<(), EraVmError> {
-        let new_context = Context::new(
+        let mut new_context = Context::new(
             program_code,
             gas_stipend,
             contract_address,
@@ -178,6 +179,7 @@ impl Execution {
             is_static,
             stipend,
         );
+        std::swap(&new_context, &mut self.current_context);
         self.running_contexts.push(new_context);
         Ok(())
     }
@@ -186,7 +188,7 @@ impl Execution {
     }

     pub fn pop_frame(&mut self) -> Result<CallFrame, ContextError> {
-        let current_context = self.current_context_mut()?;
+        let current_context = self.current_context_mut();
         if current_context.near_call_frames.is_empty() {
             let context = self.pop_context()?;
             Ok(context.frame)
@@ -205,14 +207,12 @@ impl Execution {
         Ok(())
     }

-    pub fn current_context_mut(&mut self) -> Result<&mut Context, ContextError> {
-        self.running_contexts
-            .last_mut()
-            .ok_or(ContextError::NoContract)
+    pub fn current_context_mut(&mut self) -> &mut Context {
+        &mut self.current_context
     }

-    pub fn current_context(&self) -> Result<&Context, ContextError> {
-        self.running_contexts.last().ok_or(ContextError::NoContract)
+    pub fn current_context(&self) -> &Context {
+        &self.current_context
     }

     pub fn current_frame_mut(&mut self) -> Result<&mut CallFrame, ContextError> {

Note how we no longer need to return a Result<_> when just asking for the Context.