(optimization): Reboxing also applied on Box of snapshots

Author: eytan-starkware

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::<Vec<_>>(),
-        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;

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::<core::felt252>, 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)

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<Snapshot<test::A>>([0]) -> ([1])
-store_temp<Snapshot<test::A>>([1]) -> ([1])
-struct_snapshot_deconstruct<test::A>([1]) -> ([2], [3])
-drop<Snapshot<Array<felt252>>>([3]) -> ()
-into_box<Snapshot<Array<felt252>>>([2]) -> ([4])
-return([4])
+dup<Box<Snapshot<test::A>>>([0]) -> ([0], [1])
+unbox<Snapshot<test::A>>([1]) -> ([2])
+drop<Snapshot<test::A>>([2]) -> ()
+struct_boxed_deconstruct<Snapshot<test::A>>([0]) -> ([3], [4])
+drop<Box<Snapshot<Array<felt252>>>>([4]) -> ()
+store_temp<Box<Snapshot<Array<felt252>>>>([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<Box<Snapshot<test::A>>>([0]) -> ([0], [1])
+unbox<Snapshot<test::A>>([1]) -> ([2])
+drop<Snapshot<test::A>>([2]) -> ()
+struct_boxed_deconstruct<Snapshot<test::A>>([0]) -> ([3], [4])
+drop<Box<Snapshot<Array<felt252>>>>([4]) -> ()
+store_temp<Box<Snapshot<Array<felt252>>>>([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<felt252>> {
+    BoxTrait::new(box.unbox().a)
+}
+
+//! > function_name
+foo
+
+//! > module_code
+#[derive(Drop)]
+struct A {
+    a: Array<felt252>,
+    b: Array<felt252>,
+}
+
+//! > semantic_diagnostics
+
+//! > lowering_diagnostics
 
 //! > sierra_code
 label_test::foo::0:
-unbox<Snapshot<test::A>>([0]) -> ([1])
-store_temp<Snapshot<test::A>>([1]) -> ([1])
-struct_snapshot_deconstruct<test::A>([1]) -> ([2], [3])
-drop<Snapshot<Array<felt252>>>([3]) -> ()
-into_box<Snapshot<Array<felt252>>>([2]) -> ([4])
-return([4])
+dup<Box<Snapshot<test::A>>>([0]) -> ([0], [1])
+unbox<Snapshot<test::A>>([1]) -> ([2])
+drop<Snapshot<test::A>>([2]) -> ()
+struct_boxed_deconstruct<Snapshot<test::A>>([0]) -> ([3], [4])
+drop<Box<Snapshot<Array<felt252>>>>([4]) -> ()
+store_temp<Box<Snapshot<Array<felt252>>>>([3]) -> ([3])
+return([3])

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)