(optimization): automatically detecting reboxing of struct member and applying struct_box_deconstruct libfunc

[eytan-starkware]

TL;DR

Add reboxing optimization to eliminate unnecessary box/unbox operations when accessing members of boxed structs.

Currently does not support reboxing of full struct nor Snaphots inside or outside box

Using a new flag (with no external cli at the moment) to block the optimization from being used.

What changed?

Added a new optimization phase Reboxing that identifies and optimizes patterns where code unboxes a struct, accesses a member, and then boxes that member again:

• Implemented analysis to find reboxing candidates in lowered functions
• Created transformation to replace these patterns with direct struct_boxed_deconstruct libfunc calls
• Added comprehensive test cases for various struct types and scenarios
• Integrated the optimization into the baseline optimization strategy
• Added tracing/logging support for the optimization module

[reviewable-StarkWare]

This change is 

[eytan-starkware]

• (optimization): automatically detecting reboxing of struct member and applying struct_box_deconstruct libfunc #8768 👈 (View in Graphite)
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[orizi]

Reviewable status: 0 of 12 files reviewed, 2 unresolved discussions (waiting on @giladchase and @ilyalesokhin-starkware)

---

crates/cairo-lang-filesystem/src/flag.rs line 25 at r2 (raw file):

    /// Whether to use future_sierra in the generated code.
    ///
    /// Default is false as it makes panic unprovable.

WDYM?

---

crates/cairo-lang-filesystem/src/flag.rs line 34 at r2 (raw file):

    if let Some(flag) = db.get_flag(flag) { *flag == Flag::UnsafePanic(true) } else { false }
}

doc

[ilyalesokhin-starkware]

@ilyalesokhin-starkware reviewed 1 of 11 files at r1.
Reviewable status: 1 of 12 files reviewed, 3 unresolved discussions (waiting on @eytan-starkware and @giladchase)

---

crates/cairo-lang-lowering/src/optimizations/reboxing.rs line 23 at r2 (raw file):

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum ReboxingValue {
    Nothing,

why do you need this, can't you just avoid adding the variable to the current_state hashmap?

Code quote:

Nothing,

[ilyalesokhin-starkware]

Reviewable status: 1 of 12 files reviewed, 4 unresolved discussions (waiting on @eytan-starkware and @giladchase)

---

crates/cairo-lang-lowering/src/optimizations/test_data/reboxing line 503 at r2 (raw file):

    let unboxed = bs.unbox();
    let val = unboxed.val;
    BoxTrait::new(val)

Add a test where two members are reboxed, and the boxed variable is not copyable.

Suggestion:

    let val = unboxed.val;
    BoxTrait::new(val)
    BoxTrait::new(arr)

[eytan-starkware]

Reviewable status: 1 of 12 files reviewed, 4 unresolved discussions (waiting on @giladchase, @ilyalesokhin-starkware, and @orizi)

---

crates/cairo-lang-filesystem/src/flag.rs line 25 at r2 (raw file):

Previously, orizi wrote…

WDYM?

AI/me slop 🤦‍♂️
Done

---

crates/cairo-lang-filesystem/src/flag.rs line 34 at r2 (raw file):

Previously, orizi wrote…

doc

Done

---

crates/cairo-lang-lowering/src/optimizations/reboxing.rs line 23 at r2 (raw file):

Previously, ilyalesokhin-starkware wrote…

why do you need this, can't you just avoid adding the variable to the current_state hashmap?

This is necessary for correct meet/join. Consider a situation where 3 branches are doing a meet. The first two had a meet that resulted with Nothing, if I forget then"Nothing", then the 3rd branch will pass on it's value even though it is incorrect to do so.

---

crates/cairo-lang-lowering/src/optimizations/test_data/reboxing line 503 at r2 (raw file):

Previously, ilyalesokhin-starkware wrote…

Add a test where two members are reboxed, and the boxed variable is not copyable.

If boxed variable is not coyable we do not apply the optimization which is what the test checks on a copyable member, so I dont think we should add another access here.

[ilyalesokhin-starkware]

crates/cairo-lang-lowering/src/optimizations/reboxing.rs line 23 at r2 (raw file):

Previously, eytan-starkware wrote…

This is necessary for correct meet/join. Consider a situation where 3 branches are doing a meet. The first two had a meet that resulted with Nothing, if I forget then"Nothing", then the 3rd branch will pass on it's value even though it is incorrect to do so.

Your code does not handle any joins.

[ilyalesokhin-starkware]

Reviewable status: 1 of 12 files reviewed, 4 unresolved discussions (waiting on @eytan-starkware, @giladchase, and @orizi)

---

crates/cairo-lang-lowering/src/optimizations/reboxing.rs line 23 at r2 (raw file):

Previously, ilyalesokhin-starkware wrote…

Your code does not handle any joins.

Is there a joint that you can support?

I think the remapped var should never have ReboxingValue state.

[orizi]

@orizi reviewed all commit messages.
Reviewable status: 1 of 12 files reviewed, 7 unresolved discussions (waiting on @eytan-starkware, @giladchase, and @ilyalesokhin-starkware)

---

crates/cairo-lang-lowering/src/optimizations/mod.rs line 2 at r3 (raw file):

/// Macro for debug logging with "optimization" target.
macro_rules! debug {

separate to another PR.

---

crates/cairo-lang-lowering/src/optimizations/strategy.rs line 89 at r3 (raw file):

            OptimizationPhase::SubStrategy { .. } => "SubStrategy",
        }
    }

