(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): Reboxing now also works on whole var reboxing #8994
• (optimization): Reboxing also applied on Box of snapshots #8954
• (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/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)
}

[ilyalesokhin-starkware]

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

[ilyalesokhin-starkware]

@ilyalesokhin-starkware reviewed 1 of 2 files at r2, 1 of 3 files at r8.
Reviewable status: 11 of 13 files reviewed, 6 unresolved discussions (waiting on @array, @eytan-starkware, @giladchase, and @orizi)

[eytan-starkware]

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

---

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

Previously, ilyalesokhin-starkware wrote…

?

Done

---

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

Previously, ilyalesokhin-starkware wrote…

Is it possible that this var is not droppable?

This is a bug. For now I am blocking on any non-droppable members.

---

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

Previously, ilyalesokhin-starkware wrote…

can you add this test:

Done (with slight changes)

---

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

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

Done.

[ilyalesokhin-starkware]

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

    let b = aa.a;
    (BoxTrait::new(b), aa.b)
}

give the vars better names.

Suggestion:

struct NonDrop {
    b: felt252,
}

#[derive(Copy)]
struct A {
    a: felt252,
    non_drop: B,
}

//! > function_code
fn main(a: A) -> (Box<felt252>, NonDrop) {
    let ba = BoxTrait::new(a);
    let unboxed = ba.unbox();
    let a = unboxed.a;
    (BoxTrait::new(a), unboxed.non_drop)
}

[ilyalesokhin-starkware]

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

    } else {
        // Unboxed reboxing not supported yet.
    }

use let ... else

Code quote:

            }
        } else {
            // Nested MemberOfUnboxed not supported yet.
        }
    } else {
        // Unboxed reboxing not supported yet.
    }

[ilyalesokhin-starkware]

@ilyalesokhin-starkware reviewed 5 of 5 files at r9.
Reviewable status: all files reviewed (commit messages unreviewed), 4 unresolved discussions (waiting on @array, @eytan-starkware, @giladchase, and @orizi)

[orizi]

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

---

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

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

all around.

Suggestion:

fn main(a: Box<A>) -> Box<felt252> {
    BoxTrait::new(a.b)
}

---

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

  (v6: test::Point) <- core::box::unbox::<test::Point>(v1)
  (v7: core::integer::u32, v8: core::integer::u32, v9: core::integer::u32) <- struct_destructure(v6)
  (v10: core::box::Box::<core::integer::u32>, v13: core::box::Box::<core::integer::u32>, v14: core::box::Box::<core::integer::u32>) <- struct_destructure(v1)

at least add todo to remove the double destructring.

[eytan-starkware]

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

---

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

Previously, ilyalesokhin-starkware wrote…

use let ... else

Done.

---

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

Previously, orizi wrote…

all around.

Done

---

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

Previously, orizi wrote…

at least add todo to remove the double destructring.

Done (also in reboxing.rs)

---

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

Previously, ilyalesokhin-starkware wrote…

give the vars better names.

Done.

[ilyalesokhin-starkware]

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

    // TODO(eytan-starkware): When applied, reboxing analysis should replace the existing
    // deconstruct, and add unbox statements on members as needed.

Do you mean?

Suggestion:

    // TODO(eytan-starkware): When applied, reboxing analysis should replace the existing
    // unbox + deconstruct, and add unbox statements on members as needed.

[ilyalesokhin-starkware]

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

        // If source is not member of unboxed, we are reboxing original value which is not supported
        // yet.
        return;

why not?

Code quote:

        // If source is not member of unboxed, we are reboxing original value which is not supported
        // yet.
        return;

[ilyalesokhin-starkware]

@ilyalesokhin-starkware reviewed 3 of 4 files at r10.
Reviewable status: 12 of 13 files reviewed, 4 unresolved discussions (waiting on @eytan-starkware, @giladchase, and @orizi)

[eytan-starkware]

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

---

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

Previously, ilyalesokhin-starkware wrote…

Do you mean?

I meant boxed deconstruct, so essentially a deconstruct statement on the boxed struct. The on the resulting box members we need to unbox what was previously used as an unboxed member.

---

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

Previously, ilyalesokhin-starkware wrote…

why not?

Wanted to keep the analysis simpler for now. It means I need more logic in my state tracking, and to allow tracking enums as well as structs so I am postponing

[orizi]

@orizi reviewed 1 of 3 files at r8, 1 of 5 files at r9, 1 of 4 files at r10, 1 of 1 files at r11, all commit messages.
Reviewable status: all files reviewed, 14 unresolved discussions (waiting on @eytan-starkware, @giladchase, and @ilyalesokhin-starkware)

---

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

Previously, eytan-starkware wrote…

