starknet_os: migrate test_unpack_felts (#11140)

Author: dorimedini-starkware

crates/starknet_os/src/hints/hint_implementation/stateless_compression/tests.rs

@@ -45,6 +45,81 @@ use crate::test_utils::cairo_runner::{
 
 const COMPRESSION_MODULE_PATH: &str = "starkware.starknet.core.os.data_availability.compression";
 
+/// Runs the OS unpack_felts function and returns the unpacked felts + resources used.
+fn cairo_unpack_felts(
+    compressed: &[Felt],
+    elm_bound: u128,
+    n_elms: usize,
+) -> (Vec<Felt>, ExecutionResources) {
+    let runner_config = EntryPointRunnerConfig {
+        layout: LayoutName::starknet,
+        add_main_prefix_to_entrypoint: false,
+        ..Default::default()
+    };
+    let mut implicit_args = vec![ImplicitArg::Builtin(BuiltinName::range_check)];
+    let (mut runner, program, entrypoint) = initialize_cairo_runner(
+        &runner_config,
+        OS_PROGRAM_BYTES,
+        &format!("{COMPRESSION_MODULE_PATH}.unpack_felts"),
+        &implicit_args,
+        HashMap::new(),
+    )
+    .unwrap();
+
+    // Function accepts start pointer of the compressed data, and destination pointer of the
+    // unpacked felts, as implicit arguments.
+    let compressed_start = runner
+        .vm
+        .gen_arg(&compressed.iter().map(|x| MaybeRelocatable::Int(*x)).collect::<Vec<_>>())
+        .unwrap()
+        .get_relocatable()
+        .unwrap();
+    let decompressed_start = runner.vm.add_memory_segment();
+    implicit_args.extend([
+        ImplicitArg::NonBuiltin(EndpointArg::Value(ValueArg::Single(compressed_start.into()))),
+        ImplicitArg::NonBuiltin(EndpointArg::Value(ValueArg::Single(decompressed_start.into()))),
+    ]);
+
+    let explicit_args = vec![
+        EndpointArg::Value(ValueArg::Single(n_elms.into())),
+        EndpointArg::Value(ValueArg::Single(Felt::from(elm_bound).into())),
+        EndpointArg::Value(ValueArg::Single(Felt::from(get_n_elms_per_felt(elm_bound)).into())),
+    ];
+
+    // Run the entrypoint.
+    // The unpacked data is stored in the segment starting at `decompressed_start`, the returned
+    // implicit value is the end of the unpacked data.
+    let state_reader = None;
+    let expected_explicit_return_values = vec![];
+    let (implicit_return_values, _explicit_return_values, _hint_processor) =
+        run_cairo_0_entrypoint(
+            entrypoint,
+            &explicit_args,
+            &implicit_args,
+            state_reader,
+            &mut runner,
+            &program,
+            &runner_config,
+            &expected_explicit_return_values,
+        )
+        .unwrap();
+
+    // The implicit return values are [range_check_ptr, compressed_start, decompressed_end].
+    assert_eq!(implicit_return_values.len(), 3);
+    let EndpointArg::Value(ValueArg::Single(MaybeRelocatable::RelocatableValue(decompressed_end))) =
+        implicit_return_values[2]
+    else {
+        panic!("Unexpected implicit return values, got: {implicit_return_values:?}");
+    };
+
+    // Read the compressed data from the segment and return.
+    let unpacked_data = runner
+        .vm
+        .get_integer_range(decompressed_start, (decompressed_end - decompressed_start).unwrap())
+        .unwrap();
+    (unpacked_data.into_iter().map(|f| *f).collect(), runner.get_execution_resources().unwrap())
+}
+
 /// Runs the OS compression function and returns the compressed data, plus the execution resources
 /// used to compress the data.
 fn cairo_compress(data: &[Felt]) -> (Vec<Felt>, ExecutionResources) {
@@ -244,6 +319,40 @@ fn test_usize_pack_and_unpack() {
     assert_eq!(nums, unpacked);
 }
 
+#[rstest]
+#[case::trivial(vec![], 0, None)]
+#[case::powers_of_two_plus_one((0..125).map(|i| 2_u128.pow(i) + 1).collect(), 63, Some(23))]
+#[case::small_number_range((0..2u128.pow(15)).collect(), 2usize.pow(11), Some(15))]
+#[case::large_number_range(
+    (0..2u128.pow(10)).map(|i| 2u128.pow(30) + i).collect(), 2usize.pow(7), Some(16)
+)]
+fn test_unpack_felts(
+    #[case] elms: Vec<u128>,
+    #[case] expected_compression_length: usize,
+    #[case] expected_n_steps_per_elm: Option<usize>,
+) {
+    let elm_bound = elms.iter().max().unwrap_or(&0) + 1;
+    let n_elms = elms.len();
+    let compressed = pack_in_felts(&elms, elm_bound);
+    let felt_bound = Felt::TWO.pow(251u8);
+    assert!(compressed.iter().all(|f| *f < felt_bound));
+    assert_eq!(compressed.len(), expected_compression_length);
+
+    let elm_felts: Vec<_> = elms.into_iter().map(Felt::from).collect();
+    assert_eq!(
+        elm_felts,
+        unpack_felts_to::<u128>(&compressed, n_elms, elm_bound)
+            .into_iter()
+            .map(Felt::from)
+            .collect::<Vec<Felt>>()
+    );
+    let (cairo_unpacked, resources) = cairo_unpack_felts(&compressed, elm_bound, n_elms);
+    assert_eq!(elm_felts, cairo_unpacked);
+    if n_elms > 0 {
+        assert_eq!(resources.n_steps / n_elms, expected_n_steps_per_elm.unwrap());
+    }
+}
+
 #[test]
 fn test_get_bucket_offsets() {
     let lengths = vec![2, 3, 5];