Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
IanButterworth wants to merge 10 commits into from
Closed

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
101c5cf
lowering: avoid Box for captured variables assigned in all if/elseif/…
IanButterworth Jan 4, 2026
83e6057
lowering: remove ifa candidates when used before assignment
IanButterworth Jan 4, 2026
6078029
lowering: fix edge cases in if-else box avoidance
IanButterworth Jan 5, 2026
430589b
lowering: fix if-else box optimization for sequential and nested cases
IanButterworth Jan 5, 2026
3c721d6
fix whitespace
IanButterworth Jan 5, 2026
6688a10
add test
IanButterworth Jan 6, 2026
f30d620
add test
IanButterworth Jan 6, 2026
c9173c5
make test names clearer
IanButterworth Jan 6, 2026
67b4a68
sort and group tests
IanButterworth Jan 6, 2026
4f6c518
lowering: rewrite if-else box optimization to track single-definition…
IanButterworth Jan 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions NEWS.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,9 @@ Language changes
Compiler/Runtime improvements
-----------------------------

- Captured variables that are assigned in all branches of an `if`/`elseif`/`else` statement
no longer allocate a `Core.Box`, reducing heap allocations in closures ([#60542]).

Command-line option changes
---------------------------

Expand Down
191 changes: 181 additions & 10 deletions src/julia-syntax.scm
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3915,6 +3915,8 @@ f(x) = yt(x)

;; Try to identify never-undef variables, and then clear the `captured` flag for single-assigned,
;; never-undef variables to avoid allocating unnecessary `Box`es.
;; Also handles captured variables assigned in all branches of if-else, treating the if-else
;; as a single definition point when there are no other assignments and capture is after.
(define (lambda-optimize-vars! lam)
(assert (eq? (car lam) 'lambda))
;; memoize all-methods-for to avoid O(n^2) behavior
Expand All @@ -3933,12 +3935,23 @@ f(x) = yt(x)
(decl (table))
(unused (table)) ;; variables not (yet) used (read from) in the current block
(live (table)) ;; variables that have been set in the current block
(seen (table))) ;; all variables we've seen assignments to
;; Collect candidate variables: those that are captured (and hence we want to optimize)
;; and only assigned once. This populates the initial `unused` table.
(seen (table)) ;; all variables we've seen assignments to
;; if-else optimization for captured multi-assigned variables:
(ifa-candidates (table)) ;; captured multi-assigned vars (potential optimization)
(assigned-outside (table)) ;; vars assigned outside any qualifying if-else
(captured-early (table)) ;; vars captured before their if-else completes
(ifa-completed (table)) ;; vars that completed a valid if-else (all branches assigned)
(in-ifa-for (table)) ;; during if-else visit: which vars we're tracking
(used-before-assign (table)) ;; vars used before assigned in some branch
(has-candidates #f)) ;; whether we have any ifa candidates
;; Collect candidate variables: those that are captured and single-assigned go into `unused`.
;; Captured multi-assigned vars go into ifa-candidates for if-else optimization.
(for-each (lambda (v)
(if (and (vinfo:capt v) (vinfo:sa v))
(put! unused (car v) #t)))
(if (vinfo:capt v)
(if (vinfo:sa v)
(put! unused (car v) #t)
(begin (put! ifa-candidates (car v) #t)
(set! has-candidates #t)))))
vi)
(define (restore old)
(table.foreach (lambda (k v)
Expand All @@ -3956,15 +3969,26 @@ f(x) = yt(x)
(if (and (has? unused var) (not (memq var args)))
(del! unused var)))
(define (mark-captured var)
(if (has? unused var)
(del! unused var)))
(if (and (has? unused var) (not (has? ifa-completed var)))
(del! unused var))
;; For ifa candidates: if captured before if-else completes, can't optimize
(if (and (has? ifa-candidates var)
(not (has? ifa-completed var)))
(put! captured-early var #t)))
(define (assign! var)
(if (has? unused var)
;; When a variable is assigned, move it to the live set to protect
;; it from being removed from `unused`.
(begin (put! live var #t)
(put! seen var #t)
(del! unused var))))
(del! unused var)))
;; Track assignments for ifa candidates
(if (has? ifa-candidates var)
(begin
(put! seen var #t)
;; If not currently inside an if-else for this var, it's assigned outside
(if (not (has? in-ifa-for var))
(put! assigned-outside var #t)))))
(define (declare! var)
(if (has? unused var)
(put! decl var #t)))
Expand All @@ -3976,6 +4000,19 @@ f(x) = yt(x)
(del! live k)))
(table.keys live))
(set! decl old-decls))
;; Returns table of vars assigned in ALL branches
(define (intersect-branch-tables branch-tables)
(if (null? branch-tables)
(table)
(let ((result (table.clone (car branch-tables))))
(for-each (lambda (branch-tbl)
(table.foreach
(lambda (var _)
(if (not (has? branch-tbl var))
(del! result var)))
result))
(cdr branch-tables))
result)))
(define (visit e)
;; returns whether e contained a symboliclabel
(cond ((atom? e) (if (symbol? e) (mark-used e))
Expand All @@ -3997,7 +4034,125 @@ f(x) = yt(x)
((eq? (car e) 'symboliclabel)
(kill)
#t)
((memq (car e) '(if elseif trycatch tryfinally trycatchelse))
((eq? (car e) 'if)
(let ((prev (table.clone live)))
(cond
;; if-else with exactly 3 args and we have candidates
((and (length= e 4) has-candidates)
(let ((has-label #f)
(branch-tables '())
(prev-in-ifa (table.clone in-ifa-for))
(branch-assigned (table))) ;; track assignments within current branch
;; Mark eligible candidates as being inside this if-else
(table.foreach
(lambda (var _)
(if (and (not (has? ifa-completed var))
(not (has? assigned-outside var)))
(put! in-ifa-for var #t)))
ifa-candidates)
;; Visit condition
(if (visit (cadr e)) (set! has-label #t))
(let ((pre-branch-live (table.clone live)))
;; Collect assignments from each branch
(let collect-branches ((branches (list (caddr e) (cadddr e)))
(branch-assigns '()))
(if (null? branches)
(set! branch-tables (reverse branch-assigns))
(let ((branch (car branches)))
(set! live (table.clone pre-branch-live))
(kill)
(set! branch-assigned (table))
(let ((branch-tbl (table)))
;; Visit branch, tracking assignments and use-before-assign
(let visit-branch ((expr branch))
(cond
((atom? expr)
(if (symbol? expr)
(begin
(mark-used expr)
;; Track use-before-assign for ifa candidates
(if (and (has? in-ifa-for expr)
(not (has? branch-assigned expr)))
(put! used-before-assign expr #t)))))
((lambda-opt-ignored-exprs (car expr)) #f)
((eq? (car expr) 'symboliclabel)
(set! has-label #t))
((eq? (car expr) '=)
(visit-branch (caddr expr))
(let ((var (cadr expr)))
(assign! var)
(if (has? in-ifa-for var)
(begin
(put! branch-tbl var #t)
(put! branch-assigned var #t)))))
;; Handle loops inside branches - mark vars as assigned-outside
((or (eq? (car expr) '_while) (eq? (car expr) '_do_while))
(let ((prev-in-ifa-loop (table.clone in-ifa-for)))
;; Clear in-ifa-for during loop
(set! in-ifa-for (table))
(for-each visit-branch (cdr expr))
(set! in-ifa-for prev-in-ifa-loop)))
;; Handle elseif chain
((and (eq? (car expr) 'elseif) (length= expr 4))
(visit-branch (cadr expr))
(kill)
(set! branch-assigned (table))
(let ((elseif-branch-tbl (table)))
(let visit-elseif-then ((e2 (caddr expr)))
(cond
Copy link
Member

@topolarity topolarity Jan 6, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

((atom? e2)
(if (symbol? e2)
(begin
(mark-used e2)
(if (and (has? in-ifa-for e2)
(not (has? branch-assigned e2)))
(put! used-before-assign e2 #t)))))
((lambda-opt-ignored-exprs (car e2)) #f)
((eq? (car e2) 'symboliclabel)
(set! has-label #t))
((eq? (car e2) '=)
(visit-elseif-then (caddr e2))
(let ((var (cadr e2)))
(assign! var)
(if (has? in-ifa-for var)
(begin
(put! elseif-branch-tbl var #t)
(put! branch-assigned var #t)))))
((or (eq? (car e2) '_while) (eq? (car e2) '_do_while))
(let ((prev-in-ifa-loop (table.clone in-ifa-for)))
(set! in-ifa-for (table))
(for-each visit-elseif-then (cdr e2))
(set! in-ifa-for prev-in-ifa-loop)))
(else
(for-each visit-elseif-then (cdr e2)))))
(set! branch-assigns (cons elseif-branch-tbl branch-assigns)))
(set! live (table.clone pre-branch-live))
(collect-branches (list (cadddr expr)) branch-assigns))
(else
(for-each visit-branch (cdr expr)))))
(if (not (and (pair? branch) (eq? (car branch) 'elseif)))
(collect-branches (cdr branches)
(cons branch-tbl branch-assigns)))))))
;; Mark vars assigned in ALL branches as completed (if not used-before-assign)
(let ((all-assigned (intersect-branch-tables branch-tables)))
(table.foreach
(lambda (var _)
(if (and (has? in-ifa-for var)
(not (has? captured-early var))
(not (has? used-before-assign var)))
(put! ifa-completed var #t)))
all-assigned))
(set! in-ifa-for prev-in-ifa)
(kill)
(if has-label
(begin (kill) #t)
(begin (restore prev) #f)))))
;; Default if handling
(else
(if (eager-any (lambda (e) (begin0 (visit e) (kill))) (cdr e))
(begin (kill) #t)
(begin (restore prev) #f))))))
((memq (car e) '(elseif trycatch tryfinally trycatchelse))
(let ((prev (table.clone live)))
(if (eager-any (lambda (e) (begin0 (visit e)
(kill)))
Expand All @@ -4008,8 +4163,12 @@ f(x) = yt(x)
(begin (restore prev) #f))))
((or (eq? (car e) '_while) (eq? (car e) '_do_while))
(let ((prev (table.clone live))
(decl- (table.clone decl)))
(decl- (table.clone decl))
(prev-in-ifa (table.clone in-ifa-for)))
;; Clear in-ifa-for during loop - assignments in loop are "outside"
(set! in-ifa-for (table))
(let ((result (eager-any visit (cdr e))))
(set! in-ifa-for prev-in-ifa)
(leave-loop! decl-)
(if result
#t
Expand Down Expand Up @@ -4048,10 +4207,22 @@ f(x) = yt(x)
(let ((vv (assq v vi)))
(vinfo:set-never-undef! vv #t))))
(append (table.keys live) (table.keys unused)))
;; Clear captured flag for single-assigned never-undef variables
(for-each (lambda (v)
(if (and (vinfo:sa v) (vinfo:never-undef v))
(set-car! (cddr v) (logand (caddr v) (lognot 5)))))
vi)
;; Also clear captured flag for ifa-completed vars not assigned outside or captured early
(for-each (lambda (var)
(if (and (has? ifa-completed var)
(not (has? assigned-outside var))
(not (has? captured-early var)))
(let ((vv (assq var vi)))
(if vv
(begin
(vinfo:set-never-undef! vv #t)
(set-car! (cddr vv) (logand (caddr vv) (lognot 5))))))))
(table.keys ifa-candidates))
lam))

(define (is-var-boxed? v lam)
Expand Down
Loading