Done

i see no need to separate into 2 statements.

---

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

//! > ==========================================================================

//! > Test reboxing of a non-copyable member of a snapshotted boxed struct.

is the idea of these test that more info would appear in them in the future?

---

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

fn foo(box: @Box<A>) -> Box<@Array<felt252>> {
    BoxTrait::new(box.as_snapshot().unbox().a)
}

Suggestion:

fn foo(box: Box<@A>) -> Box<@Array<felt252>> {
    BoxTrait::new(box.a)
}

---

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

//! > function_code
fn multi_member(boxed_p: Box<Point>) -> (Box<u32>, Box<u32>) {

remove the boxed_ all around.

Code quote:

fn multi_member(boxed_p: Box<Point>) -> (Box<u32>, Box<u32>) {

---

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

    let x = boxed_p.x;
    let z = boxed_p.z;
    (BoxTrait::new(x), BoxTrait::new(z))

Suggestion:

//! > function_code
fn multi_member(p: Box<Point>) -> (Box<u32>, Box<u32>) {
    (BoxTrait::new(p.x), BoxTrait::new(p.z))

---

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

    let inner = boxed_o.inner;
    BoxTrait::new(inner)
}

i thought this would be:

as this is not exactly the same as the previous options.

Suggestion:

fn nested_struct(o: Box<Outer>) -> Box<felt252> {
    BoxTrait::new(o.inner.value)
}

---

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

//! > module_code
use core::box::BoxTrait;

in prelude

Suggestion:

//! > module_code

---

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

rebox_whole

//! > TODO(eytan-starkware): Add support for whole var reboxing

sounds more like a "cancel op" sort of thing - doesn't it?

---

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

//! > ==========================================================================

//! > Test no reboxing opportunity (independent operations)

aren't there any more interesting "no reboxing" cases?

---

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

error[E0007]: Type "core::box::Box::<test::MyEnum>" has no member "A"
 --> lib.cairo:9:16
    let a = be.A;

?

---

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

//! > ==========================================================================

//! > Test reboxing with branch-specific unbox assignments

no opt happening - doc this isn't supposed to work - or add todo.

---

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

//! > ==========================================================================

//! > Test reboxing with loop break remapping

don't see a reason for this test.

[ilyalesokhin-starkware]

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

Previously, eytan-starkware wrote…

Wanted to keep the analysis simpler for now. It means I need more logic in my state tracking, and to allow tracking enums as well as structs so I am postponing

if you have unbox+box, you don't need to track the type at all.

[ilyalesokhin-starkware]

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

Previously, eytan-starkware wrote…

I meant boxed deconstruct, so essentially a deconstruct statement on the boxed struct. The on the resulting box members we need to unbox what was previously used as an unboxed member.

can you fix the TODO, I don't understand it as it currently is.

[eytan-starkware]

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

---

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

Previously, ilyalesokhin-starkware wrote…

can you fix the TODO, I don't understand it as it currently is.

Done.

---

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

Previously, ilyalesokhin-starkware wrote…

if you have unbox+box, you don't need to track the type at all.

f2f

---

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

Previously, orizi wrote…

i see no need to separate into 2 statements.

Done

---

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

Previously, orizi wrote…

remove the boxed_ all around.

Done

---

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

Previously, orizi wrote…

i thought this would be:

as this is not exactly the same as the previous options.

This wont be supported yet, but you are correct.

---

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

Previously, orizi wrote…

in prelude

Removed

---

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

Previously, orizi wrote…

sounds more like a "cancel op" sort of thing - doesn't it?

Both cancel op and reboxing could be merged eventually, but currently it will be a bit of a mess as they deal with a bit different situations (cancel op is a back analysis)

---

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

Previously, orizi wrote…

aren't there any more interesting "no reboxing" cases?

Done.

---

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

Previously, orizi wrote…

?

Oops

---

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

Previously, orizi wrote…

no opt happening - doc this isn't supposed to work - or add todo.

Done.

---

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

Previously, orizi wrote…

don't see a reason for this test.

Removed

---

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

Previously, orizi wrote…

is the idea of these test that more info would appear in them in the future?

That we will see them optimized in the future, yes

---

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

    let x = boxed_p.x;
    let z = boxed_p.z;
    (BoxTrait::new(x), BoxTrait::new(z))

Done

---

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

fn foo(box: @Box<A>) -> Box<@Array<felt252>> {
    BoxTrait::new(box.as_snapshot().unbox().a)
}

Done.

[orizi]

@orizi reviewed 3 of 4 files at r12, all commit messages.
Reviewable status: 12 of 13 files reviewed, 2 unresolved discussions (waiting on @giladchase and @ilyalesokhin-starkware)