diff --git a/src/libfuncs.rs b/src/libfuncs.rs index 970e1eb93..0179f40f7 100644 --- a/src/libfuncs.rs +++ b/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, ), diff --git a/src/libfuncs/squashed_dict.rs b/src/libfuncs/squashed_dict.rs new file mode 100644 index 000000000..1afdaf62d --- /dev/null +++ b/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, + 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, + helper: &LibfuncHelper<'ctx, 'this>, + metadata: &mut MetadataStorage, + info: &SignatureAndTypeConcreteLibfunc, +) -> Result { + 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` 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( +/// dict: SquashedFelt252Dict, +/// ) -> Array<(felt252, T, T)> nopanic; +/// ``` +pub fn build_into_entries<'ctx, 'this>( + context: &'ctx Context, + registry: &ProgramRegistry, + 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::() + .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 = 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 = 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 = 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 = 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 = 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()) + ) + ); + } + } +} diff --git a/src/metadata/runtime_bindings.rs b/src/metadata/runtime_bindings.rs index aeefd674b..6ac84cdc7 100644 --- a/src/metadata/runtime_bindings.rs +++ b/src/metadata/runtime_bindings.rs @@ -48,6 +48,8 @@ enum RuntimeBinding { DebugPrint, ExtendedEuclideanAlgorithm, CircuitArithOperation, + DictLen, + DictIntoEntries, #[cfg(feature = "with-cheatcode")] VtableCheatcode, } @@ -76,6 +78,8 @@ impl RuntimeBinding { "cairo_native__extended_euclidean_algorithm" } RuntimeBinding::CircuitArithOperation => "cairo_native__circuit_arith_operation", + RuntimeBinding::DictLen => "cairo_native__dict_len", + RuntimeBinding::DictIntoEntries => "cairo_native__dict_into_entries", #[cfg(feature = "with-cheatcode")] RuntimeBinding::VtableCheatcode => "cairo_native__vtable_cheatcode", } @@ -124,6 +128,10 @@ impl RuntimeBinding { RuntimeBinding::GetCostsBuiltin => { crate::runtime::cairo_native__get_costs_builtin as *const () } + RuntimeBinding::DictLen => crate::runtime::cairo_native__dict_len as *const (), + RuntimeBinding::DictIntoEntries => { + crate::runtime::cairo_native__dict_into_entries as *const () + } RuntimeBinding::ExtendedEuclideanAlgorithm => return None, RuntimeBinding::CircuitArithOperation => return None, #[cfg(feature = "with-cheatcode")] @@ -716,6 +724,41 @@ impl RuntimeBindingsMeta { )) } + /// Register if necessary, then invoke the `dict_into_entries()` function. + /// + /// Returns an array with the tuples of the form (felt252, T, T) by storing it + /// on `array_ptr`. + #[allow(clippy::too_many_arguments)] + pub fn dict_into_entries<'c, 'a>( + &mut self, + context: &'c Context, + helper: &LibfuncHelper<'c, 'a>, + block: &'a Block<'c>, + dict_ptr: Value<'c, 'a>, + data_prefix_offset: Value<'c, 'a>, + tuple_stride: Value<'c, 'a>, + array_ptr: Value<'c, 'a>, + location: Location<'c>, + ) -> Result> + where + 'c: 'a, + { + let function = self.build_function( + context, + helper, + block, + location, + RuntimeBinding::DictIntoEntries, + )?; + + Ok(block.append_operation( + OperationBuilder::new("llvm.call", location) + .add_operands(&[function]) + .add_operands(&[dict_ptr, data_prefix_offset, tuple_stride, array_ptr]) + .build()?, + )) + } + // Register if necessary, then invoke the `get_costs_builtin()` function. #[allow(clippy::too_many_arguments)] pub fn get_costs_builtin<'c, 'a>( @@ -799,6 +842,8 @@ pub fn setup_runtime(find_symbol_ptr: impl Fn(&str) -> Option<*mut c_void>) { RuntimeBinding::DictDup, RuntimeBinding::GetCostsBuiltin, RuntimeBinding::DebugPrint, + RuntimeBinding::DictLen, + RuntimeBinding::DictIntoEntries, #[cfg(feature = "with-cheatcode")] RuntimeBinding::VtableCheatcode, ] { diff --git a/src/runtime.rs b/src/runtime.rs index 249d562fa..15583b207 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -1,6 +1,9 @@ #![allow(non_snake_case)] -use crate::utils::BuiltinCosts; +use crate::{ + starknet::ArrayAbi, + utils::{libc_malloc, BuiltinCosts}, +}; use cairo_lang_sierra_gas::core_libfunc_cost::{ DICT_SQUASH_REPEATED_ACCESS_COST, DICT_SQUASH_UNIQUE_KEY_COST, }; @@ -25,7 +28,7 @@ use std::{ mem::{forget, ManuallyDrop}, ops::Shl, os::fd::FromRawFd, - ptr, + ptr::{self, null_mut}, rc::Rc, }; use std::{ops::Mul, vec::IntoIter}; @@ -296,6 +299,110 @@ pub unsafe extern "C" fn cairo_native__dict_get( is_present as c_int } +pub unsafe extern "C" fn cairo_native__dict_len(dict_ptr: *const FeltDict) -> u64 { + let dict_rc = Rc::from_raw(dict_ptr); + let dict_len = dict_rc.mappings.len() as u64; + forget(dict_rc); // TODO: Should we forget it? + + dict_len +} + +/// Creates an array (Array<(felt252, T, T)>) by iterating the dictionary. +unsafe fn create_mlir_array( + dict: &mut FeltDict, + data_prefix_offset: u64, + tuple_stride: u64, +) -> ArrayAbi { + let len = dict.mappings.len(); + match len { + 0 => ArrayAbi { + ptr: null_mut(), + since: 0, + until: 0, + capacity: 0, + }, + _ => { + // Pointer to the space in memory with enough memory to hold the entire array + let ptr = libc_malloc( + ((tuple_stride * dict.mappings.len() as u64) + data_prefix_offset) as usize, + ); + + // Store the reference counter + ptr.cast::().write(1); + // Store the max lenght + ptr.byte_add(size_of::()) + .cast::() + .write(len as u32); + // Move the pointer past the prefix (reference counter and max length) into where the data + // will be stored + let ptr = ptr.byte_add(data_prefix_offset as usize); + + // Get the stride for the inner types of the tuple + let key_size = Layout::new::<[u8; 32]>().pad_to_align().size(); + let generic_ty_size = dict.layout.pad_to_align().size(); + + for (key, elem_index) in &dict.mappings { + // Move the ptr to the offset of the tuple we want to modify + let key_ptr = ptr.byte_add(tuple_stride as usize * elem_index) as *mut [u8; 32]; + + // Save the key and move to the offset of the 'first_value' + *key_ptr = *key; + let first_val_ptr = key_ptr.byte_add(key_size) as *mut u8; + first_val_ptr.write_bytes(0, generic_ty_size); + + // Get the element, move to the offset of the 'last_value' and save the element in that address + let element = dict.elements.byte_add(generic_ty_size * elem_index) as *mut u8; + let last_val_ptr = first_val_ptr.byte_add(generic_ty_size); + std::ptr::copy_nonoverlapping(element, last_val_ptr, generic_ty_size); + } + + let ptr_ptr = libc_malloc(size_of::<*mut ()>()).cast::<*mut c_void>(); + ptr_ptr.write(ptr); + + ArrayAbi { + ptr: ptr_ptr, + since: 0, + until: len as u32, + capacity: len as u32, + } + } + } +} + +/// Fills each of the tuples in the array with the corresponding content. +/// +/// Receives a pointer to the dictionary and a pointer to a space in memmory with the enough space +/// to store an array of the form Array<(felt252, T, T)> that has N tuples, where N is the quantity +/// of elements in the dictionary. The dictionary is iterated and for each element, a tuple is filled with the key +/// and the value. To fill the tuples, the 'tuple_stride' is used to move the pointer and get the +/// necessary offset. +/// +/// # Caveats +/// +/// Each tuple has the form (felt252, T, T) = (key, first_value, last_value). 'last_value' is represents +/// the value of the element in the dictionary and 'first_value' is always 0. +pub unsafe extern "C" fn cairo_native__dict_into_entries( + dict_ptr: *const FeltDict, + data_prefix_offset: u64, + tuple_stride: u64, + array_ptr: *mut ArrayAbi, +) { + let dict_rc = Rc::from_raw(dict_ptr); + + // There may be multiple references to the same dictionary (snapshots), but + // as snapshots cannot access the inner dictionary, then it is safe to modify it + // without cloning it. + let dict = Rc::as_ptr(&dict_rc) + .cast_mut() + .as_mut() + .expect("rc inner pointer should never be null"); + + let arr = create_mlir_array(dict, data_prefix_offset, tuple_stride); + + *array_ptr = arr; + forget(dict_rc); +} + /// Simulates the felt252_dict_squash libfunc. /// /// # Safety diff --git a/src/types/array.rs b/src/types/array.rs index 9d73d02c1..8fff26744 100644 --- a/src/types/array.rs +++ b/src/types/array.rs @@ -105,7 +105,7 @@ pub fn build<'ctx>( /// This function clones the array shallowly. That is, it'll increment the reference counter but not /// actually clone anything. The deep clone implementation is provided in `src/libfuncs/array.rs` as /// part of some libfuncs's implementations. -fn build_dup<'ctx>( +pub fn build_dup<'ctx>( context: &'ctx Context, module: &Module<'ctx>, registry: &ProgramRegistry, @@ -198,7 +198,7 @@ fn build_dup<'ctx>( /// This function decreases the reference counter of the array by one. /// If the reference counter reaches zero, then all the resources are freed. -fn build_drop<'ctx>( +pub fn build_drop<'ctx>( context: &'ctx Context, module: &Module<'ctx>, registry: &ProgramRegistry, @@ -406,6 +406,12 @@ fn build_drop<'ctx>( Ok(region) } +/// Returns the size of the prefix in an array. This prefix contains 2 +/// integers The integers are: +/// - Reference counter: the number of references to the allocation. +/// - Max length: The number of elements present in the allocation (not necessarily the length +/// array/span being accessed, but the whole allocation). It is used to know how many elements +/// to drop when freeing the allocation. pub fn calc_data_prefix_offset(layout: Layout) -> usize { get_integer_layout(32) .extend(get_integer_layout(32))