Implement Squashed Felt252 Dict Entries libfunc

src/libfuncs.rs

@@ -69,6 +69,7 @@ mod mem;
 mod nullable;
 mod pedersen;
 mod poseidon;
+mod squashed_dict;
 mod starknet;
 mod r#struct;
 mod uint256;
@@ -176,9 +177,9 @@ impl LibfuncBuilder for CoreConcreteLibfunc {
             Self::Felt252Dict(selector) => self::felt252_dict::build(
                 context, registry, entry, location, helper, metadata, selector,
             ),
-            Self::Felt252SquashedDict(_) => {
-                native_panic!("Implement felt252_squashed_dict libfunc")
-            }
+            Self::Felt252SquashedDict(selector) => self::squashed_dict::build(
+                context, registry, entry, location, helper, metadata, selector,
+            ),
             Self::Felt252DictEntry(selector) => self::felt252_dict_entry::build(
                 context, registry, entry, location, helper, metadata, selector,
             ),

src/libfuncs/squashed_dict.rs

@@ -0,0 +1,330 @@
+use super::LibfuncHelper;
+use crate::{
+    error::{Error, Result},
+    metadata::{
+        realloc_bindings::ReallocBindingsMeta, runtime_bindings::RuntimeBindingsMeta,
+        MetadataStorage,
+    },
+    native_panic,
+    types::array::calc_data_prefix_offset,
+    utils::ProgramRegistryExt,
+};
+use cairo_lang_sierra::{
+    extensions::{
+        core::{CoreLibfunc, CoreType, CoreTypeConcrete},
+        lib_func::SignatureAndTypeConcreteLibfunc,
+        squashed_felt252_dict::SquashedFelt252DictConcreteLibfunc,
+        ConcreteLibfunc,
+    },
+    program_registry::ProgramRegistry,
+};
+use melior::{
+    dialect::llvm,
+    helpers::{ArithBlockExt, BuiltinBlockExt, LlvmBlockExt},
+    ir::{r#type::IntegerType, Block, BlockLike, Location},
+    Context,
+};
+use std::alloc::Layout;
+
+/// Select and call the correct libfunc builder function from the selector.
+pub fn build<'ctx, 'this>(
+    context: &'ctx Context,
+    registry: &ProgramRegistry<CoreType, CoreLibfunc>,
+    entry: &'this Block<'ctx>,
+    location: Location<'ctx>,
+    helper: &LibfuncHelper<'ctx, 'this>,
+    metadata: &mut MetadataStorage,
+    selector: &SquashedFelt252DictConcreteLibfunc,
+) -> Result<()> {
+    match selector {
+        SquashedFelt252DictConcreteLibfunc::IntoEntries(info) => {
+            build_into_entries(context, registry, entry, location, helper, metadata, info)
+        }
+    }
+}
+
+/// Get the layout of the tuple (felt252, T, T)
+fn get_inner_type_layout<'ctx, 'this>(
+    context: &'ctx Context,
+    registry: &ProgramRegistry<CoreType, CoreLibfunc>,
+    helper: &LibfuncHelper<'ctx, 'this>,
+    metadata: &mut MetadataStorage,
+    info: &SignatureAndTypeConcreteLibfunc,
+) -> Result<Layout> {
+    let array_ty = registry.get_type(&info.signature.branch_signatures[0].vars[0].ty)?;
+    let CoreTypeConcrete::Array(info) = array_ty else {
+        native_panic!("Received wrong type");
+    };
+    let (_, elem_layout) = registry.build_type_with_layout(context, helper, metadata, &info.ty)?;
+
+    Ok(elem_layout)
+}
+
+/// Generate MLIR operations for the `squashed_felt252_dict_entries` libfunc.
+///
+/// Receives a `SquashedFelt252Dict<T>` and returns an `Array<(felt252, T, T)>`. This
+/// array will have a tuple for each element on the dictionary. The first item represents
+/// the key of the element in the dictionary, it is followed by the first value and last
+/// value of that same element. Then (felt252, T, T) = (key, first_value, last_value).
+///
+/// # Caveats
+///
+/// In the tuple, the value that represents the first value will hold the value 0.
+///
+/// # Signature
+///
+/// ```cairo
+/// extern fn squashed_felt252_dict_entries<T>(
+///    dict: SquashedFelt252Dict<T>,
+/// ) -> Array<(felt252, T, T)> nopanic;
+/// ```
+pub fn build_into_entries<'ctx, 'this>(
+    context: &'ctx Context,
+    registry: &ProgramRegistry<CoreType, CoreLibfunc>,
+    entry: &'this Block<'ctx>,
+    location: Location<'ctx>,
+    helper: &LibfuncHelper<'ctx, 'this>,
+    metadata: &mut MetadataStorage,
+    info: &SignatureAndTypeConcreteLibfunc,
+) -> Result<()> {
+    metadata.get_or_insert_with(|| ReallocBindingsMeta::new(context, helper));
+
+    let dict_ptr = entry.arg(0)?;
+
+    // Get the size for the array (prefix + data)
+    let tuple_layout = get_inner_type_layout(context, registry, helper, metadata, info)?;
+    let data_prefix_size = calc_data_prefix_offset(tuple_layout);
+    let tuple_stride = entry.const_int_from_type(
+        context,
+        location,
+        tuple_layout.pad_to_align().size(),
+        IntegerType::new(context, 64).into(),
+    )?;
+    let data_prefix_size_value = entry.const_int(context, location, data_prefix_size, 64)?;
+    let (_, array_layout) = registry.build_type_with_layout(
+        context,
+        helper,
+        metadata,
+        &info.branch_signatures()[0].vars[0].ty,
+    )?;
+    let realloc_len = entry.const_int_from_type(
+        context,
+        location,
+        array_layout.pad_to_align().size(),
+        IntegerType::new(context, 64).into(),
+    )?;
+    // Create the pointer and alloc the necessary memory
+    let ptr_ty = llvm::r#type::pointer(context, 0);
+    let nullptr = entry.append_op_result(llvm::zero(ptr_ty, location))?;
+    let array_ptr = entry.append_op_result(ReallocBindingsMeta::realloc(
+        context,
+        nullptr,
+        realloc_len,
+        location,
+    )?)?;
+
+    // Runtime function that creates the array with its content
+    metadata
+        .get_mut::<RuntimeBindingsMeta>()
+        .ok_or(Error::MissingMetadata)?
+        .dict_into_entries(
+            context,
+            helper,
+            entry,
+            dict_ptr,
+            data_prefix_size_value,
+            tuple_stride,
+            array_ptr,
+            location,
+        )?;
+
+    // Extract the array from the pointer
+    let ptr_ty = llvm::r#type::pointer(context, 0);
+    let len_ty = IntegerType::new(context, 32).into();
+    let arr_ty = llvm::r#type::r#struct(context, &[ptr_ty, len_ty, len_ty, len_ty], false);
+    let entries_array = entry.load(context, location, array_ptr, arr_ty)?;
+
+    // Free the pointer where the array was stored
+    entry.append_operation(ReallocBindingsMeta::free(context, array_ptr, location)?);
+
+    helper.br(entry, 0, &[entries_array], location)
+}
+
+#[cfg(test)]
+mod test {
+    use crate::{jit_struct, load_cairo, utils::testing::run_program, Value};
+    use cairo_lang_sierra::program::Program;
+    use lazy_static::lazy_static;
+
+    lazy_static! {
+        static ref INTO_ENTRIES_EMPTY_DICT: (String, Program) = load_cairo! {
+            use core::dict::{Felt252Dict, Felt252DictEntryTrait, SquashedFelt252DictTrait};
+
+            fn into_entries_empty_dict() -> Array<(felt252, u8, u8)> {
+                let mut dict: Felt252Dict<u8> = Default::default();
+                dict.squash().into_entries()
+            }
+        };
+        static ref INTO_ENTRIES_U8_VALUES: (String, Program) = load_cairo! {
+            use core::dict::{Felt252Dict, Felt252DictEntryTrait, SquashedFelt252DictTrait};
+
+            fn into_entries_u8_values() -> Array<(felt252, u8, u8)> {
+                let mut dict: Felt252Dict<u8> = Default::default();
+                dict.insert(0, 0);
+                dict.insert(1, 1);
+                dict.insert(2, 2);
+                dict.squash().into_entries()
+            }
+        };
+        static ref INTO_ENTRIES_U32_VALUES: (String, Program) = load_cairo! {
+            use core::dict::{Felt252Dict, Felt252DictEntryTrait, SquashedFelt252DictTrait};
+
+            fn into_entries_u32_values() -> Array<(felt252, u32, u32)> {
+                let mut dict: Felt252Dict<u32> = Default::default();
+                dict.insert(0, 0);
+                dict.insert(1, 1);
+                dict.insert(2, 2);
+                dict.squash().into_entries()
+            }
+        };
+        static ref INTO_ENTRIES_U128_VALUES: (String, Program) = load_cairo! {
+            use core::dict::{Felt252Dict, Felt252DictEntryTrait, SquashedFelt252DictTrait};
+
+            fn into_entries_u128_values() -> Array<(felt252, u128, u128)> {
+                let mut dict: Felt252Dict<u128> = Default::default();
+                dict.insert(0, 0);
+                dict.insert(1, 1);
+                dict.insert(2, 2);
+                dict.squash().into_entries()
+            }
+        };
+        static ref INTO_ENTRIES_FELT252_VALUES: (String, Program) = load_cairo! {
+            use core::dict::{Felt252Dict, Felt252DictEntryTrait, SquashedFelt252DictTrait};
+
+            fn into_entries_felt252_values() -> Array<(felt252, felt252, felt252)> {
+                let mut dict: Felt252Dict<felt252> = Default::default();
+                dict.insert(0, 0);
+                dict.insert(1, 1);
+                dict.insert(2, 2);
+                dict.squash().into_entries()
+            }
+        };
+    }
+
+    #[test]
+    fn test_into_entries_empty_dict() {
+        let result =
+            run_program(&INTO_ENTRIES_EMPTY_DICT, "into_entries_empty_dict", &[]).return_value;
+        if let Value::Array(arr) = result {
+            assert_eq!(arr.len(), 0);
+        }
+    }
+
+    #[test]
+    fn test_into_entries_u8_values() {
+        let result =
+            run_program(&INTO_ENTRIES_U8_VALUES, "into_entries_u8_values", &[]).return_value;
+        if let Value::Array(arr) = result {
+            assert_eq!(
+                arr[0],
+                jit_struct!(Value::Felt252(0.into()), Value::Uint8(0), Value::Uint8(0))
+            );
+            assert_eq!(
+                arr[1],
+                jit_struct!(Value::Felt252(1.into()), Value::Uint8(0), Value::Uint8(1))
+            );
+            assert_eq!(
+                arr[2],
+                jit_struct!(Value::Felt252(2.into()), Value::Uint8(0), Value::Uint8(2))
+            );
+        }
+    }
+
+    #[test]
+    fn test_into_entries_u32_values() {
+        let result =
+            run_program(&INTO_ENTRIES_U32_VALUES, "into_entries_u32_values", &[]).return_value;
+        if let Value::Array(arr) = result {
+            assert_eq!(
+                arr[0],
+                jit_struct!(Value::Felt252(0.into()), Value::Uint32(0), Value::Uint32(0))
+            );
+            assert_eq!(
+                arr[1],
+                jit_struct!(Value::Felt252(1.into()), Value::Uint32(0), Value::Uint32(1))
+            );
+            assert_eq!(
+                arr[2],
+                jit_struct!(Value::Felt252(2.into()), Value::Uint32(0), Value::Uint32(2))
+            );
+        }
+    }
+
+    #[test]
+    fn test_into_entries_u128_values() {
+        let result =
+            run_program(&INTO_ENTRIES_U128_VALUES, "into_entries_u128_values", &[]).return_value;
+        if let Value::Array(arr) = result {
+            assert_eq!(
+                arr[0],
+                jit_struct!(
+                    Value::Felt252(0.into()),
+                    Value::Uint128(0),
+                    Value::Uint128(0)
+                )
+            );
+            assert_eq!(
+                arr[1],
+                jit_struct!(
+                    Value::Felt252(1.into()),
+                    Value::Uint128(0),
+                    Value::Uint128(1)
+                )
+            );
+            assert_eq!(
+                arr[2],
+                jit_struct!(
+                    Value::Felt252(2.into()),
+                    Value::Uint128(0),
+                    Value::Uint128(2)
+                )
+            );
+        }
+    }
+
+    #[test]
+    fn test_into_entries_felt252_values() {
+        let result = run_program(
+            &INTO_ENTRIES_FELT252_VALUES,
+            "into_entries_felt252_values",
+            &[],
+        )
+        .return_value;
+        if let Value::Array(arr) = result {
+            assert_eq!(
+                arr[0],
+                jit_struct!(
+                    Value::Felt252(0.into()),
+                    Value::Felt252(0.into()),
+                    Value::Felt252(0.into())
+                )
+            );
+            assert_eq!(
+                arr[1],
+                jit_struct!(
+                    Value::Felt252(1.into()),
+                    Value::Felt252(0.into()),
+                    Value::Felt252(1.into())
+                )
+            );
+            assert_eq!(
+                arr[2],
+                jit_struct!(
+                    Value::Felt252(2.into()),
+                    Value::Felt252(0.into()),
+                    Value::Felt252(2.into())
+                )
+            );
+        }
+    }
+}