reapply CairoFunctionRunner and add search_pc with related tests and error handling

Author: naor-starkware

vm/src/types/errors/program_errors.rs

@@ -19,6 +19,10 @@ pub enum ProgramError {
     StrippedProgramNoMain,
     #[error("Hint PC ({0}) is greater or equal to program length ({1})")]
     InvalidHintPc(usize, usize),
+    #[error("Identifier \"{0}\" is type alias but has no destination")]
+    AliasMissingDestination(String),
+    #[error("invalid identifier type \"{1}\" for \"{0}\": expected \"alias\" or \"function\"")]
+    InvalidIdentifierTypeForPc(String, String),
 }
 
 #[cfg(test)]
@@ -31,4 +35,28 @@ mod tests {
         let formatted_error = format!("{error}");
         assert_eq!(formatted_error, "Entrypoint my_function not found");
     }
+
+    #[test]
+    fn format_alias_missing_destination_error() {
+        let error =
+            ProgramError::AliasMissingDestination(String::from("__main__.assert_nn"));
+        let formatted_error = format!("{error}");
+        assert_eq!(
+            formatted_error,
+            "Identifier \"__main__.assert_nn\" is type alias but has no destination"
+        );
+    }
+
+    #[test]
+    fn format_invalid_identifier_type_for_pc_error() {
+        let error = ProgramError::InvalidIdentifierTypeForPc(
+            String::from("__main__.my_struct"),
+            String::from("struct"),
+        );
+        let formatted_error = format!("{error}");
+        assert_eq!(
+            formatted_error,
+            "invalid identifier type \"struct\" for \"__main__.my_struct\": expected \"alias\" or \"function\""
+        );
+    }
 }

vm/src/vm/runners/cairo_function_runner.rs

@@ -6,6 +6,7 @@
 
 use crate::hint_processor::builtin_hint_processor::builtin_hint_processor_definition::BuiltinHintProcessor;
 use crate::hint_processor::hint_processor_definition::HintProcessor;
+use crate::serde::deserialize_program::Identifier;
 use crate::types::builtin_name::BuiltinName;
 use crate::types::errors::program_errors::ProgramError;
 use crate::types::instance_definitions::mod_instance_def::ModInstanceDef;
@@ -260,17 +261,73 @@ impl CairoFunctionRunner {
             .map(|builtin_runner| MaybeRelocatable::from((builtin_runner.base() as isize, 0)))
     }
 
-    /// Gets the program counter (PC) for a function entrypoint.
+    /// Returns the program counter (PC) for a function entrypoint by name.
+    ///
+    /// Looks up the identifier `__main__.{entrypoint}` in the program, then resolves it to a PC
+    /// (following alias chains if needed) via [`get_pc_from_identifier`](Self::get_pc_from_identifier).
+    ///
+    /// # Errors
+    /// - [`ProgramError::EntrypointNotFound`] if no identifier exists for the given name.
+    /// - [`RunnerError::NoPC`] if the resolved identifier has no PC (e.g. corrupt alias).
+    /// - [`ProgramError::AliasMissingDestination`] if an alias has no `destination`.
+    /// - [`ProgramError::InvalidIdentifierTypeForPc`] if the identifier type is not `"function"` or `"alias"`.
     #[allow(clippy::result_large_err)]
-    fn get_function_pc(&self, entrypoint: &str) -> std::result::Result<usize, CairoRunError> {
+    pub(crate) fn get_function_pc(
+        &self,
+        entrypoint: &str,
+    ) -> std::result::Result<usize, CairoRunError> {
         let full_name = format!("__main__.{entrypoint}");
         let identifier = self
             .program
             .get_identifier(&full_name)
             .ok_or_else(|| ProgramError::EntrypointNotFound(entrypoint.to_string()))?;
 
-        let pc = identifier.pc.ok_or(RunnerError::NoPC)?;
+        self.get_pc_from_identifier(identifier)
+    }
 
-        Ok(pc)
+    /// Resolves an identifier to its program counter (PC), following alias chains.
+    ///
+    /// - **`function`**: returns the identifier's `pc` if present.
+    /// - **`alias`**: resolves `destination` to another identifier and recurses until a function is found.
+    /// - **Other types** (e.g. `struct`, `const`): returns [`ProgramError::InvalidIdentifierTypeForPc`].
+    ///
+    /// # Errors
+    /// - [`RunnerError::NoPC`] when the identifier is a function but has no `pc`.
+    /// - [`ProgramError::AliasMissingDestination`] when the identifier is an alias but has no `destination`.
+    /// - [`ProgramError::EntrypointNotFound`] when the alias destination is not in the program.
+    /// - [`ProgramError::InvalidIdentifierTypeForPc`] when the identifier type is not `"function"` or `"alias"`.
+    #[allow(clippy::result_large_err)]
+    fn get_pc_from_identifier(
+        &self,
+        idetifier: &Identifier,
+    ) -> std::result::Result<usize, CairoRunError> {
+        match idetifier.type_.as_deref() {
+            Some("function") => {
+                let pc = idetifier.pc.ok_or(RunnerError::NoPC)?;
+                Ok(pc)
+            }
+            Some("alias") => {
+                let destination = idetifier.destination.as_deref().ok_or(
+                    ProgramError::AliasMissingDestination(
+                        idetifier.full_name.as_deref().unwrap_or("").to_string(),
+                    ),
+                )?;
+
+                let destination_identifier = self
+                    .runner
+                    .program
+                    .get_identifier(destination)
+                    .ok_or_else(|| ProgramError::EntrypointNotFound(destination.to_string()))?;
+                self.get_pc_from_identifier(destination_identifier)
+            }
+            v => {
+                let name = idetifier
+                    .full_name
+                    .clone()
+                    .unwrap_or_else(|| "<unknown>".to_string());
+                let type_str = v.unwrap_or("<unknown>").to_string();
+                Err(ProgramError::InvalidIdentifierTypeForPc(name, type_str).into())
+            }
+        }
     }
 }

vm/src/vm/runners/cairo_function_runner_test.rs

@@ -165,6 +165,38 @@ fn run_default_cairo0_happy_path() {
     assert_eq!(function_runner.get_return_values(0).unwrap(), vec![]);
 }
 
+#[test]
+// Test that get_function_pc resolves "__main__.assert_nn" (alias) to the PC of starkware.cairo.common.math.assert_nn (0).
+fn get_function_pc_assert_nn_resolves_alias_to_pc_0() {
+    let program = load_program(include_bytes!(
+        "../../../../cairo_programs/example_program.json"
+    ));
+    let function_runner = CairoFunctionRunner::new(&program).unwrap();
+
+    let pc = function_runner.get_function_pc("assert_nn").unwrap();
+    assert_eq!(
+        pc, 0,
+        "assert_nn is an alias to starkware.cairo.common.math.assert_nn which has pc 0"
+    );
+}
+
+#[test]
+// Test that get_function_pc returns the direct PC for "__main__.assert_nn_manual_implementation" (function with pc 4).
+fn get_function_pc_assert_nn_manual_implementation_returns_pc_4() {
+    let program = load_program(include_bytes!(
+        "../../../../cairo_programs/example_program.json"
+    ));
+    let function_runner = CairoFunctionRunner::new(&program).unwrap();
+
+    let pc = function_runner
+        .get_function_pc("assert_nn_manual_implementation")
+        .unwrap();
+    assert_eq!(
+        pc, 4,
+        "assert_nn_manual_implementation is a function with pc 4"
+    );
+}
+
 #[test]
 // Test that running a missing function name returns `EntrypointNotFound`.
 fn run_missing_entrypoint_returns_entrypoint_not_found() {