---

crates/cairo-lang-lowering/src/optimizations/strategy.rs line 100 at r3 (raw file):

        lowered: &mut Lowered<'db>,
    ) -> Maybe<()> {
        debug!("Applying optimization: {}", self.name());

Suggestion:

        debug!("Applying optimization: {self:?}");

[eytan-starkware]

Reviewable status: 0 of 12 files reviewed, 7 unresolved discussions (waiting on @giladchase, @ilyalesokhin-starkware, and @orizi)

---

crates/cairo-lang-lowering/src/optimizations/mod.rs line 2 at r3 (raw file):

Previously, orizi wrote…

separate to another PR.

Done

---

crates/cairo-lang-lowering/src/optimizations/reboxing.rs line 23 at r2 (raw file):

Previously, ilyalesokhin-starkware wrote…

Is there a joint that you can support?

I think the remapped var should never have ReboxingValue state.

Oh the relevant part was deleted in migration from the dataflow approach, good catach!
Currently it should never error, but in the future I expect to follow when the vars are essentially the same value being passed around. So I think it is best to keep it.
What used to be:

Code snippet:

let source = current_state
                                .get(&call_stmt.inputs[0].var_id)
                                .unwrap_or(&ReboxingValue::Nothing);
                            if matches!(source, ReboxingValue::Nothing) {
                                continue;
                            }
                            
And a new remapper follower should now handle the meet essentially .

---

crates/cairo-lang-lowering/src/optimizations/strategy.rs line 89 at r3 (raw file):

            OptimizationPhase::SubStrategy { .. } => "SubStrategy",
        }
    }

Done.

---

