[Cairo1] Add a test to check that the output segment contains the return values when running in proof-mode

cairo1-run/src/cairo_run.rs

@@ -679,3 +679,86 @@ fn finalize_builtins(
     vm.builtins_final_stack_from_stack_pointer_dict(&builtin_name_to_stack_pointer, proof_mode)?;
     Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+    use std::path::Path;
+
+    use super::*;
+    use cairo_lang_compiler::{compile_cairo_project_at_path, CompilerConfig};
+    use cairo_vm::types::relocatable::Relocatable;
+    use rstest::rstest;
+
+    fn compile_to_sierra(filename: &str) -> SierraProgram {
+        let compiler_config = CompilerConfig {
+            replace_ids: true,
+            ..CompilerConfig::default()
+        };
+
+        compile_cairo_project_at_path(Path::new(filename), compiler_config).unwrap()
+    }
+
+    fn main_hash_panic_result(sierra_program: &SierraProgram) -> bool {
+        let main_func = find_function(sierra_program, "::main").unwrap();
+        main_func
+            .signature
+            .ret_types
+            .last()
+            .and_then(|rt| {
+                rt.debug_name
+                    .as_ref()
+                    .map(|n| n.as_ref().starts_with("core::panics::PanicResult::"))
+            })
+            .unwrap_or_default()
+    }
+
+    #[rstest]
+    #[case("../cairo_programs/cairo-1-programs/array_append.cairo")]
+    #[case("../cairo_programs/cairo-1-programs/array_get.cairo")]
+    #[case("../cairo_programs/cairo-1-programs/dictionaries.cairo")]
+    #[case("../cairo_programs/cairo-1-programs/enum_flow.cairo")]
+    #[case("../cairo_programs/cairo-1-programs/enum_match.cairo")]
+    #[case("../cairo_programs/cairo-1-programs/factorial.cairo")]
+    #[case("../cairo_programs/cairo-1-programs/fibonacci.cairo")]
+    #[case("../cairo_programs/cairo-1-programs/hello.cairo")]
+    #[case("../cairo_programs/cairo-1-programs/pedersen_example.cairo")]
+    #[case("../cairo_programs/cairo-1-programs/poseidon.cairo")]
+    #[case("../cairo_programs/cairo-1-programs/print.cairo")]
+    #[case("../cairo_programs/cairo-1-programs/array_append.cairo")]
+    #[case("../cairo_programs/cairo-1-programs/recursion.cairo")]
+    #[case("../cairo_programs/cairo-1-programs/sample.cairo")]
+    #[case("../cairo_programs/cairo-1-programs/simple_struct.cairo")]
+    #[case("../cairo_programs/cairo-1-programs/simple.cairo")]
+    #[case("../cairo_programs/cairo-1-programs/struct_span_return.cairo")]
+    fn check_append_ret_values_to_output_segment(#[case] filename: &str) {
+        // Compile to sierra
+        let sierra_program = compile_to_sierra(filename);
+        // Set proof_mode
+        let cairo_run_config = Cairo1RunConfig {
+            proof_mode: true,
+            layout: "all_cairo",
+            ..Default::default()
+        };
+        // Run program
+        let (_, vm, return_values) = cairo_run_program(&sierra_program, cairo_run_config).unwrap();
+        // When the return type is a PanicResult, we remove the panic wrapper when returning the ret values
+        // And handle the panics returning an error, so we need to add it here
+        let return_values = if main_hash_panic_result(&sierra_program) {
+            let mut rv = vec![Felt252::ZERO.into(), Felt252::ZERO.into()];
+            rv.extend_from_slice(&return_values);
+            rv
+        } else {
+            return_values
+        };
+        // Check that the output segment contains the return values
+        // The output builtin will always be the first builtin, so we know it's segment is 2
+        let output_builtin_segment = vm
+            .get_continuous_range((2, 0).into(), return_values.len())
+            .unwrap();
+        assert_eq!(output_builtin_segment, return_values, "{}", filename);
+        // Just for consistency, we will check that there are no values in the output segment after the return values
+        assert!(vm
+            .get_maybe(&Relocatable::from((2_isize, return_values.len())))
+            .is_none());
+    }
+}