Skip to content

Commit 2e3584f

Browse files
IanButterworthclaude
IanButterworth
and
claude
committed
lowering: avoid Box for captured variables assigned in all if/elseif/else branches
Enhance lambda-optimize-vars! to recognize when a captured variable is assigned in all branches of an if-else or if-elseif-else statement. Such variables are effectively single-assigned on each control flow path and don't need boxing even though they appear to be assigned multiple times syntactically. This avoids unnecessary Core.Box allocations for common patterns like: if cond1 x = a elseif cond2 x = b else x = c end return () -> x Co-Authored-By: Claude <noreply@anthropic.com>
1 parent 42ad41c commit 2e3584f

3 files changed

Lines changed: 166 additions & 4 deletions

File tree

‎NEWS.md‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,9 @@ Language changes
1818
Compiler/Runtime improvements
1919
-----------------------------
2020

21+
- Captured variables that are assigned in all branches of an `if`/`elseif`/`else` statement
22+
no longer allocate a `Core.Box`, reducing heap allocations in closures ([#60542]).
23+
2124
Command-line option changes
2225
---------------------------
2326

‎src/julia-syntax.scm‎

Lines changed: 88 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3915,6 +3915,8 @@ f(x) = yt(x)
39153915

39163916
;; Try to identify never-undef variables, and then clear the `captured` flag for single-assigned,
39173917
;; never-undef variables to avoid allocating unnecessary `Box`es.
3918+
;; Also handles the case where a variable is assigned in both branches of an if-else, making
3919+
;; it effectively single-assigned on each control path.
39183920
(define (lambda-optimize-vars! lam)
39193921
(assert (eq? (car lam) 'lambda))
39203922
;; memoize all-methods-for to avoid O(n^2) behavior
@@ -3933,12 +3935,18 @@ f(x) = yt(x)
39333935
(decl (table))
39343936
(unused (table)) ;; variables not (yet) used (read from) in the current block
39353937
(live (table)) ;; variables that have been set in the current block
3936-
(seen (table))) ;; all variables we've seen assignments to
3938+
(seen (table)) ;; all variables we've seen assignments to
3939+
(ifa (table)) ;; variables assigned in all branches of if-else ("if-assigned")
3940+
(has-ifa #f)) ;; whether ifa has any entries
39373941
;; Collect candidate variables: those that are captured (and hence we want to optimize)
39383942
;; and only assigned once. This populates the initial `unused` table.
3943+
;; Also collect captured variables assigned more than once for if-branch analysis.
39393944
(for-each (lambda (v)
3940-
(if (and (vinfo:capt v) (vinfo:sa v))
3941-
(put! unused (car v) #t)))
3945+
(if (vinfo:capt v)
3946+
(if (vinfo:sa v)
3947+
(put! unused (car v) #t)
3948+
(begin (put! ifa (car v) #t)
3949+
(set! has-ifa #t)))))
39423950
vi)
39433951
(define (restore old)
39443952
(table.foreach (lambda (k v)
@@ -3997,7 +4005,75 @@ f(x) = yt(x)
39974005
((eq? (car e) 'symboliclabel)
39984006
(kill)
39994007
#t)
4000-
((memq (car e) '(if elseif trycatch tryfinally trycatchelse))
4008+
((eq? (car e) 'if)
4009+
;; Special handling for if-else: track variables assigned in ALL branches.
4010+
;; If a captured variable is assigned in ALL branches (and not used before
4011+
;; assignment in any), it's effectively single-assigned per control path.
4012+
(let ((prev (table.clone live)))
4013+
(cond
4014+
;; if-else with exactly 3 args (cond, then, else) and we have candidates
4015+
((and (length= e 4) has-ifa)
4016+
(let ((has-label #f)
4017+
(all-assigned #f))
4018+
;; Visit condition
4019+
(if (visit (cadr e)) (set! has-label #t))
4020+
(let ((pre-branch-live (table.clone live)))
4021+
(kill)
4022+
;; Visit then-branch
4023+
(if (visit (caddr e)) (set! has-label #t))
4024+
(set! all-assigned (table.clone live))
4025+
;; Process else-branch (may be elseif chain)
4026+
(let process-else ((else-expr (cadddr e)))
4027+
(set! live (table.clone pre-branch-live))
4028+
(kill)
4029+
(cond
4030+
;; else-branch is an elseif
4031+
((and (pair? else-expr) (eq? (car else-expr) 'elseif)
4032+
(length= else-expr 4))
4033+
;; Visit elseif condition
4034+
(if (visit (cadr else-expr)) (set! has-label #t))
4035+
(let ((pre-elseif-live (table.clone live)))
4036+
(kill)
4037+
;; Visit elseif then-branch
4038+
(if (visit (caddr else-expr)) (set! has-label #t))
4039+
;; Intersect with all-assigned
4040+
(let ((branch-assigned live))
4041+
(table.foreach
4042+
(lambda (var _)
4043+
(if (not (has? branch-assigned var))
4044+
(del! all-assigned var)))
4045+
all-assigned))
4046+
;; Process nested else
4047+
(process-else (cadddr else-expr))))
4048+
;; else-branch is regular expression (final else)
4049+
(else
4050+
(if (visit else-expr) (set! has-label #t))
4051+
;; Intersect with all-assigned
4052+
(let ((branch-assigned live))
4053+
(table.foreach
4054+
(lambda (var _)
4055+
(if (not (has? branch-assigned var))
4056+
(del! all-assigned var)))
4057+
all-assigned)))))
4058+
;; Mark variables assigned in all branches as effectively single-assigned
4059+
(table.foreach
4060+
(lambda (var _)
4061+
(if (has? all-assigned var)
4062+
(begin
4063+
(put! seen var #t)
4064+
(put! unused var #t)
4065+
(del! ifa var))))
4066+
ifa)
4067+
(kill)
4068+
(if has-label
4069+
(begin (kill) #t)
4070+
(begin (restore prev) #f)))))
4071+
;; No ifa candidates - use default handling
4072+
(else
4073+
(if (eager-any (lambda (e) (begin0 (visit e) (kill))) (cdr e))
4074+
(begin (kill) #t)
4075+
(begin (restore prev) #f))))))
4076+
((memq (car e) '(elseif trycatch tryfinally trycatchelse))
40014077
(let ((prev (table.clone live)))
40024078
(if (eager-any (lambda (e) (begin0 (visit e)
40034079
(kill)))
@@ -4048,10 +4124,18 @@ f(x) = yt(x)
40484124
(let ((vv (assq v vi)))
40494125
(vinfo:set-never-undef! vv #t))))
40504126
(append (table.keys live) (table.keys unused)))
4127+
;; Clear captured flag for single-assigned never-undef variables
40514128
(for-each (lambda (v)
40524129
(if (and (vinfo:sa v) (vinfo:never-undef v))
40534130
(set-car! (cddr v) (logand (caddr v) (lognot 5)))))
40544131
vi)
4132+
;; Also clear captured flag for variables that were assigned in all branches of an if-else
4133+
;; (these are in `unused` but not `ifa`, and have never-undef set)
4134+
(for-each (lambda (var)
4135+
(let ((vv (assq var vi)))
4136+
(if (and vv (vinfo:never-undef vv) (not (has? ifa var)))
4137+
(set-car! (cddr vv) (logand (caddr vv) (lognot 5))))))
4138+
(table.keys unused))
40554139
lam))
40564140

40574141
(define (is-var-boxed? v lam)

‎test/syntax.jl‎

Lines changed: 75 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4673,3 +4673,78 @@ module M59755 end
46734673
@test M59755.v6 === 5
46744674
@test Base.binding_kind(M59755, :v6) == Base.PARTITION_KIND_CONST
46754675
end
4676+
4677+
# Test that `if cond; x = a; else x = b; end` doesn't introduce a Core.Box
4678+
# when x is captured, due to optimization in lambda-optimize-vars!
4679+
@testset "if-else assignment without boxing" begin
4680+
function if_else_nobox(cond, a, b)
4681+
if cond
4682+
x = a
4683+
else
4684+
x = b
4685+
end
4686+
return () -> x
4687+
end
4688+
closure_type = typeof(if_else_nobox(true, 1, 2))
4689+
@test fieldtype(closure_type, 1) !== Core.Box
4690+
@test if_else_nobox(true, 1, 2)() == 1
4691+
@test if_else_nobox(false, 1, 2)() == 2
4692+
4693+
# Also works with multiple variables assigned in both branches
4694+
function if_else_nobox_multi(cond, a, b, c, d)
4695+
if cond
4696+
@noinline identity(1)
4697+
x = a
4698+
y = b
4699+
else
4700+
@noinline identity(2)
4701+
x = c
4702+
y = d
4703+
end
4704+
return () -> (x, y)
4705+
end
4706+
closure_type2 = typeof(if_else_nobox_multi(true, 1, 2, 3, 4))
4707+
@test fieldtype(closure_type2, 1) !== Core.Box
4708+
@test fieldtype(closure_type2, 2) !== Core.Box
4709+
@test if_else_nobox_multi(true, 1, 2, 3, 4)() == (1, 2)
4710+
@test if_else_nobox_multi(false, 1, 2, 3, 4)() == (3, 4)
4711+
4712+
# Also works with elseif chains
4713+
function if_elseif_nobox(cond1, cond2, a, b, c)
4714+
if cond1
4715+
x = a
4716+
elseif cond2
4717+
x = b
4718+
else
4719+
x = c
4720+
end
4721+
return () -> x
4722+
end
4723+
closure_type3 = typeof(if_elseif_nobox(true, false, 1, 2, 3))
4724+
@test fieldtype(closure_type3, 1) !== Core.Box
4725+
@test if_elseif_nobox(true, false, 1, 2, 3)() == 1
4726+
@test if_elseif_nobox(false, true, 1, 2, 3)() == 2
4727+
@test if_elseif_nobox(false, false, 1, 2, 3)() == 3
4728+
4729+
# Variable assigned in only one branch must still be boxed
4730+
function if_else_onebranch(cond, a)
4731+
x = 0
4732+
if cond
4733+
x = a
4734+
end
4735+
return () -> x
4736+
end
4737+
@test fieldtype(typeof(if_else_onebranch(true, 1)), 1) === Core.Box
4738+
4739+
# Variable used before assignment in one branch must still be boxed
4740+
function if_else_usefirst(cond, a, b)
4741+
if cond
4742+
x = a
4743+
else
4744+
@noinline println(devnull, x)
4745+
x = b
4746+
end
4747+
return () -> x
4748+
end
4749+
@test fieldtype(typeof(if_else_usefirst(true, 1, 2)), 1) === Core.Box
4750+
end

0 commit comments

Comments
 (0)