crates/cairo-lang-lowering/src/optimizations/strategy.rs line 100 at r3 (raw file):

        lowered: &mut Lowered<'db>,
    ) -> Maybe<()> {
        debug!("Applying optimization: {}", self.name());

Done.

[ilyalesokhin-starkware]

crates/cairo-lang-lowering/src/optimizations/strategy.rs line 92 at r4 (raw file):

                if flag_future_sierra(db) {
                    apply_reboxing(db, lowered);
                }

Move the if inside apply_reboxing

Suggestion:

                apply_reboxing(db, lowered);
                

[orizi]

@orizi reviewed 1 of 11 files at r1, 1 of 1 files at r3, 2 of 7 files at r4.
Reviewable status: 2 of 13 files reviewed, 9 unresolved discussions (waiting on @eytan-starkware, @giladchase, and @ilyalesokhin-starkware)

---

crates/cairo-lang-lowering/src/lib.rs line 10 at r5 (raw file):

pub mod db;
pub mod destructs;

Suggestion:

pub mod concretize;
pub mod db;
pub mod destructs;

[orizi]

@orizi reviewed 1 of 7 files at r4, 4 of 5 files at r6, all commit messages.
Reviewable status: 7 of 13 files reviewed, 10 unresolved discussions (waiting on @eytan-starkware, @giladchase, and @ilyalesokhin-starkware)

---

crates/cairo-lang-lowering/src/optimizations/strategy.rs line 90 at r6 (raw file):

            OptimizationPhase::OptimizeRemappings => optimize_remappings(lowered),
            OptimizationPhase::Reboxing => {
                if flag_future_sierra(db) {

move this into the function.

Code quote:

                if flag_future_sierra(db) {

---

crates/cairo-lang-lowering/src/optimizations/reboxing.rs line 21 at r6 (raw file):

};

#[derive(Debug, Clone, PartialEq, Eq, Hash)]

doc enum and variants.

---

crates/cairo-lang-lowering/src/optimizations/reboxing.rs line 47 at r6 (raw file):

    pub reboxed_var: VariableId,
    /// Location where into_box call occurs (block_id, stmt_idx)
    pub into_box_location: (BlockId, usize),

Suggestion:

    /// Location where into_box call occurs (block_id, stmt_idx)
    pub into_box_location: StatementLocation,

---

crates/cairo-lang-lowering/src/optimizations/reboxing.rs line 79 at r6 (raw file):

    let mut current_state: OrderedHashMap<VariableId, ReboxingValue> = OrderedHashMap::default();
    let mut candidates: OrderedHashSet<ReboxCandidate> = OrderedHashSet::default();

Suggestion:

    let mut current_state = OrderedHashMap::<VariableId, ReboxingValue>::default();
    let mut candidates = OrderedHashSet::<ReboxCandidate>::default();

[eytan-starkware]

Reviewable status: 7 of 13 files reviewed, 10 unresolved discussions (waiting on @giladchase, @ilyalesokhin-starkware, and @orizi)

---

crates/cairo-lang-lowering/src/optimizations/reboxing.rs line 149 at r4 (raw file):

Previously, ilyalesokhin-starkware wrote…

use current_state.entry

Done.

---

crates/cairo-lang-lowering/src/optimizations/reboxing.rs line 257 at r4 (raw file):

Previously, ilyalesokhin-starkware wrote…

Is this the same thing?

Why yes it is

---

crates/cairo-lang-lowering/src/optimizations/reboxing.rs line 21 at r6 (raw file):

Previously, orizi wrote…

doc enum and variants.

Done.

---

crates/cairo-lang-lowering/src/optimizations/strategy.rs line 92 at r4 (raw file):

Previously, ilyalesokhin-starkware wrote…

Move the if inside apply_reboxing

Done.

---

crates/cairo-lang-lowering/src/optimizations/strategy.rs line 90 at r6 (raw file):

Previously, orizi wrote…

move this into the function.

Done.

---

crates/cairo-lang-lowering/src/optimizations/reboxing.rs line 47 at r6 (raw file):

    pub reboxed_var: VariableId,
    /// Location where into_box call occurs (block_id, stmt_idx)
    pub into_box_location: (BlockId, usize),

Done.

---

crates/cairo-lang-lowering/src/optimizations/reboxing.rs line 79 at r6 (raw file):

    let mut current_state: OrderedHashMap<VariableId, ReboxingValue> = OrderedHashMap::default();
    let mut candidates: OrderedHashSet<ReboxCandidate> = OrderedHashSet::default();

Done.

[eytan-starkware]

Reviewable status: 5 of 13 files reviewed, 10 unresolved discussions (waiting on @giladchase, @ilyalesokhin-starkware, and @orizi)

---

crates/cairo-lang-lowering/src/lib.rs line 10 at r5 (raw file):

pub mod db;
pub mod destructs;

Done.

[ilyalesokhin-starkware]

@ilyalesokhin-starkware reviewed all commit messages.
Reviewable status: 4 of 13 files reviewed, 10 unresolved discussions (waiting on @eytan-starkware, @giladchase, and @orizi)

---

crates/cairo-lang-lowering/src/optimizations/reboxing.rs line 132 at r7 (raw file):

                                };

                                update_reboxing_variable_join(&mut current_state, *output_var, res);

only BlockEnd::Goto(_, remapping) can trigger a join.

Code quote:

update_reboxing_variable_join(&mut current_state, *output_var, res);

[ilyalesokhin-starkware]

crates/cairo-lang-lowering/src/optimizations/reboxing.rs line 149 at r4 (raw file):

Previously, eytan-starkware wrote…

Done.

not done

[orizi]

@orizi reviewed 4 of 5 files at r7.
Reviewable status: 8 of 13 files reviewed, 7 unresolved discussions (waiting on @array, @eytan-starkware, @giladchase, and @ilyalesokhin-starkware)

---

crates/cairo-lang-sierra-generator/src/function_generator_test_data/struct line 63 at r7 (raw file):

drop<Box<u128>>([3]) -> ()
store_temp<Box<u128>>([4]) -> ([4])
return([4])

add as well:

fn foo(a: @box) -> Box<@array> {
BoxTrait::new(a.unbox().b)
}

//! > module_code
struct A {
a: Array,
b: Array,
}

---

crates/cairo-lang-sierra-generator/src/utils.rs line 117 at r7 (raw file):

        let concrete_enum_type =
            extract_matches!(&long_id.generic_args[0], GenericArg::Type).clone();
        // TODO: Handle snapshot of box (e.g., @Box<T>).

should probably just be handled in the lowering alg - so move to there.

Code quote:

        // TODO: Handle snapshot of box (e.g., @Box<T>).

[eytan-starkware]

Reviewable status: 8 of 13 files reviewed, 7 unresolved discussions (waiting on @array, @giladchase, @ilyalesokhin-starkware, and @orizi)

---

crates/cairo-lang-lowering/src/optimizations/reboxing.rs line 149 at r4 (raw file):

Previously, ilyalesokhin-starkware wrote…

not done

Now done

---

crates/cairo-lang-lowering/src/optimizations/reboxing.rs line 132 at r7 (raw file):

Previously, ilyalesokhin-starkware wrote…

only BlockEnd::Goto(_, remapping) can trigger a join.

Turned back to an insert

---

crates/cairo-lang-sierra-generator/src/utils.rs line 117 at r7 (raw file):

Previously, orizi wrote…

should probably just be handled in the lowering alg - so move to there.

Done.

---

crates/cairo-lang-sierra-generator/src/function_generator_test_data/struct line 63 at r7 (raw file):

Previously, orizi wrote…

add as well:

fn foo(a: @box) -> Box<@array> {
BoxTrait::new(a.unbox().b)
}

//! > module_code
struct A {
a: Array,
b: Array,
}

Done
Currently not supported, but added a todo

[ilyalesokhin-starkware]

crates/cairo-lang-lowering/src/optimizations/reboxing.rs line 166 at r8 (raw file):

        }
        Entry::Occupied(_) => { // No changes needed }
        }

Suggestion:

        Entry::Occupied(mut entry) => {
            if entry.get() != &res
                entry.insert(ReboxingValue::Revoked);
            }
        }

