feat(starknet_os): add EndpointArg and support parameters passed by value (#5024)

Author: amosStarkware

crates/starknet_committer_and_os_cli/src/os_cli/tests/python_tests.rs

@@ -197,10 +197,9 @@ use cairo_vm::hint_processor::builtin_hint_processor::hint_code::{
     VM_EXIT_SCOPE,
     XS_SAFE_DIV,
 };
-use cairo_vm::types::relocatable::MaybeRelocatable;
-use cairo_vm::vm::runners::cairo_runner::CairoArg;
 use starknet_os::hints::enum_definition::{AggregatorHint, HintExtension, OsHint};
 use starknet_os::hints::types::HintEnum;
+use starknet_os::test_utils::cairo_runner::EndpointArg;
 use starknet_os::test_utils::errors::Cairo0EntryPointRunnerError;
 use starknet_os::test_utils::utils::run_cairo_function_and_check_result;
 use strum::IntoEnumIterator;
@@ -282,7 +281,7 @@ fn compare_os_hints(input: &str) -> OsPythonTestResult {
 fn test_cairo_function(
     program_str: &str,
     function_name: &str,
-    explicit_args: &[CairoArg],
+    explicit_args: &[EndpointArg],
     expected_retdata: &Retdata,
 ) -> OsPythonTestResult {
     run_cairo_function_and_check_result(
@@ -304,7 +303,7 @@ fn run_dummy_cairo_function(input: &str) -> OsPythonTestResult {
     test_cairo_function(
         input,
         "dummy_function",
-        &[MaybeRelocatable::from(param_1).into(), MaybeRelocatable::from(param_2).into()],
+        &[param_1.into(), param_2.into()],
         &retdata![(789 + param_1).into(), param_1.into(), param_2.into()],
     )
 }
@@ -324,10 +323,10 @@ fn test_constants(input: &str) -> OsPythonTestResult {
         input,
         "test_constants",
         &[
-            MaybeRelocatable::from(max_non_compressed_contract_address).into(),
-            MaybeRelocatable::from(alias_counter_storage_key).into(),
-            MaybeRelocatable::from(initial_available_alias).into(),
-            MaybeRelocatable::from(alias_contract_address).into(),
+            max_non_compressed_contract_address.into(),
+            alias_counter_storage_key.into(),
+            initial_available_alias.into(),
+            alias_contract_address.into(),
         ],
         &retdata![],
     )

crates/starknet_os/src/test_utils/cairo_runner.rs

@@ -5,49 +5,106 @@ use cairo_vm::types::layout_name::LayoutName;
 use cairo_vm::types::program::Program;
 use cairo_vm::types::relocatable::{MaybeRelocatable, Relocatable};
 use cairo_vm::vm::runners::cairo_runner::{CairoArg, CairoRunner};
+use starknet_types_core::felt::Felt;
 
 use crate::test_utils::errors::{Cairo0EntryPointRunnerError, ExplicitArgError};
 
 #[cfg(test)]
 #[path = "cairo_runner_test.rs"]
 mod test;
 
+/// An arg passed by value (i.e., a felt, tuple, named tuple or struct).
+#[derive(Clone, Debug)]
+pub enum ValueArg {
+    Single(Felt),
+    Array(Vec<Felt>),
+    Composed(Vec<EndpointArg>),
+}
+
+/// An arg passed as a pointer. i.e., a pointer to a felt, tuple, named tuple or struct, or a
+/// pointer to a pointer.
+#[derive(Clone, Debug)]
+pub enum PointerArg {
+    Array(Vec<Felt>),
+    Composed(Vec<EndpointArg>),
+}
+
+#[derive(Clone, Debug)]
+pub enum EndpointArg {
+    Value(ValueArg),
+    Pointer(PointerArg),
+}
+
+impl From<i32> for EndpointArg {
+    fn from(value: i32) -> Self {
+        Self::Value(ValueArg::Single(value.into()))
+    }
+}
+
+impl EndpointArg {
+    /// Converts an endpoint arg into a vector of cairo args.
+    /// The cairo VM loads struct / tuple / named tuple parameters by adding each of their fields
+    /// to the stack. This is why a single endpoint arg can be converted into multiple cairo args -
+    /// an arg of type Struct {a: felt, b: felt} will be converted into a vector of two cairo args
+    /// of type felt.
+    fn to_cairo_arg_vec(endpoint_arg: &EndpointArg) -> Vec<CairoArg> {
+        match endpoint_arg {
+            EndpointArg::Value(value_arg) => match value_arg {
+                ValueArg::Single(felt) => {
+                    vec![CairoArg::Single(MaybeRelocatable::Int(*felt))]
+                }
+                ValueArg::Array(felts) => felts
+                    .iter()
+                    .map(|felt| CairoArg::Single(MaybeRelocatable::Int(*felt)))
+                    .collect(),
+                ValueArg::Composed(endpoint_args) => {
+                    endpoint_args.iter().flat_map(Self::to_cairo_arg_vec).collect()
+                }
+            },
+            EndpointArg::Pointer(pointer_arg) => match pointer_arg {
+                PointerArg::Array(felts) => vec![CairoArg::Array(
+                    felts.iter().map(|felt| MaybeRelocatable::Int(*felt)).collect(),
+                )],
+                PointerArg::Composed(endpoint_args) => vec![CairoArg::Composed(
+                    endpoint_args.iter().flat_map(Self::to_cairo_arg_vec).collect(),
+                )],
+            },
+        }
+    }
+}
+
 /// Performs basic validations on the cairo arg. Assumes the arg is not a builtin.
 /// A successful result from this function does NOT guarantee that the arguments are valid.
 fn perform_basic_validations_on_cairo_arg(
     index: usize,
     expected_arg: &Member,
-    actual_arg: &CairoArg,
+    actual_arg: &EndpointArg,
 ) -> Result<(), Cairo0EntryPointRunnerError> {
-    if matches!(actual_arg, CairoArg::Single(MaybeRelocatable::RelocatableValue(_))) {
-        Err(ExplicitArgError::SingleRelocatableParam { index, actual_arg: actual_arg.clone() })?
-    }
-    let actual_arg_is_felt = matches!(actual_arg, CairoArg::Single(MaybeRelocatable::Int(_)));
-    let actual_arg_is_single = matches!(actual_arg, CairoArg::Single(_));
+    let actual_arg_is_felt = matches!(actual_arg, EndpointArg::Value(ValueArg::Single(_)));
+    let actual_arg_is_pointer = matches!(actual_arg, EndpointArg::Pointer(_));
+    let actual_arg_is_struct_or_tuple = !actual_arg_is_felt && !actual_arg_is_pointer;
 
     let expected_arg_is_pointer = expected_arg.cairo_type.ends_with("*");
     let expected_arg_is_felt = expected_arg.cairo_type == "felt";
+    let expected_arg_is_struct_or_tuple = !expected_arg_is_felt && !expected_arg_is_pointer;
 
-    if expected_arg_is_felt != actual_arg_is_felt || expected_arg_is_pointer == actual_arg_is_single
+    if expected_arg_is_felt != actual_arg_is_felt
+        || expected_arg_is_pointer != actual_arg_is_pointer
+        || expected_arg_is_struct_or_tuple != actual_arg_is_struct_or_tuple
     {
         Err(ExplicitArgError::Mismatch {
             index,
             expected: expected_arg.clone(),
             actual: actual_arg.clone(),
         })?;
     };
-    // expected arg is tuple / named tuple / struct.
-    if !expected_arg_is_felt && !expected_arg_is_pointer {
-        // TODO(Amos): Load tuple / named tuple / struct parameters to stack and remove this error.
-        Err(ExplicitArgError::UnsupportedArgType { index, expected_arg: expected_arg.clone() })?;
-    };
     Ok(())
 }
 
 /// Performs basic validations on the explicit arguments. A successful result from this function
 /// does NOT guarantee that the arguments are valid.
 fn perform_basic_validations_on_explicit_args(
-    explicit_args: &[CairoArg],
+    explicit_args: &[EndpointArg],
     program: &Program,
     entrypoint: &str,
 ) -> Result<(), Cairo0EntryPointRunnerError> {
@@ -82,7 +139,7 @@ pub fn run_cairo_0_entry_point(
     program: &Program,
     entrypoint: &str,
     n_expected_return_values: usize,
-    explicit_args: &[CairoArg],
+    explicit_args: &[EndpointArg],
     mut hint_processor: impl HintProcessor,
 ) -> Result<Retdata, Cairo0EntryPointRunnerError> {
     // TODO(Amos): Perform complete validations.
@@ -97,7 +154,9 @@ pub fn run_cairo_0_entry_point(
     let program_base: Option<Relocatable> = None;
     cairo_runner.initialize_segments(program_base);
 
-    let entrypoint_args: Vec<&CairoArg> = explicit_args.iter().collect();
+    let entrypoint_args: Vec<CairoArg> =
+        explicit_args.iter().flat_map(EndpointArg::to_cairo_arg_vec).collect();
+    let entrypoint_args: Vec<&CairoArg> = entrypoint_args.iter().collect();
     let verify_secure = true;
     let program_segment_size: Option<usize> = None;
     // TODO(Amos): Pass implicit args to the cairo runner.

crates/starknet_os/src/test_utils/cairo_runner_test.rs

@@ -1,7 +1,6 @@
 use blockifier::retdata;
-use cairo_vm::types::relocatable::MaybeRelocatable;
-use cairo_vm::vm::runners::cairo_runner::CairoArg;
 
+use crate::test_utils::cairo_runner::{EndpointArg, PointerArg, ValueArg};
 use crate::test_utils::errors::Cairo0EntryPointRunnerError;
 use crate::test_utils::utils::run_cairo_function_and_check_result;
 
@@ -15,38 +14,90 @@ use crate::test_utils::utils::run_cairo_function_and_check_result;
 ///     b: felt,
 /// }
 ///
-/// func dummy_function(number: felt, array: felt*, tuple: felt*, simple_struct: SimpleStruct*,
-///     compound_struct: CompoundStruct*) -> (res1: felt, res2: felt, res3: felt) {
+/// func pass_felt_and_pointers(number: felt, array: felt*, tuple: felt*, simple_struct:
+///     SimpleStruct*, compound_struct: CompoundStruct*) -> (res1: felt, res2: felt, res3: felt) {
 ///     let res1 = number + array[0];
 ///     let res2 = tuple[0] + tuple[1];
 ///     let res3 = simple_struct.a + compound_struct.simple_struct.b;
 ///     return (res1=res1, res2=res2, res3=res3);
 /// }
+///
+/// func pass_structs_and_tuples(tuple: (felt, felt), named_tuple: (a: felt, b:felt), simple_struct:
+///     SimpleStruct, compound_struct: CompoundStruct) -> (res1: felt, res2: felt, res3: felt) {
+///     let res1 = tuple[0] + tuple[1];
+///     let res2 = named_tuple.a + named_tuple.b;
+///     let res3 = simple_struct.a + compound_struct.simple_struct.b;
+///     return (res1=res1, res2=res2, res3=res3);
+/// }
 const COMPILED_DUMMY_FUNCTION: &str = include_str!("compiled_dummy_function.json");
 
 #[test]
-fn test_cairo0_function_runner() -> Result<(), Cairo0EntryPointRunnerError> {
+fn test_felt_and_pointers() -> Result<(), Cairo0EntryPointRunnerError> {
     let number = 2;
     let (first_array_val, second_array_val) = (3, 4);
     let (first_tuple_val, second_tuple_val) = (5, 6);
     let (first_simple_struct_val, second_simple_struct_val) = (7, 8);
     let compound_struct_val = 9;
-    let array = CairoArg::Array(vec![first_array_val.into(), second_array_val.into()]);
-    let tuple = CairoArg::Array(vec![first_tuple_val.into(), second_tuple_val.into()]);
-    let simple_struct =
-        CairoArg::Array(vec![first_simple_struct_val.into(), second_simple_struct_val.into()]);
-    let compound_struct = CairoArg::Composed(vec![
-        MaybeRelocatable::from(compound_struct_val).into(),
+    let array = EndpointArg::Pointer(PointerArg::Array(vec![
+        first_array_val.into(),
+        second_array_val.into(),
+    ]));
+    let tuple = EndpointArg::Pointer(PointerArg::Array(vec![
+        first_tuple_val.into(),
+        second_tuple_val.into(),
+    ]));
+    let simple_struct = EndpointArg::Pointer(PointerArg::Array(vec![
+        first_simple_struct_val.into(),
+        second_simple_struct_val.into(),
+    ]));
+    let compound_struct = EndpointArg::Pointer(PointerArg::Composed(vec![
+        compound_struct_val.into(),
         simple_struct.clone(),
-    ]);
+    ]));
     run_cairo_function_and_check_result(
         COMPILED_DUMMY_FUNCTION,
-        "dummy_function",
-        &[MaybeRelocatable::from(number).into(), array, tuple, simple_struct, compound_struct],
+        "pass_felt_and_pointers",
+        &[number.into(), array, tuple, simple_struct, compound_struct],
         &retdata![
             (number + first_array_val).into(),
             (first_tuple_val + second_tuple_val).into(),
             (first_simple_struct_val + second_simple_struct_val).into()
         ],
     )
 }
+
+#[test]
+fn test_tuples_and_structs() -> Result<(), Cairo0EntryPointRunnerError> {
+    let (first_tuple_val, second_tuple_val) = (3, 4);
+    let (first_named_tuple_val, second_named_tuple_val) = (5, 6);
+    let (first_simple_struct_val, second_simple_struct_val) = (7, 8);
+    let compound_struct_val = 9;
+    let tuple =
+        EndpointArg::Value(ValueArg::Array(vec![first_tuple_val.into(), second_tuple_val.into()]));
+    let named_tuple = EndpointArg::Value(ValueArg::Array(vec![
+        first_named_tuple_val.into(),
+        second_named_tuple_val.into(),
+    ]));
+    let simple_struct = EndpointArg::Value(ValueArg::Array(vec![
+        first_simple_struct_val.into(),
+        second_simple_struct_val.into(),
+    ]));
+    let simple_struct_pointer = EndpointArg::Pointer(PointerArg::Array(vec![
+        first_simple_struct_val.into(),
+        second_simple_struct_val.into(),
+    ]));
+    let compound_struct = EndpointArg::Value(ValueArg::Composed(vec![
+        compound_struct_val.into(),
+        simple_struct_pointer,
+    ]));
+    run_cairo_function_and_check_result(
+        COMPILED_DUMMY_FUNCTION,
+        "pass_structs_and_tuples",
+        &[tuple, named_tuple, simple_struct, compound_struct],
+        &retdata![
+            (first_tuple_val + second_tuple_val).into(),
+            (first_named_tuple_val + second_named_tuple_val).into(),
+            (first_simple_struct_val + second_simple_struct_val).into()
+        ],
+    )
+}