feat(starknet_os): os_logger: implement OsLogger::exit_syscall

Author: dorimedini-starkware

crates/starknet_os/src/hint_processor/os_logger.rs

@@ -41,6 +41,8 @@ pub enum OsLoggerError {
     BuiltinsNotInSameSegment { builtin: BuiltinName, self_ptr: Relocatable, enter_ptr: Relocatable },
     #[error("Failed to build builtin pointer map: {0}.")]
     BuiltinPtrs(OsHintError),
+    #[error("Called exit_syscall with empty call stack.")]
+    CallStackEmpty,
     #[error("SyscallTrace should be finalized only once.")]
     DoubleFinalize,
     #[error("Failed to fetch identifier data for struct {0}.")]
@@ -60,6 +62,8 @@ pub enum OsLoggerError {
     ResourceAccessBeforeFinalize,
     #[error("The {0} syscall is not supposed to have an inner syscall.")]
     UnexpectedParentSyscall(String),
+    #[error("Unexpected syscall {actual:?}, expected {expected:?}.")]
+    UnexpectedSyscall { expected: SyscallSelector, actual: SyscallSelector },
     #[error("{0}")]
     UnknownBuiltin(String),
     #[error("Builtin {0} is not in the known sizes mapping {:?}.", BUILTIN_INSTANCE_SIZES)]
@@ -98,6 +102,10 @@ impl SyscallTrace {
     pub fn new(selector: SyscallSelector, is_deprecated: bool, tab_count: usize) -> Self {
         Self { selector, is_deprecated, tab_count, inner_syscalls: Vec::new(), resources: None }
     }
+
+    pub fn push_inner_syscall(&mut self, inner: SyscallTrace) {
+        self.inner_syscalls.push(inner);
+    }
 }
 
 impl ResourceFinalizer for SyscallTrace {
@@ -110,10 +118,10 @@ impl ResourceFinalizer for SyscallTrace {
     }
 }
 
-impl TryFrom<SyscallTrace> for String {
+impl TryFrom<&SyscallTrace> for String {
     type Error = OsLoggerError;
 
-    fn try_from(trace: SyscallTrace) -> OsLoggerResult<Self> {
+    fn try_from(trace: &SyscallTrace) -> OsLoggerResult<Self> {
         let deprecated_prefix = if trace.is_deprecated { "deprecated " } else { "" };
         let indentation = "  ".repeat(trace.tab_count + 1);
         let resources = trace.get_resources()?;
@@ -145,7 +153,6 @@ impl TryFrom<SyscallTrace> for String {
 pub struct OsTransactionTrace {
     tx_type: TransactionType,
     tx_hash: TransactionHash,
-    #[allow(dead_code)]
     syscalls: Vec<SyscallTrace>,
     resources: Option<ExecutionResources>,
 }
@@ -154,6 +161,10 @@ impl OsTransactionTrace {
     pub fn new(tx_type: TransactionType, tx_hash: TransactionHash) -> Self {
         Self { tx_type, tx_hash, syscalls: Vec::new(), resources: None }
     }
+
+    pub fn push_syscall(&mut self, syscall: SyscallTrace) {
+        self.syscalls.push(syscall);
+    }
 }
 
 impl ResourceFinalizer for OsTransactionTrace {
@@ -404,4 +415,51 @@ impl OsLogger {
 
         Ok(())
     }
+
+    #[allow(clippy::too_many_arguments)]
+    pub fn exit_syscall(
+        &mut self,
+        selector: SyscallSelector,
+        n_steps: usize,
+        range_check_ptr: Relocatable,
+        ids_data: &HashMap<String, HintReference>,
+        vm: &VirtualMachine,
+        ap_tracking: &ApTracking,
+        os_program: &Program,
+    ) -> OsLoggerResult<()> {
+        let mut current_syscall = self.syscall_stack.pop().ok_or(OsLoggerError::CallStackEmpty)?;
+        let enter_resources_counter =
+            self.resource_counter_stack.pop().ok_or(OsLoggerError::CallStackEmpty)?;
+        // A sanity check to ensure we store the syscall we work on.
+        if selector != current_syscall.selector {
+            return Err(OsLoggerError::UnexpectedSyscall {
+                actual: selector,
+                expected: current_syscall.selector,
+            });
+        }
+
+        let exit_resources_counter =
+            ResourceCounter::new(n_steps, range_check_ptr, ids_data, vm, ap_tracking, os_program)?;
+
+        current_syscall
+            .finalize_resources(exit_resources_counter.sub_counter(&enter_resources_counter)?)?;
+
+        if current_syscall.selector.is_calling_syscall() {
+            self.log(&format!("Exiting {}.", String::try_from(&current_syscall)?), false);
+        }
+
+        match self.syscall_stack.last_mut() {
+            Some(last_call) => {
+                last_call.push_inner_syscall(current_syscall);
+            }
+            None => {
+                self.current_tx
+                    .as_mut()
+                    .ok_or(OsLoggerError::NotInTxContext)?
+                    .push_syscall(current_syscall);
+            }
+        }
+
+        Ok(())
+    }
 }