Eagerly check listing/mapping in iterables (#752)

There is currently a bug around resolving variables within the iterable of a for
generator or spread syntax (https://github.com/apple/pkl/issues/741)

Normally, mappings/listings are type-checked lazily. However, this results in the said
bug getting widened, for any object members declared in the iterable.

As a workaround for now, prevent the bug from being any worse by ensuring that these
object members are eagerly typechecked.

pkl-core/src/main/java/org/pkl/core/ast/VmModifier.java

@@ -55,6 +55,9 @@ public final class VmModifier {
 
   public static final int GLOB = 0x1000;
 
+  // To be removed when https://github.com/apple/pkl/issues/741 is fixed
+  public static final int IS_IN_ITERABLE = 0x100000;
+
   // modifier sets
 
   public static final int NONE = 0;
@@ -134,6 +137,10 @@ public final class VmModifier {
     return (modifiers & ENTRY) != 0;
   }
 
+  public static boolean isInIterable(int modifiers) {
+    return (modifiers & IS_IN_ITERABLE) != 0;
+  }
+
   public static boolean isType(int modifiers) {
     return (modifiers & (CLASS | TYPE_ALIAS | IMPORT)) != 0 && (modifiers & GLOB) == 0;
   }

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

@@ -900,8 +900,12 @@ public final class AstBuilder extends AbstractAstBuilder<Object> {
 
   @Override
   public GeneratorMemberNode visitObjectSpread(ObjectSpreadContext ctx) {
-    return GeneratorSpreadNodeGen.create(
-        createSourceSection(ctx), visitExpr(ctx.expr()), ctx.QSPREAD() != null);
+    var scope = symbolTable.getCurrentScope();
+    var visitingIterable = scope.isVisitingIterable();
+    scope.setVisitingIterable(true);
+    var expr = visitExpr(ctx.expr());
+    scope.setVisitingIterable(visitingIterable);
+    return GeneratorSpreadNodeGen.create(createSourceSection(ctx), expr, ctx.QSPREAD() != null);
   }
 
   private void insertWriteForGeneratorVarsToFrameSlotsNode(@Nullable MemberNode memberNode) {
@@ -992,7 +996,11 @@ public final class AstBuilder extends AbstractAstBuilder<Object> {
           ignoreT1 ? null : visitTypeAnnotation(ctx.t1.typedIdentifier().typeAnnotation());
     }
 
+    var scope = symbolTable.getCurrentScope();
+    var visitingIterable = scope.isVisitingIterable();
+    scope.setVisitingIterable(true);
     var iterableNode = visitExpr(ctx.e);
+    scope.setVisitingIterable(visitingIterable);
     var memberNodes = doVisitForWhenBody(ctx.objectBody());
     if (keyVariableSlot != -1) {
       currentScope.popForGeneratorVariable();
@@ -1197,11 +1205,15 @@ public final class AstBuilder extends AbstractAstBuilder<Object> {
         scope -> {
           var elementNode = visitExpr(ctx.expr());
 
+          var modifier =
+              scope.isVisitingIterable()
+                  ? VmModifier.ELEMENT | VmModifier.IS_IN_ITERABLE
+                  : VmModifier.ELEMENT;
           var member =
               new ObjectMember(
                   createSourceSection(ctx),
                   elementNode.getSourceSection(),
-                  VmModifier.ELEMENT,
+                  modifier,
                   null,
                   scope.getQualifiedName());
 
@@ -1252,13 +1264,13 @@ public final class AstBuilder extends AbstractAstBuilder<Object> {
       @Nullable ExprContext valueCtx,
       List<? extends ObjectBodyContext> objectBodyCtxs) {
     return scope -> {
+      var modifier =
+          scope.isVisitingIterable()
+              ? VmModifier.ENTRY | VmModifier.IS_IN_ITERABLE
+              : VmModifier.ENTRY;
       var member =
           new ObjectMember(
-              sourceSection,
-              keyNode.getSourceSection(),
-              VmModifier.ENTRY,
-              null,
-              scope.getQualifiedName());
+              sourceSection, keyNode.getSourceSection(), modifier, null, scope.getQualifiedName());
 
       if (valueCtx != null) { // ["key"] = value
         var valueNode = visitExpr(valueCtx);
@@ -1338,6 +1350,10 @@ public final class AstBuilder extends AbstractAstBuilder<Object> {
       result += modifier;
     }
 
+    if (symbolTable.getCurrentScope().isVisitingIterable()) {
+      result += VmModifier.IS_IN_ITERABLE;
+    }
+
     // flag modifier combinations that are never valid right away
 
     if (VmModifier.isExternal(result) && !ModuleKeys.isStdLibModule(moduleKey)) {

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

@@ -171,6 +171,7 @@ public final class SymbolTable {
     private int entryCount = 0;
     private final FrameDescriptor.Builder frameDescriptorBuilder;
     private final ConstLevel constLevel;
+    private boolean isVisitingIterable;
 
     private Scope(
         @Nullable Scope parent,
@@ -187,6 +188,7 @@ public final class SymbolTable {
           parent != null && parent.constLevel.biggerOrEquals(constLevel)
               ? parent.constLevel
               : constLevel;
+      this.isVisitingIterable = parent != null && parent.isVisitingIterable;
     }
 
     public final @Nullable Scope getParent() {
@@ -337,6 +339,14 @@ public final class SymbolTable {
     public ConstLevel getConstLevel() {
       return constLevel;
     }
+
+    public void setVisitingIterable(boolean isVisitingIterable) {
+      this.isVisitingIterable = isVisitingIterable;
+    }
+
+    public boolean isVisitingIterable() {
+      return isVisitingIterable;
+    }
   }
 
   private interface LexicalScope {}

pkl-core/src/main/java/org/pkl/core/ast/member/Member.java

@@ -130,4 +130,21 @@ public abstract class Member {
   public final boolean isLocalOrExternalOrAbstract() {
     return VmModifier.isLocalOrExternalOrAbstract(modifiers);
   }
+
+  /**
+   * Tells if this member is declared inside the iterable of a for-generator, or an object spread.
+   * <p>
+   * This is {@code true} for {@code new {}} within:
+   *
+   * <pre>
+   * {@code
+   * for (x in new Listing { new {} }) {
+   *                         ^^^^^^
+   *   // etc
+   * }
+   * </pre>
+   */
+  public boolean isInIterable() {
+    return VmModifier.isInIterable(modifiers);
+  }
 }

pkl-core/src/main/java/org/pkl/core/ast/member/ObjectMember.java

@@ -27,6 +27,7 @@ import org.pkl.core.runtime.VmUtils;
 import org.pkl.core.util.Nullable;
 
 public final class ObjectMember extends Member {
+
   @CompilationFinal private @Nullable Object constantValue;
   @CompilationFinal private @Nullable MemberNode memberNode;
 

pkl-core/src/main/java/org/pkl/core/ast/member/PropertyTypeNode.java

@@ -61,9 +61,25 @@ public final class PropertyTypeNode extends PklRootNode {
     return qualifiedPropertyName;
   }
 
+  private boolean isInIterable(VirtualFrame frame) {
+    var args = frame.getArguments();
+    return args.length >= 4 && args[3] instanceof Boolean b && b;
+  }
+
   @Override
   public Object execute(VirtualFrame frame) {
     try {
+      if (isInIterable(frame)) {
+        // There is currently a bug around resolving variables within the iterable of a for
+        // generator or spread syntax (https://github.com/apple/pkl/issues/741)
+        //
+        // Normally, mappings/listings are type-checked lazily. However, this results in said
+        // bug getting widened, for any object members declared in the iterable.
+        //
+        // As a workaround for now, prevent the bug from being any worse by ensuring that these
+        // object members are eagerly typechecked.
+        return typeNode.executeEagerly(frame, frame.getArguments()[2]);
+      }
       return typeNode.execute(frame, frame.getArguments()[2]);
     } catch (VmTypeMismatchException e) {
       CompilerDirectives.transferToInterpreter();

pkl-core/src/main/java/org/pkl/core/ast/member/TypeCheckedPropertyNode.java

@@ -55,7 +55,8 @@ public abstract class TypeCheckedPropertyNode extends RegularMemberNode {
 
     // TODO: propagate SUPER_CALL_MARKER to disable constraint (but not type) check
     if (callNode != null && VmUtils.shouldRunTypeCheck(frame)) {
-      return callNode.call(VmUtils.getReceiverOrNull(frame), property.getOwner(), result);
+      return callNode.call(
+          VmUtils.getReceiverOrNull(frame), property.getOwner(), result, member.isInIterable());
     }
 
     return result;
@@ -75,7 +76,8 @@ public abstract class TypeCheckedPropertyNode extends RegularMemberNode {
             typeAnnNode.getCallTarget(),
             VmUtils.getReceiverOrNull(frame),
             property.getOwner(),
-            result);
+            result,
+            member.isInIterable());
       }
     }
 

pkl-core/src/main/java/org/pkl/core/ast/member/TypedPropertyNode.java

@@ -47,7 +47,10 @@ public final class TypedPropertyNode extends RegularMemberNode {
     var propertyValue = executeBody(frame);
     if (VmUtils.shouldRunTypeCheck(frame)) {
       return typeCheckCallNode.call(
-          VmUtils.getReceiver(frame), VmUtils.getOwner(frame), propertyValue);
+          VmUtils.getReceiver(frame),
+          VmUtils.getOwner(frame),
+          propertyValue,
+          member.isInIterable());
     }
     return propertyValue;
   }

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

@@ -0,0 +1,33 @@
+class Aviary {
+  birds: Listing<Bird>
+}
+
+class Bird {
+  age: Int
+}
+
+function Bird(_age: Int) = new Bird { age = _age }
+
+res1 {
+  for (i in IntSeq(1, 1)) {
+    for (j in 
+      new Aviary {
+        birds {
+          Bird(i)
+        }
+      }.birds
+    ) {
+      j
+    }
+  }
+}
+
+res2 {
+  for (i in IntSeq(1, 1)) {
+    ...new Aviary {
+      birds {
+        Bird(i)
+      }
+    }.birds
+  }
+}

pkl-core/src/test/files/LanguageSnippetTests/output/generators/forGeneratorNestedReference.pcf

@@ -0,0 +1,10 @@
+res1 {
+  new {
+    age = 1
+  }
+}
+res2 {
+  new {
+    age = 1
+  }
+}