[Cairo 1] Filter implicit arguments from return value (#1775)

* Filter implicit arguments from return value * Fix logic * Fix logic * Add CHANGELOG entry * Update cairo1-run/src/cairo_run.rs Co-authored-by: Mario Rugiero <[email protected]> * fmt --------- Co-authored-by: Mario Rugiero <[email protected]>
lambdaclass · May 28, 2024 · cc9baac · cc9baac
1 parent 26043e6
commit cc9baac
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 10 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,6 +2,8 @@
 
 #### Upcoming Changes
 
+* feat: Filter implicit arguments from return value in cairo1-run crate [#1775](https://github.com/lambdaclass/cairo-vm/pull/1775)
+
 * feat(BREAKING): Serialize inputs into output segment in cairo1-run crate:
   * Checks that only `Array<Felt252>` can be received by the program main function when running with with either `--proof_mode` or `--append_return_values`.
   * Copies the input value to the output segment right after the output in the format `[array_len, arr[0], arr[1],.., arr[n]]`.

diff --git a/cairo1-run/src/cairo_run.rs b/cairo1-run/src/cairo_run.rs
@@ -133,7 +133,16 @@ pub fn cairo_run_program(
     let initial_gas = 9999999999999_usize;
 
     // Fetch return type data
-    let return_type_id = main_func.signature.ret_types.last();
+    let return_type_id = match main_func.signature.ret_types.last() {
+        // We need to check if the last return type is indeed the function's return value and not an implicit return value
+        return_type @ Some(concrete_ty)
+            if get_info(&sierra_program_registry, concrete_ty)
+                .is_some_and(|info| !is_implicit_generic_id(&info.long_id.generic_id)) =>
+        {
+            return_type
+        }
+        _ => None,
+    };
 
     if (cairo_run_config.proof_mode || cairo_run_config.append_return_values)
         && !check_only_array_felt_input_type(
@@ -862,15 +871,7 @@ fn check_only_array_felt_input_type(
         .filter(|ty| {
             let info = get_info(sierra_program_registry, ty).unwrap();
             let generic_ty = &info.long_id.generic_id;
-            !(generic_ty != &SegmentArenaType::ID
-                || generic_ty == &GasBuiltinType::ID
-                || generic_ty == &BitwiseType::ID
-                || generic_ty == &EcOpType::ID
-                || generic_ty == &PedersenType::ID
-                || generic_ty == &PoseidonType::ID
-                || generic_ty == &RangeCheckType::ID
-                || generic_ty == &SegmentArenaType::ID
-                || generic_ty == &SystemType::ID)
+            !is_implicit_generic_id(generic_ty)
         })
         .collect_vec();
     if arg_types.is_empty() {
@@ -885,6 +886,22 @@ fn check_only_array_felt_input_type(
         false
     }
 }
+
+// Returns true if the generic id corresponds to an implicit argument (aka a builtin, gas, or system type)
+fn is_implicit_generic_id(generic_ty: &GenericTypeId) -> bool {
+    [
+        SegmentArenaType::ID,
+        GasBuiltinType::ID,
+        BitwiseType::ID,
+        EcOpType::ID,
+        PedersenType::ID,
+        PoseidonType::ID,
+        RangeCheckType::ID,
+        SegmentArenaType::ID,
+        SystemType::ID,
+    ]
+    .contains(generic_ty)
+}
 // Checks that the return type is either an Array<Felt252> or a PanicResult<Array<Felt252>> type
 fn check_only_array_felt_return_type(
     return_type_id: Option<&ConcreteTypeId>,