diff --git a/src/executor.rs b/src/executor.rs index cd9827650..6f0b9a8aa 100644 --- a/src/executor.rs +++ b/src/executor.rs @@ -14,10 +14,10 @@ use crate::{ SEGMENT_ARENA_BUILTIN_SIZE, }, native_panic, - runtime::{BLAKE_CALL_COUNT, BUILTIN_COSTS}, + runtime::{BLAKE_CALL_COUNT, BOX_ARENA, BUILTIN_COSTS}, starknet::{handler::StarknetSyscallHandlerCallbacks, StarknetSyscallHandler}, types::TypeBuilder, - utils::{libc_free, BuiltinCosts, RangeExt}, + utils::{BuiltinCosts, RangeExt}, values::Value, }; use bumpalo::Bump; @@ -151,7 +151,6 @@ fn invoke_dynamic( #[cfg(feature = "with-cheatcode")] syscall_handler.as_mut().map(|h| h as *mut _ as *mut ()), ); - // Generate argument list. let mut iter = args.iter(); for item in function_signature.param_types.iter().filter_map(|type_id| { @@ -365,6 +364,7 @@ fn invoke_dynamic( pub(crate) struct InvocationGuard { builtin_costs: BuiltinCosts, blake_call_count: u64, + box_arena: Bump, #[cfg(feature = "with-cheatcode")] syscall_handler: Option<*mut ()>, } @@ -379,6 +379,7 @@ impl InvocationGuard { Self { builtin_costs: BUILTIN_COSTS.replace(builtin_costs), blake_call_count: BLAKE_CALL_COUNT.with(|c| c.replace(0)), + box_arena: BOX_ARENA.with(|c| std::mem::replace(&mut *c.borrow_mut(), Bump::new())), #[cfg(feature = "with-cheatcode")] syscall_handler: syscall_handler.map(|ptr| { let previous_value = crate::starknet::SYSCALL_HANDLER_VTABLE.get(); @@ -393,6 +394,7 @@ impl Drop for InvocationGuard { fn drop(&mut self) { BUILTIN_COSTS.set(self.builtin_costs); BLAKE_CALL_COUNT.with(|c| c.set(self.blake_call_count)); + BOX_ARENA.with(|c| std::mem::swap(&mut *c.borrow_mut(), &mut self.box_arena)); #[cfg(feature = "with-cheatcode")] if let Some(previous_value) = self.syscall_handler { crate::starknet::SYSCALL_HANDLER_VTABLE.set(previous_value); @@ -436,7 +438,6 @@ fn parse_result( let ptr = return_ptr.unwrap_or_else(|| NonNull::new_unchecked(ret_registers[0] as *mut ())); let value = Value::from_ptr(ptr, &info.ty, registry, true)?; - libc_free(ptr.cast().as_ptr()); Ok(value) }, CoreTypeConcrete::EcPoint(_) | CoreTypeConcrete::EcState(_) => Ok(Value::from_ptr( @@ -574,7 +575,6 @@ fn parse_result( } else { let ptr = NonNull::new_unchecked(ptr); let value = Value::from_ptr(ptr, &info.ty, registry, true)?; - libc_free(ptr.as_ptr().cast()); Ok(value) } }, diff --git a/src/libfuncs/array.rs b/src/libfuncs/array.rs index e58f53630..c81dc4f25 100644 --- a/src/libfuncs/array.rs +++ b/src/libfuncs/array.rs @@ -203,7 +203,7 @@ pub fn build_span_from_tuple<'ctx, 'this>( location, )?)?; - // Move the data into the array and free the original tuple. Since the tuple and the array are + // Move the data into the array. Since the tuple and the array are // represented the same way, a simple memcpy is enough. entry.memcpy( context, @@ -212,11 +212,6 @@ pub fn build_span_from_tuple<'ctx, 'this>( data_ptr, array_len_bytes_val, ); - entry.append_operation(ReallocBindingsMeta::free( - context, - entry.argument(0)?.into(), - location, - )?); // Allocate metadata struct: { refcount: u32, max_len: u32, data_ptr: *mut u8 } let metadata_size = entry.const_int(context, location, calc_metadata_size(), 64)?; diff --git a/src/libfuncs/blake.rs b/src/libfuncs/blake.rs index 42cf086f4..58d90f19f 100644 --- a/src/libfuncs/blake.rs +++ b/src/libfuncs/blake.rs @@ -1,7 +1,7 @@ use cairo_lang_sierra::{ extensions::{ blake::BlakeConcreteLibfunc, - core::{CoreLibfunc, CoreType}, + core::{CoreLibfunc, CoreType, CoreTypeConcrete}, lib_func::SignatureOnlyConcreteLibfunc, }, program_registry::ProgramRegistry, @@ -15,6 +15,8 @@ use melior::{ use crate::{ error::{panic::ToNativeAssertError, Result}, metadata::{runtime_bindings::RuntimeBindingsMeta, MetadataStorage}, + native_panic, + types::TypeBuilder, }; use super::LibfuncHelper; @@ -57,12 +59,12 @@ pub fn build<'ctx, 'this>( #[allow(clippy::too_many_arguments)] fn build_blake_operation<'ctx, 'this>( context: &'ctx Context, - _registry: &ProgramRegistry, + registry: &ProgramRegistry, entry: &'this Block<'ctx>, location: Location<'ctx>, helper: &LibfuncHelper<'ctx, 'this>, metadata: &mut MetadataStorage, - _info: &SignatureOnlyConcreteLibfunc, + info: &SignatureOnlyConcreteLibfunc, finalize: bool, ) -> Result<()> { let state_ptr = entry.arg(0)?; @@ -70,14 +72,29 @@ fn build_blake_operation<'ctx, 'this>( let message = entry.arg(2)?; let k_finalize = entry.const_int(context, location, finalize as u8, 1)?; + // The state parameter is `Box<[u32; 8]>`; we need the layout of the inner + // `[u32; 8]` to size the output slot, not the layout of the box pointer. + let CoreTypeConcrete::Box(box_info) = + registry.get_type(&info.signature.param_signatures[0].ty)? + else { + native_panic!("blake state parameter should be a Box"); + }; + let inner_layout = registry.get_type(&box_info.ty)?.layout(registry)?; + let size = entry.const_int(context, location, inner_layout.size(), 64)?; + let align = entry.const_int(context, location, inner_layout.align(), 64)?; + let runtime_bindings = metadata .get_mut::() .to_native_assert_error("runtime library should be available")?; + let out_state_ptr = + runtime_bindings.box_alloc(context, helper, entry, location, size, align)?; + runtime_bindings.libfunc_blake_compress( context, helper, entry, + out_state_ptr, state_ptr, message, bytes_count, @@ -85,7 +102,7 @@ fn build_blake_operation<'ctx, 'this>( location, )?; - helper.br(entry, 0, &[state_ptr], location)?; + helper.br(entry, 0, &[out_state_ptr], location)?; Ok(()) } diff --git a/src/libfuncs/box.rs b/src/libfuncs/box.rs index 1f73ffb84..fc8a2591c 100644 --- a/src/libfuncs/box.rs +++ b/src/libfuncs/box.rs @@ -1,13 +1,17 @@ //! # Box libfuncs //! //! A heap allocated value, which is internally a pointer that can't be null. +//! +//! Allocations are made from the per-invocation arena (`cairo_native__box_alloc` runtime +//! function) and are freed all at once when the program invocation completes. Individual +//! boxes are therefore never freed on unbox. use std::alloc::Layout; use super::LibfuncHelper; use crate::{ error::Result, - metadata::{realloc_bindings::ReallocBindingsMeta, MetadataStorage}, + metadata::{runtime_bindings::RuntimeBindingsMeta, MetadataStorage}, types::TypeBuilder, utils::ProgramRegistryExt, }; @@ -21,12 +25,11 @@ use cairo_lang_sierra::{ program_registry::ProgramRegistry, }; use melior::{ - dialect::{ - llvm::{self, r#type::pointer, LoadStoreOptions}, - ods, - }, + dialect::llvm::{self, LoadStoreOptions}, helpers::{ArithBlockExt, BuiltinBlockExt}, - ir::{attribute::IntegerAttribute, r#type::IntegerType, Block, BlockLike, Location, Value}, + ir::{ + attribute::IntegerAttribute, r#type::IntegerType, Block, BlockLike, Location, Module, Value, + }, Context, }; @@ -69,37 +72,38 @@ pub fn build_into_box<'ctx, 'this>( metadata: &mut MetadataStorage, info: &SignatureAndTypeConcreteLibfunc, ) -> Result<()> { - if metadata.get::().is_none() { - metadata.insert(ReallocBindingsMeta::new(context, helper)); - } - let inner_type = registry.get_type(&info.ty)?; let inner_layout = inner_type.layout(registry)?; - let ptr = into_box(context, entry, location, entry.arg(0)?, inner_layout)?; + let ptr = into_box( + context, + helper.module, + entry, + location, + entry.arg(0)?, + inner_layout, + metadata, + )?; helper.br(entry, 0, &[ptr], location) } /// Receives a value and inserts it into a box +/// Allocate an arena slot and store `inner_val` into it, returning a pointer. pub fn into_box<'ctx, 'this>( context: &'ctx Context, + module: &Module<'ctx>, entry: &'this Block<'ctx>, location: Location<'ctx>, inner_val: Value<'ctx, 'this>, inner_layout: Layout, + metadata: &mut MetadataStorage, ) -> Result> { - let value_len = entry.const_int(context, location, inner_layout.pad_to_align().size(), 64)?; - let ptr = entry - .append_operation(ods::llvm::mlir_zero(context, pointer(context, 0), location).into()) - .result(0)? - .into(); - let ptr = entry - .append_operation(ReallocBindingsMeta::realloc( - context, ptr, value_len, location, - )?) - .result(0)? - .into(); + let size = entry.const_int(context, location, inner_layout.size(), 64)?; + let align = entry.const_int(context, location, inner_layout.align(), 64)?; + + let rtb = metadata.get_or_insert_with(RuntimeBindingsMeta::default); + let ptr = rtb.box_alloc(context, module, entry, location, size, align)?; entry.append_operation(llvm::store( context, @@ -125,8 +129,6 @@ pub fn build_unbox<'ctx, 'this>( metadata: &mut MetadataStorage, info: &SignatureAndTypeConcreteLibfunc, ) -> Result<()> { - metadata.get_or_insert_with(|| ReallocBindingsMeta::new(context, helper)); - let value = unbox( context, registry, entry, location, helper, metadata, &info.ty, )?; @@ -162,8 +164,6 @@ pub fn unbox<'ctx, 'this>( .result(0)? .into(); - entry.append_operation(ReallocBindingsMeta::free(context, entry.arg(0)?, location)?); - Ok(value) } diff --git a/src/libfuncs/const.rs b/src/libfuncs/const.rs index 9ff3f00a2..e8db4906b 100644 --- a/src/libfuncs/const.rs +++ b/src/libfuncs/const.rs @@ -3,8 +3,8 @@ use super::LibfuncHelper; use crate::{ error::{Error, Result}, - libfuncs::{r#enum::build_enum_value, r#struct::build_struct_value}, - metadata::{realloc_bindings::ReallocBindingsMeta, MetadataStorage}, + libfuncs::{r#box::into_box, r#enum::build_enum_value, r#struct::build_struct_value}, + metadata::MetadataStorage, native_panic, types::TypeBuilder, utils::{felt_to_unsigned, ProgramRegistryExt, RangeExt}, @@ -23,7 +23,7 @@ use cairo_lang_sierra::{ program_registry::ProgramRegistry, }; use melior::{ - dialect::llvm::{self, r#type::pointer}, + dialect::llvm::{self}, helpers::{ArithBlockExt, BuiltinBlockExt, LlvmBlockExt}, ir::{r#type::IntegerType, Block, Location, Value}, Context, @@ -59,10 +59,6 @@ pub fn build_const_as_box<'ctx, 'this>( metadata: &mut MetadataStorage, info: &ConstAsBoxConcreteLibfunc, ) -> Result<()> { - if metadata.get::().is_none() { - metadata.insert(ReallocBindingsMeta::new(context, helper)); - } - let const_type_outer = registry.get_type(&info.const_type)?; // Create constant @@ -78,16 +74,15 @@ pub fn build_const_as_box<'ctx, 'this>( let const_ty = registry.get_type(&const_type.inner_ty)?; let inner_layout = const_ty.layout(registry)?; - // Create box - let value_len = entry.const_int(context, location, inner_layout.pad_to_align().size(), 64)?; - - let ptr = entry.append_op_result(llvm::zero(pointer(context, 0), location))?; - let ptr = entry.append_op_result(ReallocBindingsMeta::realloc( - context, ptr, value_len, location, - )?)?; - - // Store constant in box - entry.store(context, location, ptr, value)?; + let ptr = into_box( + context, + helper.module, + entry, + location, + value, + inner_layout, + metadata, + )?; helper.br(entry, 0, &[ptr], location) } diff --git a/src/libfuncs/enum.rs b/src/libfuncs/enum.rs index b45f7a371..ceca93114 100644 --- a/src/libfuncs/enum.rs +++ b/src/libfuncs/enum.rs @@ -6,10 +6,7 @@ use super::LibfuncHelper; use crate::{ error::{panic::ToNativeAssertError, Error, Result}, libfuncs::r#box::into_box, - metadata::{ - enum_snapshot_variants::EnumSnapshotVariantsMeta, realloc_bindings::ReallocBindingsMeta, - MetadataStorage, - }, + metadata::{enum_snapshot_variants::EnumSnapshotVariantsMeta, MetadataStorage}, native_assert, native_panic, types::TypeBuilder, utils::ProgramRegistryExt, @@ -592,9 +589,6 @@ pub fn build_boxed_match<'ctx, 'this>( metadata: &mut MetadataStorage, info: &EnumBoxedMatchConcreteLibfunc, ) -> Result<()> { - metadata.get_or_insert_with(|| ReallocBindingsMeta::new(context, helper)); - - // Get the variant type IDs from the concrete libfunc info let variant_ids = &info.variants; match variant_ids.len() { @@ -626,7 +620,6 @@ pub fn build_boxed_match<'ctx, 'this>( variant_ids, )?; - // Tag is at offset 0 in the box, so load it directly from the pointer let tag_val = entry.load(context, location, entry.arg(0)?, tag_ty)?; let default_block = helper.append_block(Block::new(&[])); @@ -666,7 +659,6 @@ pub fn build_boxed_match<'ctx, 'this>( default_block.append_operation(llvm::unreachable(location)); } - // Enum variants. for (i, (block, (payload_ty, payload_layout))) in variant_blocks.into_iter().zip(variant_tys).enumerate() { @@ -683,13 +675,6 @@ pub fn build_boxed_match<'ctx, 'this>( block.load(context, location, ptr, payload_ty)? }; - // Free the input box - block.append_operation(ReallocBindingsMeta::free( - context, - entry.arg(0)?, - location, - )?); - // Get the output variant type layout for boxing let output_variant_type_id = &info.branch_signatures()[i].vars[0].ty; let CoreTypeConcrete::Box(output_box_info) = @@ -704,8 +689,15 @@ pub fn build_boxed_match<'ctx, 'this>( &output_box_info.ty, )?; - // Box the payload - let boxed_payload = into_box(context, block, location, payload_val, output_layout)?; + let boxed_payload = into_box( + context, + helper.module, + block, + location, + payload_val, + output_layout, + metadata, + )?; helper.br(block, i, &[boxed_payload], location)?; } diff --git a/src/libfuncs/struct.rs b/src/libfuncs/struct.rs index d643edbd7..3ccb0d13f 100644 --- a/src/libfuncs/struct.rs +++ b/src/libfuncs/struct.rs @@ -4,7 +4,7 @@ use super::LibfuncHelper; use crate::{ error::Result, libfuncs::r#box::{into_box, unbox}, - metadata::{realloc_bindings::ReallocBindingsMeta, MetadataStorage}, + metadata::MetadataStorage, native_panic, types::TypeBuilder, utils::ProgramRegistryExt, @@ -180,8 +180,6 @@ pub fn build_boxed_deconstruct<'ctx, 'this>( metadata: &mut MetadataStorage, info: &ConcreteStructBoxedDeconstructLibfunc, ) -> Result<()> { - metadata.get_or_insert_with(|| ReallocBindingsMeta::new(context, helper)); - // Unbox the container let CoreTypeConcrete::Box(box_info) = registry.get_type(&info.param_signatures()[0].ty)? else { native_panic!("Should receibe a Box type as argument"); @@ -205,7 +203,15 @@ pub fn build_boxed_deconstruct<'ctx, 'this>( let (_, member_layout) = registry.build_type_with_layout(context, helper, metadata, member_type_id)?; // Box the member - let member = into_box(context, entry, location, member, member_layout)?; + let member = into_box( + context, + helper.module, + entry, + location, + member, + member_layout, + metadata, + )?; fields.push(member); } diff --git a/src/metadata/runtime_bindings.rs b/src/metadata/runtime_bindings.rs index a052c5778..f7bc107d7 100644 --- a/src/metadata/runtime_bindings.rs +++ b/src/metadata/runtime_bindings.rs @@ -58,6 +58,7 @@ enum RuntimeBinding { QM31Sub, QM31Mul, QM31Div, + BoxAlloc, #[cfg(feature = "with-cheatcode")] VtableCheatcode, } @@ -98,6 +99,7 @@ impl RuntimeBinding { RuntimeBinding::QM31Sub => "cairo_native__libfunc__qm31__qm31_sub", RuntimeBinding::QM31Mul => "cairo_native__libfunc__qm31__qm31_mul", RuntimeBinding::QM31Div => "cairo_native__libfunc__qm31__qm31_div", + RuntimeBinding::BoxAlloc => "cairo_native__box_alloc", #[cfg(feature = "with-cheatcode")] RuntimeBinding::VtableCheatcode => "cairo_native__vtable_cheatcode", } @@ -168,6 +170,7 @@ impl RuntimeBinding { | RuntimeBinding::U252ExtendedEuclideanAlgorithm | RuntimeBinding::U384ExtendedEuclideanAlgorithm => return None, RuntimeBinding::CircuitArithOperation => return None, + RuntimeBinding::BoxAlloc => crate::runtime::cairo_native__box_alloc as *const (), #[cfg(feature = "with-cheatcode")] RuntimeBinding::VtableCheatcode => { crate::starknet::cairo_native__vtable_cheatcode as *const () @@ -512,6 +515,7 @@ impl RuntimeBindingsMeta { context: &'c Context, module: &Module, block: &'a Block<'c>, + out_state: Value<'c, 'a>, state: Value<'c, 'a>, message: Value<'c, 'a>, count_bytes: Value<'c, 'a>, @@ -532,7 +536,7 @@ impl RuntimeBindingsMeta { Ok(block.append_operation( OperationBuilder::new("llvm.call", location) .add_operands(&[function]) - .add_operands(&[state, message, count_bytes, finalize]) + .add_operands(&[out_state, state, message, count_bytes, finalize]) .build()?, )) } @@ -758,6 +762,33 @@ impl RuntimeBindingsMeta { Ok(block.load(context, location, res_ptr, qm31_ty)?) } + /// Register the `cairo_native__box_alloc` global if necessary, then invoke + /// it to allocate `size` bytes with alignment `align` from the per-invocation + /// arena. Returns an opaque pointer to the allocated memory. + pub fn box_alloc<'c, 'a>( + &mut self, + context: &'c Context, + module: &Module, + block: &'a Block<'c>, + location: Location<'c>, + size: Value<'c, 'a>, + align: Value<'c, 'a>, + ) -> Result> + where + 'c: 'a, + { + let function = + self.build_function(context, module, block, location, RuntimeBinding::BoxAlloc)?; + + Ok(block.append_op_result( + OperationBuilder::new("llvm.call", location) + .add_operands(&[function]) + .add_operands(&[size, align]) + .add_results(&[llvm::r#type::pointer(context, 0)]) + .build()?, + )?) + } + /// Register if necessary, then invoke the `dict_alloc_new()` function. /// /// Returns a opaque pointer as the result. @@ -1049,6 +1080,7 @@ pub fn setup_runtime(find_symbol_ptr: impl Fn(&str) -> Option<*mut c_void>) { RuntimeBinding::QM31Sub, RuntimeBinding::QM31Mul, RuntimeBinding::QM31Div, + RuntimeBinding::BoxAlloc, #[cfg(feature = "with-cheatcode")] RuntimeBinding::VtableCheatcode, ] { diff --git a/src/runtime.rs b/src/runtime.rs index be34e766e..befcb8163 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -5,6 +5,7 @@ use crate::{ types::array::ArrayMetadata, utils::{blake_utils, libc_malloc, BuiltinCosts}, }; +use bumpalo::Bump; use cairo_lang_sierra_gas::core_libfunc_cost::{ DICT_SQUASH_REPEATED_ACCESS_COST, DICT_SQUASH_UNIQUE_KEY_COST, }; @@ -23,7 +24,7 @@ use starknet_types_core::{ }; use std::{ alloc::{dealloc, realloc, Layout}, - cell::Cell, + cell::{Cell, RefCell}, collections::{hash_map::Entry, HashMap}, ffi::{c_int, c_void}, fs::File, @@ -36,6 +37,20 @@ use std::{ }; use std::{ops::Mul, vec::IntoIter}; +// Thread-local handle to the box arena currently active for this invocation. +thread_local! { + pub(crate) static BOX_ARENA: RefCell = RefCell::new(Bump::new()); +} + +/// Allocate `size` bytes with `align` alignment from the per-invocation box arena. +pub unsafe extern "C" fn cairo_native__box_alloc(size: u64, align: u64) -> *mut u8 { + BOX_ARENA.with(|arena| { + let layout = Layout::from_size_align(size as usize, align as usize) + .expect("cairo_native__box_alloc: invalid layout"); + arena.borrow_mut().alloc_layout(layout).as_ptr() + }) +} + lazy_static! { pub static ref HALF_PRIME: Felt = Felt::from_dec_str( "1809251394333065606848661391547535052811553607665798349986546028067936010240" @@ -151,7 +166,8 @@ pub unsafe extern "C" fn cairo_native__libfunc__hades_permutation( } pub unsafe extern "C" fn cairo_native__libfunc__blake_compress( - state: &mut [u32; 8], + out_state: &mut [u32; 8], + state: &[u32; 8], message: &[u32; 16], count_bytes: u32, finalize: bool, @@ -165,7 +181,7 @@ pub unsafe extern "C" fn cairo_native__libfunc__blake_compress( 0, ); - *state = new_state; + *out_state = new_state; // Track blake invocations: Blake doesn't have an implicit counter argument // like buffer-based builtins, so we count calls here directly. diff --git a/src/types/box.rs b/src/types/box.rs index 192c6b576..cc12a37a6 100644 --- a/src/types/box.rs +++ b/src/types/box.rs @@ -17,7 +17,7 @@ use crate::{ error::Result, metadata::{ drop_overrides::DropOverridesMeta, dup_overrides::DupOverridesMeta, - realloc_bindings::ReallocBindingsMeta, MetadataStorage, + runtime_bindings::RuntimeBindingsMeta, MetadataStorage, }, types::TypeBuilder, utils::ProgramRegistryExt, @@ -30,12 +30,9 @@ use cairo_lang_sierra::{ program_registry::ProgramRegistry, }; use melior::{ - dialect::{func, llvm, ods}, + dialect::{func, llvm}, helpers::{ArithBlockExt, BuiltinBlockExt, LlvmBlockExt}, - ir::{ - attribute::IntegerAttribute, r#type::IntegerType, Block, BlockLike, Location, Module, - Region, Type, - }, + ir::{Block, BlockLike, Location, Module, Region, Type}, Context, }; @@ -56,10 +53,12 @@ pub fn build<'ctx>( metadata, info.self_ty(), |metadata| { - // There's no need to build the type here because it'll always be built within - // `build_dup`. - - Ok(Some(build_dup(context, module, registry, metadata, &info)?)) + registry.build_type(context, module, metadata, &info.ty)?; + if DupOverridesMeta::is_overriden(metadata, &info.ty) { + Ok(Some(build_dup(context, module, registry, metadata, &info)?)) + } else { + Ok(None) + } }, )?; DropOverridesMeta::register_with( @@ -69,12 +68,14 @@ pub fn build<'ctx>( metadata, info.self_ty(), |metadata| { - // There's no need to build the type here because it'll always be built within - // `build_drop`. - - Ok(Some(build_drop( - context, module, registry, metadata, &info, - )?)) + registry.build_type(context, module, metadata, &info.ty)?; + if DropOverridesMeta::is_overriden(metadata, &info.ty) { + Ok(Some(build_drop( + context, module, registry, metadata, &info, + )?)) + } else { + Ok(None) + } }, )?; @@ -89,49 +90,30 @@ fn build_dup<'ctx>( info: &WithSelf, ) -> Result> { let location = Location::unknown(context); - if metadata.get::().is_none() { - metadata.insert(ReallocBindingsMeta::new(context, module)); - } let inner_ty = registry.get_type(&info.ty)?; - let inner_len = inner_ty.layout(registry)?.pad_to_align().size(); + let inner_layout = inner_ty.layout(registry)?; + let inner_len = inner_layout.size(); + let inner_align = inner_layout.align(); let inner_ty = inner_ty.build(context, module, registry, metadata, &info.ty)?; let region = Region::new(); let entry = region.append_block(Block::new(&[(llvm::r#type::pointer(context, 0), location)])); - let null_ptr = - entry.append_op_result(llvm::zero(llvm::r#type::pointer(context, 0), location))?; - let inner_len_val = entry.const_int(context, location, inner_len, 64)?; + let size_val = entry.const_int(context, location, inner_len, 64)?; + let align_val = entry.const_int(context, location, inner_align, 64)?; let src_value = entry.arg(0)?; - let dst_value = entry.append_op_result(ReallocBindingsMeta::realloc( - context, - null_ptr, - inner_len_val, - location, - )?)?; + // build_dup is only registered when the inner type has a dup override. + let rtb = metadata.get_or_insert_with(RuntimeBindingsMeta::default); + let dst_value = rtb.box_alloc(context, module, &entry, location, size_val, align_val)?; - if DupOverridesMeta::is_overriden(metadata, &info.ty) { - let value = entry.load(context, location, src_value, inner_ty)?; - let values = DupOverridesMeta::invoke_override( - context, registry, module, &entry, &entry, location, metadata, &info.ty, value, - )?; - entry.store(context, location, src_value, values.0)?; - entry.store(context, location, dst_value, values.1)?; - } else { - entry.append_operation( - ods::llvm::intr_memcpy_inline( - context, - dst_value, - src_value, - IntegerAttribute::new(IntegerType::new(context, 64).into(), inner_len as i64), - IntegerAttribute::new(IntegerType::new(context, 1).into(), 0), - location, - ) - .into(), - ); - } + let value = entry.load(context, location, src_value, inner_ty)?; + let values = DupOverridesMeta::invoke_override( + context, registry, module, &entry, &entry, location, metadata, &info.ty, value, + )?; + entry.store(context, location, src_value, values.0)?; + entry.store(context, location, dst_value, values.1)?; entry.append_operation(func::r#return(&[src_value, dst_value], location)); Ok(region) @@ -145,24 +127,19 @@ fn build_drop<'ctx>( info: &WithSelf, ) -> Result> { let location = Location::unknown(context); - if metadata.get::().is_none() { - metadata.insert(ReallocBindingsMeta::new(context, module)); - } let inner_ty = registry.build_type(context, module, metadata, &info.ty)?; let region = Region::new(); let entry = region.append_block(Block::new(&[(llvm::r#type::pointer(context, 0), location)])); + // build_drop is only registered when the inner type has a drop override. let value = entry.arg(0)?; - if DropOverridesMeta::is_overriden(metadata, &info.ty) { - let value = entry.load(context, location, value, inner_ty)?; - DropOverridesMeta::invoke_override( - context, registry, module, &entry, &entry, location, metadata, &info.ty, value, - )?; - } + let value = entry.load(context, location, value, inner_ty)?; + DropOverridesMeta::invoke_override( + context, registry, module, &entry, &entry, location, metadata, &info.ty, value, + )?; - entry.append_operation(ReallocBindingsMeta::free(context, value, location)?); entry.append_operation(func::r#return(&[], location)); Ok(region) } diff --git a/src/types/nullable.rs b/src/types/nullable.rs index c340ccb3d..eebf8e45d 100644 --- a/src/types/nullable.rs +++ b/src/types/nullable.rs @@ -10,7 +10,7 @@ use crate::{ error::Result, metadata::{ drop_overrides::DropOverridesMeta, dup_overrides::DupOverridesMeta, - realloc_bindings::ReallocBindingsMeta, MetadataStorage, + runtime_bindings::RuntimeBindingsMeta, MetadataStorage, }, utils::ProgramRegistryExt, }; @@ -49,10 +49,12 @@ pub fn build<'ctx>( metadata, info.self_ty(), |metadata| { - // There's no need to build the type here because it'll always be built within - // `build_dup`. - - Ok(Some(build_dup(context, module, registry, metadata, &info)?)) + registry.build_type(context, module, metadata, &info.ty)?; + if DupOverridesMeta::is_overriden(metadata, &info.ty) { + Ok(Some(build_dup(context, module, registry, metadata, &info)?)) + } else { + Ok(None) + } }, )?; DropOverridesMeta::register_with( @@ -62,12 +64,14 @@ pub fn build<'ctx>( metadata, info.self_ty(), |metadata| { - // There's no need to build the type here because it'll always be built within - // `build_drop`. - - Ok(Some(build_drop( - context, module, registry, metadata, &info, - )?)) + registry.build_type(context, module, metadata, &info.ty)?; + if DropOverridesMeta::is_overriden(metadata, &info.ty) { + Ok(Some(build_drop( + context, module, registry, metadata, &info, + )?)) + } else { + Ok(None) + } }, )?; @@ -83,12 +87,11 @@ fn build_dup<'ctx>( info: &WithSelf, ) -> Result> { let location = Location::unknown(context); - if metadata.get::().is_none() { - metadata.insert(ReallocBindingsMeta::new(context, module)); - } let inner_ty = registry.get_type(&info.ty)?; - let inner_len = inner_ty.layout(registry)?.pad_to_align().size(); + let inner_layout = inner_ty.layout(registry)?; + let inner_len = inner_layout.size(); + let inner_align = inner_layout.align(); let inner_ty = inner_ty.build(context, module, registry, metadata, &info.ty)?; let region = Region::new(); @@ -96,7 +99,6 @@ fn build_dup<'ctx>( let null_ptr = entry.append_op_result(llvm::zero(llvm::r#type::pointer(context, 0), location))?; - let inner_len_val = entry.const_int(context, location, inner_len, 64)?; let src_value = entry.arg(0)?; let src_is_null = entry.append_op_result( @@ -125,41 +127,33 @@ fn build_dup<'ctx>( )); { - let dst_value = block_realloc.append_op_result(ReallocBindingsMeta::realloc( + // build_dup is only registered when the inner type has a dup override. + let size_val = block_realloc.const_int(context, location, inner_len, 64)?; + let align_val = block_realloc.const_int(context, location, inner_align, 64)?; + let rtb = metadata.get_or_insert_with(RuntimeBindingsMeta::default); + let dst_value = rtb.box_alloc( context, - null_ptr, - inner_len_val, + module, + &block_realloc, location, - )?)?; + size_val, + align_val, + )?; - if DupOverridesMeta::is_overriden(metadata, &info.ty) { - let value = block_realloc.load(context, location, src_value, inner_ty)?; - let values = DupOverridesMeta::invoke_override( - context, - registry, - module, - &block_realloc, - &block_realloc, - location, - metadata, - &info.ty, - value, - )?; - block_realloc.store(context, location, src_value, values.0)?; - block_realloc.store(context, location, dst_value, values.1)?; - } else { - block_realloc.append_operation( - ods::llvm::intr_memcpy_inline( - context, - dst_value, - src_value, - IntegerAttribute::new(IntegerType::new(context, 64).into(), inner_len as i64), - IntegerAttribute::new(IntegerType::new(context, 1).into(), 0), - location, - ) - .into(), - ); - } + let value = block_realloc.load(context, location, src_value, inner_ty)?; + let values = DupOverridesMeta::invoke_override( + context, + registry, + module, + &block_realloc, + &block_realloc, + location, + metadata, + &info.ty, + value, + )?; + block_realloc.store(context, location, src_value, values.0)?; + block_realloc.store(context, location, dst_value, values.1)?; block_realloc.append_operation(cf::br(&block_finish, &[dst_value], location)); } @@ -176,10 +170,8 @@ fn build_drop<'ctx>( info: &WithSelf, ) -> Result> { let location = Location::unknown(context); - if metadata.get::().is_none() { - metadata.insert(ReallocBindingsMeta::new(context, module)); - } + // build_drop is only registered when the inner type has a drop override. let inner_ty = registry.build_type(context, module, metadata, &info.ty)?; let region = Region::new(); @@ -201,37 +193,33 @@ fn build_drop<'ctx>( .into(), )?; - let block_free = region.append_block(Block::new(&[])); - let block_finish = - region.append_block(Block::new(&[(llvm::r#type::pointer(context, 0), location)])); + let block_drop = region.append_block(Block::new(&[])); + let block_finish = region.append_block(Block::new(&[])); entry.append_operation(cf::cond_br( context, is_null, &block_finish, - &block_free, - &[null_ptr], + &block_drop, + &[], &[], location, )); { - if DropOverridesMeta::is_overriden(metadata, &info.ty) { - let value = block_free.load(context, location, value, inner_ty)?; - DropOverridesMeta::invoke_override( - context, - registry, - module, - &block_free, - &block_free, - location, - metadata, - &info.ty, - value, - )?; - } - - block_free.append_operation(ReallocBindingsMeta::free(context, value, location)?); - block_free.append_operation(func::r#return(&[], location)); + // No free: the pointer lives in the arena and is reclaimed at invocation end. + let inner_value = block_drop.load(context, location, value, inner_ty)?; + DropOverridesMeta::invoke_override( + context, + registry, + module, + &block_drop, + &block_drop, + location, + metadata, + &info.ty, + inner_value, + )?; + block_drop.append_operation(cf::br(&block_finish, &[], location)); } block_finish.append_operation(func::r#return(&[], location)); diff --git a/src/values.rs b/src/values.rs index 644f1e9da..20d59c424 100644 --- a/src/values.rs +++ b/src/values.rs @@ -707,13 +707,7 @@ impl Value { } CoreTypeConcrete::Box(info) => { let inner = *ptr.cast::>().as_ptr(); - let value = Self::from_ptr(inner, &info.ty, registry, should_drop)?; - - if should_drop { - libc_free(inner.as_ptr().cast()); - } - - value + Self::from_ptr(inner, &info.ty, registry, should_drop)? } CoreTypeConcrete::EcPoint(_) => { let data = ptr.cast::<[[u8; 32]; 2]>().as_mut(); @@ -769,18 +763,12 @@ impl Value { if inner_ptr.is_null() { Self::Null } else { - let value = Self::from_ptr( + Self::from_ptr( NonNull::new_unchecked(inner_ptr).cast(), &info.ty, registry, should_drop, - )?; - - if should_drop { - libc_free(inner_ptr.cast()); - } - - value + )? } } CoreTypeConcrete::Uninitialized(_) => { diff --git a/test_data/contracts/box_arena/callee.cairo b/test_data/contracts/box_arena/callee.cairo new file mode 100644 index 000000000..0db7c85bf --- /dev/null +++ b/test_data/contracts/box_arena/callee.cairo @@ -0,0 +1,18 @@ +#[inline(never)] +fn make_box(v: felt252) -> Box { + BoxTrait::new(v) +} + +#[starknet::contract] +mod callee { + use super::make_box; + + #[storage] + struct Storage {} + + #[external(v0)] + fn add_one(ref self: ContractState, x: felt252) -> felt252 { + let boxed: Box = make_box(x + 1); + boxed.unbox() + } +} diff --git a/test_data/contracts/box_arena/caller.cairo b/test_data/contracts/box_arena/caller.cairo new file mode 100644 index 000000000..fd2bc5244 --- /dev/null +++ b/test_data/contracts/box_arena/caller.cairo @@ -0,0 +1,45 @@ +use starknet::ContractAddress; +use starknet::SyscallResultTrait; +use starknet::syscalls::call_contract_syscall; + +#[inline(never)] +fn make_box(v: felt252) -> Box { + BoxTrait::new(v) +} + +#[starknet::contract] +mod caller { + use super::make_box; + use starknet::ContractAddress; + use starknet::SyscallResultTrait; + use starknet::syscalls::call_contract_syscall; + + #[storage] + struct Storage {} + + // Deliberately holds a box LIVE across the nested `call_contract_syscall`. + #[external(v0)] + fn proxy_add_one( + ref self: ContractState, + target: ContractAddress, + selector: felt252, + x: felt252, + ) -> (felt252, felt252) { + // Alloc #1 — lives across the syscall. + let boxed_x: Box = make_box(x); + + // Pass raw `x` to the callee so the syscall itself doesn't consume boxed_x. + let result = call_contract_syscall(target, selector, array![x].span()) + .unwrap_syscall(); + + // After the callee's reset, allocate a fresh box holding the callee's + // return value. Because the arena was reset, this lands at offset 0 — + // exactly where boxed_x was. + let clobber: Box = make_box(*result[0]); + + let recovered = boxed_x.unbox(); + let clobber_val = clobber.unbox(); + + (recovered, clobber_val) + } +} diff --git a/tests/tests/starknet/box_arena.rs b/tests/tests/starknet/box_arena.rs new file mode 100644 index 000000000..7227029e6 --- /dev/null +++ b/tests/tests/starknet/box_arena.rs @@ -0,0 +1,318 @@ +//! Regression test for the per-invocation box arena. +//! +//! A caller contract allocates two `Box` values, then invokes a callee +//! contract via `call_contract_syscall`, then continues using its boxes after +//! the syscall returns. With a shared/global arena (the pre-fix design) the +//! callee's invocation-end reset would free the caller's live boxes; a +//! subsequent allocation in the caller would land on the same slot, producing +//! a use-after-free. With per-invocation arenas swapped in/out by +//! `InvocationGuard`, the caller's boxes survive the nested call intact. +//! +//! The contract-to-contract dispatch is handled by [`MultiContractHandler`], +//! a syscall handler that routes `call_contract` to a registered +//! [`AotContractExecutor`] by contract address. + +use cairo_lang_starknet_classes::{ + casm_contract_class::ENTRY_POINT_COST, contract_class::ContractClass, +}; +use cairo_native::{ + executor::AotContractExecutor, + starknet::{ + ExecutionInfo, ExecutionInfoV2, ExecutionInfoV3, Secp256k1Point, Secp256r1Point, + StarknetSyscallHandler, SyscallResult, U256, + }, + utils::testing::load_contract, + OptLevel, +}; +use starknet_types_core::felt::Felt; +use std::{collections::HashMap, sync::Arc}; + +struct MultiContractHandler { + contracts: HashMap>, +} + +impl MultiContractHandler { + fn new() -> Self { + Self { + contracts: HashMap::new(), + } + } + + fn register(&mut self, address: Felt, executor: AotContractExecutor) { + self.contracts.insert(address, Arc::new(executor)); + } +} + +impl StarknetSyscallHandler for &mut MultiContractHandler { + fn call_contract( + &mut self, + address: Felt, + entry_point_selector: Felt, + calldata: &[Felt], + remaining_gas: &mut u64, + ) -> SyscallResult> { + let executor = self + .contracts + .get(&address) + .unwrap_or_else(|| panic!("no contract registered at address {address}")) + .clone(); + + let result = executor + .run( + entry_point_selector, + calldata, + *remaining_gas, + None, + &mut **self, + ) + .expect("nested contract execution failed"); + + *remaining_gas = result.remaining_gas; + + if result.failure_flag { + return Err(result.return_values); + } + Ok(result.return_values) + } + + fn get_block_hash( + &mut self, + _block_number: u64, + _remaining_gas: &mut u64, + ) -> SyscallResult { + unimplemented!() + } + fn get_execution_info(&mut self, _remaining_gas: &mut u64) -> SyscallResult { + unimplemented!() + } + fn get_execution_info_v2( + &mut self, + _remaining_gas: &mut u64, + ) -> SyscallResult { + unimplemented!() + } + fn get_execution_info_v3( + &mut self, + _remaining_gas: &mut u64, + ) -> SyscallResult { + unimplemented!() + } + fn deploy( + &mut self, + _class_hash: Felt, + _contract_address_salt: Felt, + _calldata: &[Felt], + _deploy_from_zero: bool, + _remaining_gas: &mut u64, + ) -> SyscallResult<(Felt, Vec)> { + unimplemented!() + } + fn replace_class(&mut self, _class_hash: Felt, _remaining_gas: &mut u64) -> SyscallResult<()> { + unimplemented!() + } + fn library_call( + &mut self, + _class_hash: Felt, + _function_selector: Felt, + _calldata: &[Felt], + _remaining_gas: &mut u64, + ) -> SyscallResult> { + unimplemented!() + } + fn storage_read( + &mut self, + _address_domain: u32, + _address: Felt, + _remaining_gas: &mut u64, + ) -> SyscallResult { + unimplemented!() + } + fn storage_write( + &mut self, + _address_domain: u32, + _address: Felt, + _value: Felt, + _remaining_gas: &mut u64, + ) -> SyscallResult<()> { + unimplemented!() + } + fn emit_event( + &mut self, + _keys: &[Felt], + _data: &[Felt], + _remaining_gas: &mut u64, + ) -> SyscallResult<()> { + unimplemented!() + } + fn send_message_to_l1( + &mut self, + _to_address: Felt, + _payload: &[Felt], + _remaining_gas: &mut u64, + ) -> SyscallResult<()> { + unimplemented!() + } + fn keccak(&mut self, _input: &[u64], _remaining_gas: &mut u64) -> SyscallResult { + unimplemented!() + } + fn secp256k1_new( + &mut self, + _x: U256, + _y: U256, + _remaining_gas: &mut u64, + ) -> SyscallResult> { + unimplemented!() + } + fn secp256k1_add( + &mut self, + _p0: Secp256k1Point, + _p1: Secp256k1Point, + _remaining_gas: &mut u64, + ) -> SyscallResult { + unimplemented!() + } + fn secp256k1_mul( + &mut self, + _p: Secp256k1Point, + _m: U256, + _remaining_gas: &mut u64, + ) -> SyscallResult { + unimplemented!() + } + fn secp256k1_get_point_from_x( + &mut self, + _x: U256, + _y_parity: bool, + _remaining_gas: &mut u64, + ) -> SyscallResult> { + unimplemented!() + } + fn secp256k1_get_xy( + &mut self, + _p: Secp256k1Point, + _remaining_gas: &mut u64, + ) -> SyscallResult<(U256, U256)> { + unimplemented!() + } + fn secp256r1_new( + &mut self, + _x: U256, + _y: U256, + _remaining_gas: &mut u64, + ) -> SyscallResult> { + unimplemented!() + } + fn secp256r1_add( + &mut self, + _p0: Secp256r1Point, + _p1: Secp256r1Point, + _remaining_gas: &mut u64, + ) -> SyscallResult { + unimplemented!() + } + fn secp256r1_mul( + &mut self, + _p: Secp256r1Point, + _m: U256, + _remaining_gas: &mut u64, + ) -> SyscallResult { + unimplemented!() + } + fn secp256r1_get_point_from_x( + &mut self, + _x: U256, + _y_parity: bool, + _remaining_gas: &mut u64, + ) -> SyscallResult> { + unimplemented!() + } + fn secp256r1_get_xy( + &mut self, + _p: Secp256r1Point, + _remaining_gas: &mut u64, + ) -> SyscallResult<(U256, U256)> { + unimplemented!() + } + fn sha256_process_block( + &mut self, + _state: &mut [u32; 8], + _block: &[u32; 16], + _remaining_gas: &mut u64, + ) -> SyscallResult<()> { + unimplemented!() + } + fn get_class_hash_at( + &mut self, + _contract_address: Felt, + _remaining_gas: &mut u64, + ) -> SyscallResult { + unimplemented!() + } + fn meta_tx_v0( + &mut self, + _address: Felt, + _entry_point_selector: Felt, + _calldata: &[Felt], + _signature: &[Felt], + _remaining_gas: &mut u64, + ) -> SyscallResult> { + unimplemented!() + } +} + +fn build_executor(contract: &ContractClass) -> (AotContractExecutor, Felt) { + let extracted = contract + .extract_sierra_program(false) + .expect("failed to extract sierra program"); + + let executor = AotContractExecutor::new( + &extracted.program, + &contract.entry_points_by_type, + extracted.sierra_version, + OptLevel::Default, + None, + ) + .expect("failed to build AotContractExecutor"); + + let selector = Felt::from( + &contract + .entry_points_by_type + .external + .first() + .expect("contract should have an external entry point") + .selector, + ); + + (executor, selector) +} + +#[test] +fn test_boxes_stay_valid_across_contract_calls() { + let callee = load_contract("test_data_artifacts/contracts/box_arena/callee.contract.json"); + let caller = load_contract("test_data_artifacts/contracts/box_arena/caller.contract.json"); + + let (callee_executor, add_one_selector) = build_executor(&callee); + let (caller_executor, proxy_selector) = build_executor(&caller); + + let callee_address = Felt::from(0x42); + let mut handler = MultiContractHandler::new(); + handler.register(callee_address, callee_executor); + + let input = Felt::from(5); + let result = caller_executor + .run( + proxy_selector, + &[callee_address, add_one_selector, input], + u64::MAX - ENTRY_POINT_COST as u64, + None, + &mut handler, + ) + .expect("caller contract execution failed"); + + assert!( + !result.failure_flag, + "caller execution panicked: {:?}", + result.error_msg + ); + assert_eq!(result.return_values, vec![Felt::from(5), Felt::from(6)]); +} diff --git a/tests/tests/starknet/mod.rs b/tests/tests/starknet/mod.rs index 9a62585d5..86e370dcf 100644 --- a/tests/tests/starknet/mod.rs +++ b/tests/tests/starknet/mod.rs @@ -1,3 +1,4 @@ +mod box_arena; mod keccak; mod secp256; mod u256;