[ilyalesokhin-starkware]

crates/cairo-lang-lowering/src/optimizations/reboxing.rs line 267 at r8 (raw file):

    let generic_args = concrete_box.generic_args(db);
    // TODO: Why isnt this a box??????

?

Code quote:

 // TODO: Why isnt this a box??????

[ilyalesokhin-starkware]

crates/cairo-lang-lowering/src/optimizations/reboxing.rs line 322 at r8 (raw file):

            let box_ty = cairo_lang_semantic::corelib::core_box_ty(db, member_ty);
            let out_location = variables[output_var].location;
            let var = variables.alloc(Variable::with_default_context(db, box_ty, out_location));

Is it possible that this var is not droppable?

Code quote:

  let var = variables.alloc(Variable::with_default_context(db, box_ty, out_location));

[ilyalesokhin-starkware]

crates/cairo-lang-lowering/src/optimizations/test_data/reboxing line 23 at r8 (raw file):

    let b = ba.b;
    BoxTrait::new(b)
}

can you add this test:

Suggestion:

#[derive(Copy)]
struct A {
    a: NonDrop,
    b: felt252,
}

//! > function
fn main(a: A) -> (Box<felt252>, NonDrop) {
    let ba = BoxTrait::new(a);
    let a = ba.a;
    let b = ba.b;
    (BoxTrait::new(b), x)
}