From 94a3d3b3387239f7340e48f9065cff483309db78 Mon Sep 17 00:00:00 2001
From: fmoletta <99273364+fmoletta@users.noreply.github.com>
Date: Wed, 13 Mar 2024 11:41:21 -0300
Subject: [PATCH] Fix error handling in `initialize_state` (#1657)

* Fix error handling in initialize_state

* Add changelog entry

* Refactor code according to suggestion
---
 CHANGELOG.md                      |  4 +++-
 vm/src/vm/runners/cairo_runner.rs | 35 +++++++++++--------------------
 2 files changed, 15 insertions(+), 24 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1547001753..fee05bc0f3 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,7 +3,9 @@
 #### Upcoming Changes
 * feat: add a `--tracer` option which hosts a web server that shows the line by line execution of cairo code along with memory registers [#1265](https://github.com/lambdaclass/cairo-vm/pull/1265)
 
-* feat: Make air public inputs deserializable [#1648](https://github.com/lambdaclass/cairo-vm/pull/1648)
+* feat: Fix error handling in `initialize_state`[#1657](https://github.com/lambdaclass/cairo-vm/pull/1657)
+
+* feat: Make air public inputs deserializable [#1657](https://github.com/lambdaclass/cairo-vm/pull/1648)
 
 * feat: Show only layout builtins in air private input [#1651](https://github.com/lambdaclass/cairo-vm/pull/1651)
 
diff --git a/vm/src/vm/runners/cairo_runner.rs b/vm/src/vm/runners/cairo_runner.rs
index 4dcd533cc9..7becd79eea 100644
--- a/vm/src/vm/runners/cairo_runner.rs
+++ b/vm/src/vm/runners/cairo_runner.rs
@@ -436,30 +436,19 @@ impl CairoRunner {
         entrypoint: usize,
         stack: Vec<MaybeRelocatable>,
     ) -> Result<(), RunnerError> {
-        if let Some(prog_base) = self.program_base {
-            let initial_pc = Relocatable {
-                segment_index: prog_base.segment_index,
-                offset: prog_base.offset + entrypoint,
-            };
-            self.initial_pc = Some(initial_pc);
-            vm.load_data(prog_base, &self.program.shared_program_data.data)
-                .map_err(RunnerError::MemoryInitializationError)?;
-
-            // Mark all addresses from the program segment as accessed
-            let base = self
-                .program_base
-                .unwrap_or_else(|| Relocatable::from((0, 0)));
-            for i in 0..self.program.shared_program_data.data.len() {
-                vm.segments.memory.mark_as_accessed((base + i)?);
-            }
-        }
-        if let Some(exec_base) = self.execution_base {
-            vm.segments
-                .load_data(exec_base, &stack)
-                .map_err(RunnerError::MemoryInitializationError)?;
-        } else {
-            return Err(RunnerError::NoProgBase);
+        let prog_base = self.program_base.ok_or(RunnerError::NoProgBase)?;
+        let exec_base = self.execution_base.ok_or(RunnerError::NoExecBase)?;
+        self.initial_pc = Some((prog_base + entrypoint)?);
+        vm.load_data(prog_base, &self.program.shared_program_data.data)
+            .map_err(RunnerError::MemoryInitializationError)?;
+
+        // Mark all addresses from the program segment as accessed
+        for i in 0..self.program.shared_program_data.data.len() {
+            vm.segments.memory.mark_as_accessed((prog_base + i)?);
         }
+        vm.segments
+            .load_data(exec_base, &stack)
+            .map_err(RunnerError::MemoryInitializationError)?;
         Ok(())
     }