(optimization): Reboxing now also works on whole var reboxing

Author: eytan-starkware

crates/cairo-lang-lowering/src/optimizations/reboxing.rs

@@ -215,13 +215,23 @@ fn apply_reboxing_candidate<'db>(
         candidate.reboxed_var.index()
     );
 
-    // TODO(eytan-starkware): Handle snapshot of box (e.g., @Box<T>).
-    // Only support MemberOfUnboxed where source is Unboxed for now.
-    let ReboxingValue::MemberOfUnboxed { source, member } = &candidate.source else {
-        // If source is not member of unboxed, we are reboxing original value which is not supported
-        // yet.
-        return;
+    match &candidate.source {
+        ReboxingValue::Revoked => (),
+        ReboxingValue::Unboxed(id) => replace_variable_usages(lowered, candidate.reboxed_var, *id),
+        ReboxingValue::MemberOfUnboxed { source, member } => {
+            replace_into_box_call(db, lowered, candidate, source, *member)
+        }
     };
+}
+
+/// Replaces the call to `into_box` with a call to `struct_boxed_deconstruct`.
+fn replace_into_box_call<'db>(
+    db: &'db dyn Database,
+    lowered: &mut Lowered<'db>,
+    candidate: &ReboxCandidate,
+    source: &Rc<ReboxingValue>,
+    member: usize,
+) {
     let ReboxingValue::Unboxed(source_var) = **source else {
         // When source of the value is not `Unboxes`, it is a nested MemberOfUnboxed, which is not
         // supported yet.
@@ -232,7 +242,7 @@ fn apply_reboxing_candidate<'db>(
         db,
         &mut lowered.variables,
         source_var,
-        *member,
+        member,
         candidate.reboxed_var,
         &lowered.blocks[candidate.into_box_location.0].statements[candidate.into_box_location.1],
     ) {
@@ -244,6 +254,59 @@ fn apply_reboxing_candidate<'db>(
     }
 }
 
+/// Replaces all usages of old_var with new_var throughout the lowered function.
+fn replace_variable_usages(lowered: &mut Lowered<'_>, old_var: VariableId, new_var: VariableId) {
+    if old_var == new_var {
+        return;
+    }
+
+    for block in lowered.blocks.iter_mut() {
+        for stmt in &mut block.statements {
+            for input in stmt.inputs_mut() {
+                if input.var_id == old_var {
+                    input.var_id = new_var;
+                }
+            }
+        }
+
+        match &mut block.end {
+            BlockEnd::Return(var_usages, _) => {
+                for return_var in var_usages.iter_mut() {
+                    if return_var.var_id == old_var {
+                        return_var.var_id = new_var;
+                    }
+                }
+            }
+            BlockEnd::Goto(_, remapping) => {
+                // Collect all changes needed
+                let mut changes = Vec::new();
+                for (dst, src_usage) in remapping.iter() {
+                    if *dst == old_var {
+                        changes.push((*dst, Some(new_var), *src_usage));
+                    } else if src_usage.var_id == old_var {
+                        let mut new_usage = *src_usage;
+                        new_usage.var_id = new_var;
+                        changes.push((*dst, None, new_usage));
+                    }
+                }
+
+                // Apply changes
+                for (old_key, new_key_opt, new_usage) in changes {
+                    if let Some(new_key) = new_key_opt {
+                        // Key needs to change
+                        remapping.swap_remove(&old_key);
+                        remapping.insert(new_key, new_usage);
+                    } else {
+                        // Only value needs to change
+                        remapping.insert(old_key, new_usage);
+                    }
+                }
+            }
+            _ => {}
+        }
+    }
+}
+
 /// Creates a struct_boxed_deconstruct call statement.
 /// Returns None if the call cannot be created.
 fn create_struct_boxed_deconstruct_call<'db>(

crates/cairo-lang-lowering/src/optimizations/test_data/reboxing

@@ -229,8 +229,6 @@ test_reboxing_analysis
 //! > function_name
 rebox_whole
 
-//! > TODO(eytan-starkware): Add support for whole var reboxing
-
 //! > module_code
 struct Simple {
     value: felt252,
@@ -264,7 +262,7 @@ Statements:
   (v1: test::Simple) <- core::box::unbox::<test::Simple>(v0)
   (v2: core::box::Box::<test::Simple>) <- core::box::into_box::<test::Simple>(v1)
 End:
-  Return(v2)
+  Return(v0)
 
 //! > ==========================================================================