Don't box captured variables that are not assigned after closure declaration

[uniment]

Using conditional logic and/or reassignments to initialize a variable before capturing it is a common pattern, and it's the source of most unnecessary boxes for which #15276 continues to be mentioned, so it's attractive to attempt a solution.

Consider the following example that defines two functionally equivalent inner functions (modified from Performance Tips):

julia> function abmult(r1::Int)
           if r1 < 0
               r1 = -r1
           end
           r2 = r1

           f1 = x -> x * r1
           f2 = x -> x * r2
           return f1, f2
       end
abmult (generic function with 1 method)

julia> abmult(1)
(var"#3#5"(Core.Box(1)), var"#4#6"{Int64}(1))

Notice that, throughout the lifetimes of both closures f1 and f2, neither capture r1 nor r2 can ever take on a new value—there's never an assignment subsequent to the closures' declarations in their parent scope, nor in their function bodies. f1 and f2 behave entirely identically (except for performance); clearly, the language semantics can be satisfied without boxing r1 here.

As far as I can tell from experimentation, the current rule for boxing a capture is to box if:
a) the variable has more than one syntactical assignment in parent scope,
b) the variable is syntactically assigned within a conditional in parent scope,
c) the variable is defined only after the closure is declared, or if
d) the variable is syntactically assigned within a closure that captures it.
However, this boxing policy (specifically, (a) and (b)) decides to box too aggressively in many cases, such as this.

To address this, I propose a different boxing rule: a capture should be boxed iff there is a syntactical assignment to it, which is syntactically reachable after any of its closures' instantiation. This is the actual boxing rule that we are reaching for—the only reason to box a capture is if its value could change during its closures' lifetimes; the currently implemented rule is merely an imperfect approximation of this.

To make this proposal more concrete, consider the following procedure that aims to implement it. For illustration purposes this is not optimized. Here I assume working with IR code so that each instruction is a single line, prior to insertion of boxes:

For a given variable x that is known to be captured by a closure f (and possibly closures g, h, etc.):

1. If x is syntactically assigned to, anywhere within f, g, h, or etc., then box x.
2. If x is not syntactically assigned to, anywhere within its parent scope, then x is a global; stop.
3. Initialize an empty record of explored paths (a Set of line numbers).
4. Call check_branch on the line on which f is instantiated in its parent scope.
5. function check_branch(starting_line): Starting at starting_line, proceed line-by-line:
   a. If the current line is in the list of explored paths, then stop exploring this path (return). (to avoid infinite loops)
   b. Push the current line into the list of explored paths.
   c. If this line is an assignment to x, then box x.
   d. If the current line branches, goto _ if not _, then explore both paths (i.e. call check_branch on the line where goto goes to).
   e. If the end of the parent scope has been reached, then stop (return).
6. If x is captured by additional closures g, h, etc., call check_branch on the line where each closure is initialized.
7. If all the above has finished without boxing, then x is an unboxed capture.

This should be $\mathcal O(n)$ in the number of instructions.

[vtjnash]

Duplicate of #15276

And yes, it sounds like the rule you are attempting to define is traditionally referred to as dominator analysis. There is essentially a rudimentary pass, but it gets confused by cases such as your example. An accurate dominator tree computation on the AST during lowering would be able to correctly identify cases such as yours.

[uniment]

#15276 is a compendium of capture boxing problems, whereas this issue specifically focuses on one solution. This proposal will solve a majority of the persisting and new problems, but the remainder will require other approaches.

An accurate dominator tree computation on the AST during lowering would be able to correctly identify cases such as yours.

This proposal is a simple tree traversal, not a dominator analysis, and thus has linear complexity instead of quadratic or cubic. The detail is in determining what is the correct tree to traverse. (see sample algorithm in OP)

@Seelengrab reopen this please.

[Seelengrab]

I'll reopen, but want to add that I'm not sure having a different issue to try to solve a particular version of the problem is good for tracking the overall state of the issue. Duplicates usually get closed because of that. To my knowledge, what's going to solve the issue is an eventual redesign of the lowering stage, without piling on ever-more bandaids for special cases.

[uniment]

Thanks for reopening.

I'd like to point out that this is not an additional special case. Instead, this proposal outlines the exact rule to determine whether a capture should be boxed, and it would replace what is currently a collection of bandaids.

With this rule implemented (alongside a rule for named functions described in #53295), the only way to improve #15276 further is to have lowering perform type inference and parameterize the box. That will be a separate effort.

For tracking purposes, would it be better to make this proposal as a comment in #15276 instead of a separate issue?

[Seelengrab]

For tracking purposes, would it be better to make this proposal as a comment in #15276 instead of a separate issue?

I think so, yes. Alternatively, a PR implementing this to see how it performs on a run of PkgEval would be ideal - I don't think there's too many people who can comment on how well this will work ahead of time anyway.

[JeffBezanson]

I like separate issues in this particular case, because there are actually several different patterns of how bindings can be used and captured, and some are much easier to handle than others. #15276 is on thin ice since it's probably not possible to fix 100%.