Fix resolution of for-generator vars within iteratee

pkl-core/src/main/java/org/pkl/core/ast/builder/AstBuilder.java

@@ -900,7 +900,7 @@ public GeneratorMemberNode visitObjectEntry(ObjectEntryContext ctx) {
   @Override
   public GeneratorMemberNode visitObjectSpread(ObjectSpreadContext ctx) {
     return GeneratorSpreadNodeGen.create(
-        createSourceSection(ctx), visitExpr(ctx.expr()), ctx.QSPREAD() != null);
+        createSourceSection(ctx), visitGeneratorIterable(ctx.expr()), ctx.QSPREAD() != null);
   }
 
   private void insertWriteForGeneratorVarsToFrameSlotsNode(@Nullable MemberNode memberNode) {
@@ -964,6 +964,26 @@ private static boolean isIgnored(@Nullable ParameterContext param) {
     return param != null && param.UNDERSCORE() != null;
   }
 
+  private GeneratorIterableNode visitGeneratorIterable(ExprContext ctx) {
+    var enclosingForGeneratorVariables = symbolTable.getCurrentScope().getForGeneratorVariables();
+    var frameDescriptor = symbolTable.getCurrentScope().buildFrameDescriptor();
+    var auxiliarySlots = new int[enclosingForGeneratorVariables.size()];
+    var i = 0;
+    for (var forGeneratorVariable : enclosingForGeneratorVariables) {
+      var auxiliarySlot = frameDescriptor.findOrAddAuxiliarySlot(forGeneratorVariable);
+      auxiliarySlots[i] = auxiliarySlot;
+      i++;
+    }
+    var expr = visitExpr(ctx);
+    return new GeneratorIterableNode(
+        language,
+        frameDescriptor,
+        auxiliarySlots,
+        createSourceSection(ctx),
+        symbolTable.getCurrentScope().getQualifiedName(),
+        expr);
+  }
+
   @Override
   public GeneratorForNode visitForGenerator(ForGeneratorContext ctx) {
     checkClosingDelimiter(ctx.err, ")", ctx.e.stop);
@@ -976,6 +996,7 @@ public GeneratorForNode visitForGenerator(ForGeneratorContext ctx) {
     var ignoreT1 = isIgnored(ctx.t1);
     var ignoreT2 = ctx.t2 == null ? ignoreT1 : isIgnored(ctx.t2);
 
+    var iterableNode = visitGeneratorIterable(ctx.e);
     if (ctx.t2 != null) {
       keyVariableSlot = ignoreT1 ? -1 : pushForGeneratorVariableContext(ctx.t1);
       valueVariableSlot = ignoreT2 ? -1 : pushForGeneratorVariableContext(ctx.t2);
@@ -991,7 +1012,6 @@ public GeneratorForNode visitForGenerator(ForGeneratorContext ctx) {
           ignoreT1 ? null : visitTypeAnnotation(ctx.t1.typedIdentifier().typeAnnotation());
     }
 
-    var iterableNode = visitExpr(ctx.e);
     var memberNodes = doVisitForWhenBody(ctx.objectBody());
     if (keyVariableSlot != -1) {
       currentScope.popForGeneratorVariable();

pkl-core/src/main/java/org/pkl/core/ast/expression/generator/GeneratorForNode.java

@@ -23,7 +23,6 @@
 import com.oracle.truffle.api.nodes.ExplodeLoop;
 import com.oracle.truffle.api.source.SourceSection;
 import java.util.*;
-import org.pkl.core.ast.ExpressionNode;
 import org.pkl.core.ast.type.TypeNode;
 import org.pkl.core.ast.type.UnresolvedTypeNode;
 import org.pkl.core.ast.type.VmTypeMismatchException;
@@ -32,10 +31,9 @@
 import org.pkl.core.util.Nullable;
 import org.pkl.core.util.Pair;
 
-public abstract class GeneratorForNode extends GeneratorMemberNode {
+public abstract class GeneratorForNode extends GeneratorIteratingNode {
   private final int keySlot;
   private final int valueSlot;
-  @Child private ExpressionNode iterableNode;
   @Child private @Nullable UnresolvedTypeNode unresolvedKeyTypeNode;
   @Child private @Nullable UnresolvedTypeNode unresolvedValueTypeNode;
   @Children private final GeneratorMemberNode[] childNodes;
@@ -46,14 +44,14 @@ public GeneratorForNode(
       SourceSection sourceSection,
       int keySlot,
       int valueSlot,
-      ExpressionNode iterableNode,
+      GeneratorIterableNode iterableNode,
       @Nullable UnresolvedTypeNode unresolvedKeyTypeNode,
       @Nullable UnresolvedTypeNode unresolvedValueTypeNode,
       GeneratorMemberNode[] childNodes,
       boolean hasKeyIdentifier,
       boolean hasValueIdentifier) {
 
-    super(sourceSection);
+    super(sourceSection, iterableNode);
     this.keySlot = keySlot;
     this.valueSlot = valueSlot;
     this.iterableNode = iterableNode;
@@ -79,7 +77,7 @@ protected abstract void executeWithIterable(
 
   @Override
   public final void execute(VirtualFrame frame, Object parent, ObjectData data) {
-    executeWithIterable(frame, parent, data, iterableNode.executeGeneric(frame));
+    executeWithIterable(frame, parent, data, evalIterable(frame, data));
   }
 
   @Specialization

pkl-core/src/main/java/org/pkl/core/ast/expression/generator/GeneratorIterableNode.java

@@ -0,0 +1,80 @@
+/*
+ * Copyright © 2024 Apple Inc. and the Pkl project authors. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.pkl.core.ast.expression.generator;
+
+import com.oracle.truffle.api.frame.FrameDescriptor;
+import com.oracle.truffle.api.frame.VirtualFrame;
+import com.oracle.truffle.api.source.SourceSection;
+import org.pkl.core.ast.ExpressionNode;
+import org.pkl.core.ast.PklRootNode;
+import org.pkl.core.ast.frame.WriteNthFrameSlotNode;
+import org.pkl.core.ast.frame.WriteNthFrameSlotNodeGen;
+import org.pkl.core.runtime.VmLanguage;
+import org.pkl.core.util.Nullable;
+
+/** The iterable node of a for generator, or spread syntax. */
+public class GeneratorIterableNode extends PklRootNode {
+
+  private final int[] auxiliarySlots;
+  private final SourceSection sourceSection;
+  private final String qualifiedName;
+  @Child private ExpressionNode bodyNode;
+  @Child private WriteNthFrameSlotNode writeFrameSlotNode = WriteNthFrameSlotNodeGen.create();
+
+  public GeneratorIterableNode(
+      VmLanguage language,
+      FrameDescriptor descriptor,
+      int[] auxiliarySlots,
+      SourceSection sourceSection,
+      String qualifiedName,
+      ExpressionNode bodyNode) {
+
+    super(language, descriptor);
+    this.auxiliarySlots = auxiliarySlots;
+    this.sourceSection = sourceSection;
+    this.qualifiedName = qualifiedName;
+    this.bodyNode = bodyNode;
+  }
+
+  @Override
+  public SourceSection getSourceSection() {
+    return sourceSection;
+  }
+
+  @Override
+  public @Nullable String getName() {
+    return qualifiedName;
+  }
+
+  @Override
+  public Object execute(VirtualFrame frame) {
+    var forGeneratorVariables = (Object[]) frame.getArguments()[2];
+    assert forGeneratorVariables.length == auxiliarySlots.length;
+    for (var i = 0; i < auxiliarySlots.length; i++) {
+      frame.setAuxiliarySlot(auxiliarySlots[i], forGeneratorVariables[i]);
+    }
+    var frameSlotValues = (Object[]) frame.getArguments()[3];
+    for (var i = 0; i < frameSlotValues.length; i++) {
+      if (i >= frame.getFrameDescriptor().getNumberOfSlots()) {
+        break;
+      }
+      if (frameSlotValues[i] != null) {
+        writeFrameSlotNode.executeWithSlotAndValue(frame, i, frameSlotValues[i]);
+      }
+    }
+    return bodyNode.executeGeneric(frame);
+  }
+}

pkl-core/src/main/java/org/pkl/core/ast/expression/generator/GeneratorIteratingNode.java

@@ -0,0 +1,63 @@
+/*
+ * Copyright © 2024 Apple Inc. and the Pkl project authors. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.pkl.core.ast.expression.generator;
+
+import com.oracle.truffle.api.CompilerDirectives;
+import com.oracle.truffle.api.frame.FrameSlotKind;
+import com.oracle.truffle.api.frame.VirtualFrame;
+import com.oracle.truffle.api.nodes.IndirectCallNode;
+import com.oracle.truffle.api.source.SourceSection;
+import org.pkl.core.runtime.VmUtils;
+import org.pkl.core.util.LateInit;
+
+public abstract class GeneratorIteratingNode extends GeneratorMemberNode {
+
+  @Child protected GeneratorIterableNode iterableNode;
+  @Child @LateInit private IndirectCallNode callNode;
+
+  protected GeneratorIteratingNode(
+      SourceSection sourceSection, GeneratorIterableNode iterableNode) {
+    super(sourceSection);
+    this.iterableNode = iterableNode;
+  }
+
+  protected Object evalIterable(VirtualFrame frame, ObjectData data) {
+    var currentForVariables = data.getCurrentForBindings();
+    if (currentForVariables == null) {
+      currentForVariables = new Object[0];
+    }
+    var arguments = new Object[4];
+    arguments[0] = VmUtils.getReceiver(frame);
+    arguments[1] = VmUtils.getOwner(frame);
+    arguments[2] = currentForVariables;
+    var frameSlotValues = new Object[frame.getFrameDescriptor().getNumberOfSlots()];
+    for (var i = 0; i < frame.getFrameDescriptor().getNumberOfSlots(); i++) {
+      if (frame.getTag(i) != FrameSlotKind.Illegal.tag) {
+        frameSlotValues[i] = frame.getValue(i);
+      }
+    }
+    arguments[3] = frameSlotValues;
+    return getCallNode().call(iterableNode.getCallTarget(), arguments);
+  }
+
+  private IndirectCallNode getCallNode() {
+    if (callNode == null) {
+      CompilerDirectives.transferToInterpreterAndInvalidate();
+      callNode = IndirectCallNode.create();
+    }
+    return callNode;
+  }
+}

pkl-core/src/main/java/org/pkl/core/ast/expression/generator/GeneratorMemberNode.java

@@ -96,6 +96,10 @@ public static final class ObjectData {
 
     private Object @Nullable [] currentForBindings;
 
+    public Object @Nullable [] getCurrentForBindings() {
+      return currentForBindings;
+    }
+
     @TruffleBoundary
     Object @Nullable [] addForBinding(Object value) {
       var result = currentForBindings;

pkl-core/src/main/java/org/pkl/core/ast/expression/generator/GeneratorSpreadNode.java

@@ -25,7 +25,6 @@
 import com.oracle.truffle.api.dsl.Specialization;
 import com.oracle.truffle.api.frame.VirtualFrame;
 import com.oracle.truffle.api.source.SourceSection;
-import org.pkl.core.ast.ExpressionNode;
 import org.pkl.core.ast.member.ObjectMember;
 import org.pkl.core.runtime.BaseModule;
 import org.pkl.core.runtime.Identifier;
@@ -45,14 +44,12 @@
 import org.pkl.core.util.MutableLong;
 
 @ImportStatic(BaseModule.class)
-public abstract class GeneratorSpreadNode extends GeneratorMemberNode {
-  @Child private ExpressionNode iterableNode;
+public abstract class GeneratorSpreadNode extends GeneratorIteratingNode {
   private final boolean nullable;
 
   public GeneratorSpreadNode(
-      SourceSection sourceSection, ExpressionNode iterableNode, boolean nullable) {
-    super(sourceSection);
-    this.iterableNode = iterableNode;
+      SourceSection sourceSection, GeneratorIterableNode iterableNode, boolean nullable) {
+    super(sourceSection, iterableNode);
     this.nullable = nullable;
   }
 
@@ -61,7 +58,7 @@ protected abstract void executeWithIterable(
 
   @Override
   public final void execute(VirtualFrame frame, Object parent, ObjectData data) {
-    executeWithIterable(frame, parent, data, iterableNode.executeGeneric(frame));
+    executeWithIterable(frame, parent, data, evalIterable(frame, data));
   }
 
   @Specialization

pkl-core/src/main/java/org/pkl/core/ast/frame/WriteNthFrameSlotNode.java

@@ -0,0 +1,65 @@
+/*
+ * Copyright © 2024 Apple Inc. and the Pkl project authors. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.pkl.core.ast.frame;
+
+import com.oracle.truffle.api.dsl.Specialization;
+import com.oracle.truffle.api.frame.FrameSlotKind;
+import com.oracle.truffle.api.frame.VirtualFrame;
+import org.pkl.core.ast.PklNode;
+
+public abstract class WriteNthFrameSlotNode extends PklNode {
+
+  public abstract void executeWithSlotAndValue(VirtualFrame frame, int slot, Object value);
+
+  @Specialization(guards = "isBooleanOrIllegal(frame, slot)")
+  protected void evalBoolean(VirtualFrame frame, int slot, boolean value) {
+    frame.getFrameDescriptor().setSlotKind(slot, FrameSlotKind.Boolean);
+    frame.setBoolean(slot, value);
+  }
+
+  @Specialization(guards = "isIntOrIllegal(frame, slot)")
+  protected void evalInt(VirtualFrame frame, int slot, int value) {
+    frame.getFrameDescriptor().setSlotKind(slot, FrameSlotKind.Long);
+    frame.setLong(slot, value);
+  }
+
+  @Specialization(guards = "isFloatOrIllegal(frame, slot)")
+  protected void evalFloat(VirtualFrame frame, int slot, double value) {
+    frame.getFrameDescriptor().setSlotKind(slot, FrameSlotKind.Double);
+    frame.setDouble(slot, value);
+  }
+
+  @Specialization(replaces = {"evalInt", "evalFloat", "evalBoolean"})
+  protected void evalGeneric(VirtualFrame frame, int slot, Object value) {
+    frame.getFrameDescriptor().setSlotKind(slot, FrameSlotKind.Object);
+    frame.setObject(slot, value);
+  }
+
+  protected final boolean isIntOrIllegal(VirtualFrame frame, int slot) {
+    var kind = frame.getFrameDescriptor().getSlotKind(slot);
+    return kind == FrameSlotKind.Long || kind == FrameSlotKind.Illegal;
+  }
+
+  protected final boolean isFloatOrIllegal(VirtualFrame frame, int slot) {
+    var kind = frame.getFrameDescriptor().getSlotKind(slot);
+    return kind == FrameSlotKind.Double || kind == FrameSlotKind.Illegal;
+  }
+
+  protected final boolean isBooleanOrIllegal(VirtualFrame frame, int slot) {
+    var kind = frame.getFrameDescriptor().getSlotKind(slot);
+    return kind == FrameSlotKind.Boolean || kind == FrameSlotKind.Illegal;
+  }
+}

pkl-core/src/test/files/LanguageSnippetTests/input/generators/forGeneratorNestedReference.pkl

@@ -0,0 +1,31 @@
+res1 {
+  for (i in IntSeq(1, 1)) {
+    for (j in new Dynamic { res { new { i } } }.res) {
+      j
+    }
+  }
+}
+
+res2 {
+  for (i in new Listing { 1 }) {
+    for (j in new Dynamic { res { new { i } } }.res) {
+      j
+    }
+  }
+}
+
+res3 {
+  for (i in new Mapping { ["hi"] = 1 }) {
+    for (j in new Dynamic { res { new { i } } }.res) {
+      j
+    }
+  }
+}
+
+res4 {
+  for (i in List(1)) {
+    for (j in new Dynamic { res { new { i } } }.res) {
+      j
+    }
+  }
+}