RFR: 8341782: Allow lambda capture of basic for() loop variables as with enhanced for()

[Archie Cobbs]

On Wed, 9 Oct 2024 09:13:59 GMT, Maurizio Cimadamore <mcimadamore at openjdk.org> wrote:

>> This PR changes the compiler to allow basic `for()` loop iteration variables to be captured by lambdas in the loop body, even when their "effectively final" status is invalidated by modifications in the header of the loop.
>> 
>> This allows code like this to compile:
>> 
>> for (int i = 1; i <= 3; i++) {
>>     Runnable r = () -> System.out.println(i);
>> }
>
> src/jdk.compiler/share/classes/com/sun/tools/javac/comp/Flow.java line 2733:
> 
>> 2731:                 try {
>> 2732:                     scan(tree.body);
>> 2733:                     for (int adr = forLoopVarMin; adr < forLoopVarMax; adr++) {
> 
> Question: since we already have a dedicated flag for loop variables - could we set the flag on _all_ loop variables (before the body), and then remove it in `letInit` (same as we do for effectively final flag) ? This way, at the end of the loop, the variables that still have the flag set will be the ones that have not been reassigned, correct?

Yes - that could be a simpler way to do it. It would allow for the elimination of the `reassignedInForLoopBodyPrev` bit set, which is currently serving the same purpose (but in reverse). However, there would be two new concerns:
* We would need a new flag that says "we are currently in the for() loop body" and not, for example, in the step, where we should leave the flag alone.
* We would also need to ensure we don't get confused with nested `for()` loops. So the aforementioned flag can't just be a boolean (because otherwise which `for()` loop are we talking about?). For example, a symbol address range instead.

That would end up looking something like this... WDYT?

index 29ab8435ada..aa908902059 100644
--- a/src/jdk.compiler/share/classes/com/sun/tools/javac/comp/Flow.java
+++ b/src/jdk.compiler/share/classes/com/sun/tools/javac/comp/Flow.java
@@ -2146,6 +2148,12 @@ public class AssignAnalyzer extends BaseAnalyzer {
          */
         WriteableScope unrefdResources;
 
+        /** In a basic for() loop body, this is the address range of the loop variables.
+         *  If we see a reassignment to any of these variables, we clear the variable's
+         *  FOR_LOOP_BODY_MAY_CAPTURE flag (where "reassigned" means "assigned when not DU").
+         */
+        AddressRange forLoopVarsInBody;
+
         /** Modified when processing a loop body the second time for DU analysis. */
         FlowKind flowKind = FlowKind.NORMAL;
 
@@ -2231,6 +2239,11 @@ void newVar(JCVariableDecl varDecl) {
          */
         void letInit(DiagnosticPosition pos, VarSymbol sym) {
             if (sym.adr >= firstadr && trackable(sym)) {
+
+                // In basic for() loop bodies, variables that are reassigned can no longer be captured
+                if (forLoopVarsInBody != null && forLoopVarsInBody.test(sym.adr) && !uninits.isMember(sym.adr))
+                    sym.flags_field &= ~FOR_LOOP_BODY_MAY_CAPTURE;
+
                 if ((sym.flags() & EFFECTIVELY_FINAL) != 0) {
                     if (!uninits.isMember(sym.adr)) {
                         //assignment targeting an effectively final variable
@@ -2691,6 +2704,7 @@ public void visitForLoop(JCForLoop tree) {
             flowKind = FlowKind.NORMAL;
             int nextadrPrev = nextadr;
             scan(tree.init);
+            final AddressRange loopVars = new AddressRange(nextadrPrev, nextadr);
             final Bits initsSkip = new Bits(true);
             final Bits uninitsSkip = new Bits(true);
             pendingExits = new ListBuffer<>();
@@ -2712,7 +2726,16 @@ public void visitForLoop(JCForLoop tree) {
                     uninitsSkip.assign(uninits);
                     uninitsSkip.inclRange(firstadr, nextadr);
                 }
-                scan(tree.body);
+                final AddressRange forLoopVarsInBodyPrev = forLoopVarsInBody;
+                forLoopVarsInBody = loopVars;
+                try {
+                    loopVars.stream()
+                      .mapToObj(adr -> vardecls[adr].sym)
+                      .forEach(sym -> sym.flags_field |= FOR_LOOP_BODY_MAY_CAPTURE);
+                    scan(tree.body);
+                } finally {
+                    forLoopVarsInBody = forLoopVarsInBodyPrev;
+                }
                 resolveContinues(tree);
                 scan(tree.step);
                 if (log.nerrors != prevErrors ||
@@ -3270,6 +3293,20 @@ public void analyzeTree(Env<?> env, JCTree tree, TreeMaker make) {
         }
     }
 
+    /** A contiguous range of symbol addresses.
+     */
+    record AddressRange(int min, int max) implements IntPredicate {
+
+        @Override
+        public boolean test(int adr) {
+            return adr >= this.min() && adr < this.max();
+        }
+
+        public IntStream stream() {
+            return IntStream.range(this.min(), this.max());
+        }
+    }
+
     /**
      * This pass implements the last step of the dataflow analysis, namely
      * the effectively-final analysis check. This checks that every local variable

-------------

PR Review Comment: https://git.openjdk.org/jdk/pull/21415#discussion_r1793750314