lowering: avoid Core.Box for captured variables assigned in all if/elseif/else branches

[IanButterworth]

An attempt at doing the Core.Box avoidance mentioned in #60479 (comment) automatically @KristofferC

Developed with Claude:

---

Variables that are captured by a closure and assigned in all branches of an
if/elseif/else statement are now treated as effectively single-assigned,
avoiding unnecessary Core.Box heap allocations.

The optimization is disabled when:

• @goto is present (could skip assignments)
• Variable is used before assignment in any branch
• Variable is captured before the if-else block

[IanButterworth]

Not that I'd expect much, but it doesn't seem to have much of a performance improvement wrt. CI timing. Comparison of this build against recent timing

[topolarity]

I'm a big fan of this improvement, but the test coverage is a little thin TBH

Ideally, a change like this would have tests for correctness (not optimizing when it's unsafe) in addition to the tests you added for precision (optimizing when it is safe). I think this would benefit from more test coverage or (very) thorough review to make sure the approach here is sound.

[mlechu]

I'm not sure this makes sense. Could you maybe explain how the added unboxing logic works? My concern is that throughout lambda-optimize-vars, "single assign" (vinfo:sa) always means "only one syntactic assignment in the whole thunk" but your PR description and the LLM comments imply it's some sort of more local condition.

Easy counterexample:

function f()
    if rand() > 0.5
        x = 1
    else
        x = 2
    end
    function g()
        @show x
    end
    if true
        x = 123
    end
    g()
end; f()

[IanButterworth]

Thanks for the reviews! The tests are now expanded, and claude seems to have figured out getting them to pass.

Perhaps the main value here is developing out and agreeing on the tests.
Please feel free to disregard or redirect the implementation if there is a better strategy.

[xal-0]

I worry that if I provide another counterexample, Claude will produce an even more convoluted "analysis" that hacks around the test

function f()
    x = 1
    function g()
        @show x
    end
    if rand() > 0.5
        x = 2
    else
        x = 3
    end
    g()
end

[IanButterworth]

Yeah that fails. I'll add it as a failing test and leave the PR where it is

[adienes]

here is another

function f()
    x = 0
    if true
        x = 1
        g = () -> x
        x = 2
    else
        x = 3
    end
    return g()
end

in particular, I would think this optimization is only possible if the variable is never captured again within one of the branches before assignment

[IanButterworth]

Thanks. I added that test too.

I also got claude to review the tests and reviews here, and come up with a better strategy.

Here's the summary doc it made, and strategy D is now implemented, and passes tests locally:

Box Avoidance Optimization Strategies

Problem Statement

When a variable is captured by a closure and assigned multiple times, Julia allocates a Core.Box to hold the value so the closure can see updates. This PR attempts to avoid the Box when the variable is assigned in all branches of an if-else (effectively single-assigned per control path).

Test Cases

Optimized cases (should NOT box)

#	Test Name	Pattern	Description
1	ifa_basic	if c; x=a else x=b end; ()->x	Basic case
2	ifa_multi_vars	Same with two vars	Multiple vars
3	ifa_elseif	if c1; x=a elseif c2; x=b else x=c end	Elseif chain
4	ifa_nested	Nested if-else, all paths assign	Nested structure
5	ifa_deeply_nested	3 levels deep, all paths assign	Deep nesting

Boxed cases - incomplete branch coverage

#	Test Name	Pattern	Description
6	box_missing_else	x=0; if c; x=a end; ()->x	Missing else branch

Boxed cases - use/capture timing

#	Test Name	Pattern	Description
7	box_use_before_assign	if c; x=a else print(x); x=b end	Use before assign in branch
8	box_capture_before_any_assign	f=()->x; if c; x=a else x=b end	Capture before any assignment
9	box_assign_and_capture_before_ifelse	x=1; g=()->x; if c; x=2 else x=3 end	Assign+capture before if-else
10	box_capture_inside_branch_then_assign	if true; x=1; g=()->x; x=2 else x=3 end	Capture then assign in branch

Boxed cases - control flow

#	Test Name	Pattern	Description
11	box_goto_skips_assign	if c; @goto skip; x=a else x=b end; @label skip	Goto can skip assignments

Boxed cases - loops

#	Test Name	Pattern	Description
12	box_for_loop_in_branch	if c; x=a; for i...; x=i end else x=b end	For loop in branch
13	box_while_loop_in_branch	if c; x=a; while...; x=i end else x=b end	While loop in branch
14	box_ifelse_inside_loop	for i...; if c; x=a else x=b end end; ()->x	If-else inside loop

Boxed cases - multiple assignments

#	Test Name	Pattern	Description
15	box_assign_after_capture	if c; x=1 else x=2 end; g=()->x; x=123	Assign after capture

Strategies

Strategy A: Original Attempt (Complex State Tracking)

Track multiple tables (ifa, ifa-promoted, ifa-used-first, ifa-assigned) to handle various cases.

Pros: Attempts to handle many cases
Cons: Complex, had bugs with cases 9 and 10, hard to reason about correctness

Test	Result
1-5, 6-8, 11-15	✅ Pass
9	❌ FAIL (returns Int64, should be Box)
10	❌ FAIL

Strategy B: Only-Assignment-Site (Strict)

Rule: Only optimize if the if-else block contains the ONLY assignments to the variable in the entire function.

Check: total_assignments == branches_in_if_else (exactly one assignment per branch, no assignments outside)

Implementation:

1. Count total assignments to each captured var
2. During if-else visit, count assignments per branch
3. Only promote if total assignments equals branch count

Test	Result	Reasoning
1	✅	2 assignments total, 2 branches, all inside if-else
2	✅	Same logic per var
3	✅	3 assignments total, 3 branches
4	❌	Has x=0 outside → don't optimize
5	✅	2 assignments, 2 branches (use-before-assign is separate issue)
6	✅	2 assignments, 2 branches... but capture is before! Need additional check
7	✅	Would need goto check
8	✅	4 assignments in nested structure, 4 leaf branches
9	❌	Loop creates multiple assignments
10	❌	While creates multiple assignments
11	✅	8 assignments, 8 branches
12	❌	Loop around if-else
13	❌	x=123 is outside the optimized if-else
14	❌	x=1 is outside the if-else
15	❌	x=2 is after capture within branch

Pros: Simple to understand, conservative
Cons: Still needs capture-timing check, may miss cases 5,6,7

Strategy C: Capture-After-Only-If-Else (Very Strict)

Rule: Only optimize if:

1. Variable has NO assignments outside the if-else
2. Variable is captured ONLY after the if-else completes
3. No loops/gotos that could cause re-execution

Implementation:

• Track if var was ever assigned outside any if-else
• Track if var was ever captured before/during if-else
• Disqualify if either is true

Test	Result	Reasoning
1-3	✅	Clean case
4	❌	x=0 before
5	❌	Use (which implies potential capture context) before assign
6	❌	Capture before if-else
7	❌	Goto present
8, 11	✅	Nested but clean
9-10	❌	Loop in branch
12	❌	If-else in loop
13	❌	Second if-else
14	❌	Assign before if-else
15	❌	Capture inside branch

Pros: Very conservative, likely correct
Cons: May miss some valid optimizations

Strategy D: Single-Definition-Region (Implemented)

Rule: The if-else block must be the "single definition region" for the variable:

1. No assignments to var outside the if-else (before or after)
2. No use of var before assignment in any branch
3. Exactly one assignment per branch (all branches must assign)
4. Capture must occur after the if-else completes
5. No loops inside branches (which would create multiple assignments)

This treats the if-else as a single φ-node (single definition point).

Implementation:
Uses these tracking tables:

• ifa-candidates: captured multi-assigned vars (potential optimization targets)
• assigned-outside: vars assigned outside any qualifying if-else
• captured-early: vars captured before their if-else completes
• ifa-completed: vars that completed a valid if-else (all branches assigned)
• in-ifa-for: during if-else visit, which vars we're tracking
• used-before-assign: vars used before assigned in some branch

During visit:

• When entering if-else, mark eligible candidates in in-ifa-for
• Track assignments per branch; intersect to find vars assigned in ALL branches
• Track use-before-assign within branches
• Clear in-ifa-for when entering loops (assignments in loops are "outside")
• Mark vars as ifa-completed if assigned in all branches, not captured early, not used-before-assign

At end:

• Promote ifa-completed vars that aren't assigned-outside or captured-early

Test	Pass?
1-5	✅ (optimized)
6-15	✅ (correctly boxed)

Final Implementation

Strategy D was implemented. Key aspects:

• Clear semantic meaning (if-else acts as single definition point)
• Handles all 15 test cases correctly
• Conservative but covers the main use case
• Adds ~130 lines to lambda-optimize-vars!

The key insight: we're not trying to do general dataflow analysis. We're identifying a specific pattern where an if-else block IS the variable's only definition, making it effectively single-assigned.

[topolarity]

I'm not really convinced by this iteration-via-Claude-via-third-party approach to solving this problem.

This algorithm is still quite limited, as demonstrated by the fact that it handles:

if cond
    x = 1
else
    x = 2
end
return ()->x

but not:

x = 2
if cond
    x = 1
end
return ()->x

In fact many of the added tests are examples of cases that would be supported by a more complete control-flow analysis, not of required safety or good behavior. Having several serious bugs found so far does not increase my confidence.

Here is (yet) another one:

function ifa_basic(cond, a, b)
    local x
    while true
        if cond
            x = a
        else
            break
            x = b
        end
        break
    end
    return ()->x
end

With your change ifa_basic(false, 1, 2) throws UndefVarError because lowering believes it to be always-assigned and eligible to be unboxed.

I think we probably need to go a different direction here. This approach introduces a substantial amount of ad-hoc control flow analysis to closure conversion, and I don't think the extra complexity justifies such a specific optimization, especially given the bugs so far.

[topolarity]

Much of this code is heavily duplicated above - it could use re-factoring

[topolarity]

Another imprecise case

[topolarity]

Again this is testing that the analysis is imprecise (does not optimize when it would be legal to)

[topolarity]

Another imprecise case

[topolarity]

Another imprecise case

[topolarity]

Many (most) of these are situations where the optimization is actually valid, but the algorithm has decided not to optimize anyway.

Useful tests need to cover cases where the optimization is actually unsafe to apply.

These are just asserting that the algorithm is imprecise, which is not useful.

[adienes]

while we're throwing AI at it, here's three more failure modes gemini helped me get:

function f()
    if true
        if false
            x = 1
        end
    else
        x = 2
    end
    return () -> x
end
@test f() isa Function

function f()
    @goto mark
    if true
        x = 1
        @label mark
    else
        x = 2
    end
    return () -> x
end
@test f() isa Function

function f()
    if true
        x = 1
        g = () -> x
        x = 2
    else
        g = () -> 0
        x = 3
    end
    return g
end
@test f()() == 2

[IanButterworth]

Ok. I think I'm not being effective in trying to help this along. Sorry for the noise.