diff --git a/crates/cairo-lang-lowering/src/optimizations/reboxing.rs b/crates/cairo-lang-lowering/src/optimizations/reboxing.rs index 1441c6689cb..1a59a03d151 100644 --- a/crates/cairo-lang-lowering/src/optimizations/reboxing.rs +++ b/crates/cairo-lang-lowering/src/optimizations/reboxing.rs @@ -7,10 +7,11 @@ use std::rc::Rc; use cairo_lang_filesystem::flag::flag_future_sierra; use cairo_lang_semantic::helper::ModuleHelper; use cairo_lang_semantic::items::structure::StructSemantic; -use cairo_lang_semantic::types::{TypesSemantic, peel_snapshots}; +use cairo_lang_semantic::types::{TypesSemantic, peel_snapshots, wrap_in_snapshots}; use cairo_lang_semantic::{ConcreteTypeId, GenericArgumentId, TypeLongId}; use cairo_lang_utils::ordered_hash_map::{Entry, OrderedHashMap}; use cairo_lang_utils::ordered_hash_set::OrderedHashSet; +use itertools::Itertools; use salsa::Database; use crate::borrow_check::analysis::StatementLocation; @@ -257,22 +258,17 @@ fn create_struct_boxed_deconstruct_call<'db>( } let (n_snapshots, struct_ty) = peel_snapshots(db, *inner_ty); - // TODO(eytan-starkware): Support snapshots of structs in reboxing optimization. - // Currently we give up if the struct is wrapped in snapshots. - if n_snapshots > 0 { - trace!("Skipping reboxing for snapshotted struct (n_snapshots={})", n_snapshots); - return None; - } - - // Extract member types from struct or tuple + trace!("Extracted struct or tuple type: {:?}", struct_ty); let member_types = match struct_ty { TypeLongId::Concrete(ConcreteTypeId::Struct(struct_id)) => db .concrete_struct_members(struct_id) .ok()? .iter() - .map(|(_, member)| member.ty) - .collect::>(), - TypeLongId::Tuple(inner_types) => inner_types, + .map(|(_, member)| wrap_in_snapshots(db, member.ty, n_snapshots)) + .collect_vec(), + TypeLongId::Tuple(inner_types) => { + inner_types.into_iter().map(|ty| wrap_in_snapshots(db, ty, n_snapshots)).collect() + } _ => { trace!("Unsupported type for reboxing: {:?}", struct_ty); return None; diff --git a/crates/cairo-lang-lowering/src/optimizations/test_data/reboxing b/crates/cairo-lang-lowering/src/optimizations/test_data/reboxing index cd7da13cd44..a8e712cfd18 100644 --- a/crates/cairo-lang-lowering/src/optimizations/test_data/reboxing +++ b/crates/cairo-lang-lowering/src/optimizations/test_data/reboxing @@ -215,7 +215,7 @@ blk0 (root): Statements: (v1: @test::Data) <- core::box::unbox::<@test::Data>(v0) (v2: @test::NonCopy, v3: @core::felt252) <- struct_destructure(v1) - (v4: core::box::Box::<@test::NonCopy>) <- core::box::into_box::<@test::NonCopy>(v2) + (v4: core::box::Box::<@test::NonCopy>, v5: core::box::Box::<@core::felt252>) <- struct_destructure(v0) End: Return(v4) @@ -690,3 +690,64 @@ Statements: (v5: (core::box::Box::, test::NonDrop)) <- struct_construct(v4, v3) End: Return(v5) + +//! > ========================================================================== + +//! > Test reboxing with snapshot of non-drop struct and member + +//! > test_runner_name +test_reboxing_analysis + +//! > function_name +main + +//! > module_code +use core::box::BoxTrait; + +#[derive(Copy)] +struct NonDrop { + b: felt252, +} + +#[derive(Copy)] +struct A { + a: felt252, + non_drop: NonDrop, +} + +//! > function_code +fn main(a: Box<@A>) -> (Box<@felt252>, @NonDrop) { + let a = a.unbox(); + (BoxTrait::new(a.a), a.non_drop) +} + +//! > semantic_diagnostics + +//! > lowering_diagnostics + +//! > candidates +v4 + +//! > before +Parameters: v0: core::box::Box::<@test::A> +blk0 (root): +Statements: + (v1: @test::A) <- core::box::unbox::<@test::A>(v0) + (v2: @core::felt252, v3: @test::NonDrop) <- struct_destructure(v1) + (v4: core::box::Box::<@core::felt252>) <- core::box::into_box::<@core::felt252>(v2) + (v5: @core::felt252, v6: @test::NonDrop) <- struct_destructure(v1) + (v7: (core::box::Box::<@core::felt252>, @test::NonDrop)) <- struct_construct(v4, v6) +End: + Return(v7) + +//! > after +Parameters: v0: core::box::Box::<@test::A> +blk0 (root): +Statements: + (v1: @test::A) <- core::box::unbox::<@test::A>(v0) + (v2: @core::felt252, v3: @test::NonDrop) <- struct_destructure(v1) + (v4: core::box::Box::<@core::felt252>, v8: core::box::Box::<@test::NonDrop>) <- struct_destructure(v0) + (v5: @core::felt252, v6: @test::NonDrop) <- struct_destructure(v1) + (v7: (core::box::Box::<@core::felt252>, @test::NonDrop)) <- struct_construct(v4, v6) +End: + Return(v7) diff --git a/crates/cairo-lang-sierra-generator/src/function_generator_test_data/struct b/crates/cairo-lang-sierra-generator/src/function_generator_test_data/struct index aaed758a173..5f0c512aed6 100644 --- a/crates/cairo-lang-sierra-generator/src/function_generator_test_data/struct +++ b/crates/cairo-lang-sierra-generator/src/function_generator_test_data/struct @@ -118,16 +118,15 @@ struct A { //! > lowering_diagnostics -//! > TODO(eytan-starkware): Rebox on non-copy structs. - //! > sierra_code label_test::foo::0: -unbox>([0]) -> ([1]) -store_temp>([1]) -> ([1]) -struct_snapshot_deconstruct([1]) -> ([2], [3]) -drop>>([3]) -> () -into_box>>([2]) -> ([4]) -return([4]) +dup>>([0]) -> ([0], [1]) +unbox>([1]) -> ([2]) +drop>([2]) -> () +struct_boxed_deconstruct>([0]) -> ([3], [4]) +drop>>>([4]) -> () +store_temp>>>([3]) -> ([3]) +return([3]) //! > ========================================================================== @@ -155,15 +154,48 @@ struct A { //! > lowering_diagnostics -//! > TODO(eytan-starkware): We want reboxing to apply to sierra in the future, +//! > sierra_code +label_test::foo::0: +dup>>([0]) -> ([0], [1]) +unbox>([1]) -> ([2]) +drop>([2]) -> () +struct_boxed_deconstruct>([0]) -> ([3], [4]) +drop>>>([4]) -> () +store_temp>>>([3]) -> ([3]) +return([3]) + +//! > ========================================================================== + +//! > Test reboxing of a boxed, repeatedly snapshotted struct. + +//! > test_runner_name +test_function_generator(future_sierra:true) -//! > so we will see struct_boxed_deconstruct. +//! > function_code +fn foo(box: Box<@@@A>) -> Box<@@@Array> { + BoxTrait::new(box.unbox().a) +} + +//! > function_name +foo + +//! > module_code +#[derive(Drop)] +struct A { + a: Array, + b: Array, +} + +//! > semantic_diagnostics + +//! > lowering_diagnostics //! > sierra_code label_test::foo::0: -unbox>([0]) -> ([1]) -store_temp>([1]) -> ([1]) -struct_snapshot_deconstruct([1]) -> ([2], [3]) -drop>>([3]) -> () -into_box>>([2]) -> ([4]) -return([4]) +dup>>([0]) -> ([0], [1]) +unbox>([1]) -> ([2]) +drop>([2]) -> () +struct_boxed_deconstruct>([0]) -> ([3], [4]) +drop>>>([4]) -> () +store_temp>>>([3]) -> ([3]) +return([3]) diff --git a/crates/cairo-lang-sierra-generator/src/utils.rs b/crates/cairo-lang-sierra-generator/src/utils.rs index b1c0f6e2b5d..0690bafa25a 100644 --- a/crates/cairo-lang-sierra-generator/src/utils.rs +++ b/crates/cairo-lang-sierra-generator/src/utils.rs @@ -112,9 +112,8 @@ pub fn struct_deconstruct_libfunc_id( let is_snapshot = long_id.generic_id == SnapshotType::id(); let is_box = long_id.generic_id == BoxType::id(); Ok(if is_snapshot { - let concrete_enum_type = - extract_matches!(&long_id.generic_args[0], GenericArg::Type).clone(); - get_libfunc_id_with_generic_arg(db, "struct_snapshot_deconstruct", concrete_enum_type) + let concrete_type = extract_matches!(&long_id.generic_args[0], GenericArg::Type).clone(); + get_libfunc_id_with_generic_arg(db, "struct_snapshot_deconstruct", concrete_type) } else if is_box { let inner_ty = extract_matches!(&long_id.generic_args[0], GenericArg::Type).clone(); get_libfunc_id_with_generic_arg(db, "struct_boxed_deconstruct", inner_ty)