Improve lazy type checking of listings and mappings (#789)

Motivation:
- simplify implementation of lazy type checking
- fix correctness issues of lazy type checking (#785)

Changes:
- implement listing/mapping type cast via amendment (`parent`) instead of delegation (`delegate`)
- handle type checking of *computed* elements/entries in the same way as type checking of computed properties
  - ElementOrEntryNode is the equivalent of TypeCheckedPropertyNode
- remove fields VmListingOrMapping.delegate/typeNodeFrame/cachedMembers/checkedMembers
- fix #785 by executing all type casts between a member's owner and receiver
- fix #823 by storing owner and receiver directly
  instead of storing the mutable frame containing them (typeNodeFrame)
- remove overrides of VmObject methods that are no longer required
  - good for Truffle partial evaluation and JVM inlining
- revert a85a173 except for added tests
- move `VmUtils.setOwner` and `VmUtils.setReceiver` and make them private
  - these methods aren't generally safe to use

Result:
- simpler code with greater optimization potential
  - VmListingOrMapping can now have both a type node and new members
- fewer changes to surrounding code
- smaller memory footprint
- better performance in some cases
- fixes #785
- fixes #823

Potential future optimizations:
- avoid lazy type checking overhead for untyped listings/mappings
- improve efficiency of forcing a typed listing/mapping
  - currently, lazy type checking will traverse the parent chain once per member,
    reducing the performance benefit of shallow-forcing
	  a listing/mapping over evaluating each member individually
- avoid creating an intermediate untyped listing/mapping in the following cases:
  - `new Listing<X> {...}`
  - amendment of `property: Listing<X>`

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

@@ -1221,7 +1221,7 @@ private ObjectMember doVisitObjectElement(ObjectElementContext ctx) {
             member.initConstantValue(constantNode);
           } else {
             member.initMemberNode(
-                new UntypedObjectMemberNode(
+                ElementOrEntryNodeGen.create(
                     language, scope.buildFrameDescriptor(), member, elementNode));
           }
 
@@ -1278,7 +1278,7 @@ private Function<EntryScope, Pair<ExpressionNode, ObjectMember>> objectMemberIns
           member.initConstantValue(constantNode);
         } else {
           member.initMemberNode(
-              new UntypedObjectMemberNode(
+              ElementOrEntryNodeGen.create(
                   language, scope.buildFrameDescriptor(), member, valueNode));
         }
       } else { // ["key"] { ... }
@@ -1287,7 +1287,7 @@ private Function<EntryScope, Pair<ExpressionNode, ObjectMember>> objectMemberIns
                 objectBodyCtxs,
                 new ReadSuperEntryNode(unavailableSourceSection(), new GetMemberKeyNode()));
         member.initMemberNode(
-            new UntypedObjectMemberNode(
+            ElementOrEntryNodeGen.create(
                 language, scope.buildFrameDescriptor(), member, objectBody));
       }
 
@@ -2446,6 +2446,7 @@ public UnresolvedTypeNode visitDeclaredType(DeclaredTypeContext ctx) {
 
     return new UnresolvedTypeNode.Parameterized(
         createSourceSection(ctx),
+        language,
         doVisitTypeName(idCtx),
         argCtx.ts.stream().map(this::visitType).toArray(UnresolvedTypeNode[]::new));
   }

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

@@ -103,7 +103,7 @@ private void addMembers(VirtualFrame frame, VmObject parent, ObjectData data) {
             var callTarget = member.getCallTarget();
             value = callTarget.call(parent, owner, key);
           }
-          owner.setCachedValue(key, value, member);
+          owner.setCachedValue(key, value);
         }
 
         frame.setAuxiliarySlot(customThisSlot, value);

pkl-core/src/main/java/org/pkl/core/ast/expression/member/ReadLocalPropertyNode.java

@@ -71,7 +71,7 @@ public Object executeGeneric(VirtualFrame frame) {
 
     if (result == null) {
       result = callNode.call(objReceiver, owner, property.getName());
-      objReceiver.setCachedValue(property, result, property);
+      objReceiver.setCachedValue(property, result);
     }
 
     return result;

pkl-core/src/main/java/org/pkl/core/ast/expression/primary/ResolveVariableNode.java

@@ -184,12 +184,12 @@ private ExpressionNode doResolve(VirtualFrame frame) {
       if (member != null) {
         var constantValue = member.getConstantValue();
         if (constantValue != null) {
-          baseModule.setCachedValue(variableName, constantValue, member);
+          baseModule.setCachedValue(variableName, constantValue);
           return new ConstantValueNode(sourceSection, constantValue);
         }
 
         var computedValue = member.getCallTarget().call(baseModule, baseModule);
-        baseModule.setCachedValue(variableName, computedValue, member);
+        baseModule.setCachedValue(variableName, computedValue);
         return new ConstantValueNode(sourceSection, computedValue);
       }
     }

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

@@ -0,0 +1,73 @@
+/*
+ * 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.member;
+
+import com.oracle.truffle.api.dsl.Cached;
+import com.oracle.truffle.api.dsl.Cached.Shared;
+import com.oracle.truffle.api.dsl.Executed;
+import com.oracle.truffle.api.dsl.Specialization;
+import com.oracle.truffle.api.frame.FrameDescriptor;
+import com.oracle.truffle.api.frame.VirtualFrame;
+import com.oracle.truffle.api.nodes.IndirectCallNode;
+import org.pkl.core.ast.ExpressionNode;
+import org.pkl.core.ast.expression.primary.GetReceiverNode;
+import org.pkl.core.runtime.VmDynamic;
+import org.pkl.core.runtime.VmLanguage;
+import org.pkl.core.runtime.VmListing;
+import org.pkl.core.runtime.VmMapping;
+import org.pkl.core.runtime.VmUtils;
+import org.pkl.core.util.Nullable;
+
+/** Equivalent of {@link TypedPropertyNode} for elements/entries. */
+public abstract class ElementOrEntryNode extends RegularMemberNode {
+  @Child @Executed protected ExpressionNode receiverNode = new GetReceiverNode();
+
+  protected ElementOrEntryNode(
+      @Nullable VmLanguage language,
+      FrameDescriptor descriptor,
+      ObjectMember member,
+      ExpressionNode bodyNode) {
+
+    super(language, descriptor, member, bodyNode);
+  }
+
+  @Specialization
+  protected Object evalListing(
+      VirtualFrame frame,
+      VmListing receiver,
+      @Cached("create()") @Shared("callNode") IndirectCallNode callNode) {
+    var result = executeBody(frame);
+    return VmUtils.shouldRunTypeCheck(frame)
+        ? receiver.executeTypeCasts(result, VmUtils.getOwner(frame), callNode, null, null)
+        : result;
+  }
+
+  @Specialization
+  protected Object evalMapping(
+      VirtualFrame frame,
+      VmMapping receiver,
+      @Cached("create()") @Shared("callNode") IndirectCallNode callNode) {
+    var result = executeBody(frame);
+    return VmUtils.shouldRunTypeCheck(frame)
+        ? receiver.executeTypeCasts(result, VmUtils.getOwner(frame), callNode, null, null)
+        : result;
+  }
+
+  @Specialization
+  protected Object evalDynamic(VirtualFrame frame, VmDynamic ignored) {
+    return executeBody(frame);
+  }
+}

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

@@ -26,7 +26,7 @@
 import org.pkl.core.util.Nullable;
 
 /** Performs a typecast on a Mapping entry value, or a Listing element. */
-public class ListingOrMappingTypeCastNode extends PklRootNode {
+public final class ListingOrMappingTypeCastNode extends PklRootNode {
 
   @Child private TypeNode typeNode;
   private final String qualifiedName;

pkl-core/src/main/java/org/pkl/core/ast/type/ResolveDeclaredTypeNode.java

@@ -73,7 +73,7 @@ protected VmTyped getImport(
     var result = module.getCachedValue(importName);
     if (result == null) {
       result = callNode.call(member.getCallTarget(), module, module, importName);
-      module.setCachedValue(importName, result, member);
+      module.setCachedValue(importName, result);
     }
     return (VmTyped) result;
   }
@@ -94,7 +94,7 @@ protected VmTyped getImport(
     var result = module.getCachedValue(typeName);
     if (result == null) {
       result = callNode.call(member.getCallTarget(), module, module, typeName);
-      module.setCachedValue(typeName, result, member);
+      module.setCachedValue(typeName, result);
     }
     return result;
   }

pkl-core/src/main/java/org/pkl/core/ast/type/TypeNode.java

@@ -21,6 +21,7 @@
 import com.oracle.truffle.api.dsl.Cached;
 import com.oracle.truffle.api.dsl.Fallback;
 import com.oracle.truffle.api.dsl.Specialization;
+import com.oracle.truffle.api.frame.Frame;
 import com.oracle.truffle.api.frame.FrameDescriptor;
 import com.oracle.truffle.api.frame.FrameSlotKind;
 import com.oracle.truffle.api.frame.VirtualFrame;
@@ -1418,16 +1419,27 @@ protected boolean isParametric() {
   }
 
   public static final class ListingTypeNode extends ListingOrMappingTypeNode {
-    public ListingTypeNode(SourceSection sourceSection, TypeNode valueTypeNode) {
-      super(sourceSection, null, valueTypeNode);
+    public ListingTypeNode(
+        SourceSection sourceSection, VmLanguage language, TypeNode valueTypeNode) {
+      super(sourceSection, language, null, valueTypeNode);
     }
 
     @Override
     public Object execute(VirtualFrame frame, Object value) {
       if (!(value instanceof VmListing vmListing)) {
         throw typeMismatch(value, BaseModule.getListingClass());
       }
-      return doTypeCast(frame, vmListing);
+      if (vmListing.isValueTypeKnownSubtypeOf(valueTypeNode)) {
+        return vmListing;
+      }
+      return new VmListing(
+          vmListing.getEnclosingFrame(),
+          vmListing,
+          EconomicMaps.emptyMap(),
+          vmListing.getLength(),
+          getValueTypeCastNode(),
+          VmUtils.getReceiver(frame),
+          VmUtils.getOwner(frame));
     }
 
     @Override
@@ -1470,9 +1482,12 @@ protected boolean acceptTypeNode(TypeNodeConsumer consumer) {
 
   public static final class MappingTypeNode extends ListingOrMappingTypeNode {
     public MappingTypeNode(
-        SourceSection sourceSection, TypeNode keyTypeNode, TypeNode valueTypeNode) {
+        SourceSection sourceSection,
+        VmLanguage language,
+        TypeNode keyTypeNode,
+        TypeNode valueTypeNode) {
 
-      super(sourceSection, keyTypeNode, valueTypeNode);
+      super(sourceSection, language, keyTypeNode, valueTypeNode);
     }
 
     @Override
@@ -1482,7 +1497,16 @@ public Object execute(VirtualFrame frame, Object value) {
       }
       // execute type checks on mapping keys
       doEagerCheck(frame, vmMapping, false, true);
-      return doTypeCast(frame, vmMapping);
+      if (vmMapping.isValueTypeKnownSubtypeOf(valueTypeNode)) {
+        return vmMapping;
+      }
+      return new VmMapping(
+          vmMapping.getEnclosingFrame(),
+          vmMapping,
+          EconomicMaps.emptyMap(),
+          getValueTypeCastNode(),
+          VmUtils.getReceiver(frame),
+          VmUtils.getOwner(frame));
     }
 
     @Override
@@ -1530,17 +1554,22 @@ protected boolean acceptTypeNode(TypeNodeConsumer consumer) {
   }
 
   public abstract static class ListingOrMappingTypeNode extends ObjectSlotTypeNode {
+    private final VmLanguage language;
     @Child protected @Nullable TypeNode keyTypeNode;
     @Child protected TypeNode valueTypeNode;
-    @Child @Nullable protected ListingOrMappingTypeCastNode listingOrMappingTypeCastNode;
+    @Child @Nullable protected ListingOrMappingTypeCastNode valueTypeCastNode;
 
     private final boolean skipKeyTypeChecks;
     private final boolean skipValueTypeChecks;
 
     protected ListingOrMappingTypeNode(
-        SourceSection sourceSection, @Nullable TypeNode keyTypeNode, TypeNode valueTypeNode) {
+        SourceSection sourceSection,
+        VmLanguage language,
+        @Nullable TypeNode keyTypeNode,
+        TypeNode valueTypeNode) {
 
       super(sourceSection);
+      this.language = language;
       this.keyTypeNode = keyTypeNode;
       this.valueTypeNode = valueTypeNode;
 
@@ -1560,17 +1589,14 @@ public TypeNode getValueTypeNode() {
       return valueTypeNode;
     }
 
-    protected ListingOrMappingTypeCastNode getListingOrMappingTypeCastNode() {
-      if (listingOrMappingTypeCastNode == null) {
+    protected ListingOrMappingTypeCastNode getValueTypeCastNode() {
+      if (valueTypeCastNode == null) {
         CompilerDirectives.transferToInterpreterAndInvalidate();
-        listingOrMappingTypeCastNode =
+        valueTypeCastNode =
             new ListingOrMappingTypeCastNode(
-                VmLanguage.get(this),
-                getRootNode().getFrameDescriptor(),
-                valueTypeNode,
-                getRootNode().getName());
+                language, new FrameDescriptor(), valueTypeNode, getRootNode().getName());
       }
-      return listingOrMappingTypeCastNode;
+      return valueTypeCastNode;
     }
 
     // either (if defaultMemberValue != null):
@@ -1651,15 +1677,6 @@ public final Object createDefaultValue(
           EconomicMaps.of(Identifier.DEFAULT, defaultMember));
     }
 
-    protected <T extends VmListingOrMapping<T>> T doTypeCast(VirtualFrame frame, T original) {
-      // optimization: don't create new object if the original already has the same typecheck, or if
-      // this typecheck is a no-op.
-      if (isNoopTypeCheck() || original.hasSameChecksAs(valueTypeNode)) {
-        return original;
-      }
-      return original.withCheckedMembers(getListingOrMappingTypeCastNode(), frame.materialize());
-    }
-
     protected void doEagerCheck(VirtualFrame frame, VmObject object) {
       doEagerCheck(frame, object, skipKeyTypeChecks, skipValueTypeChecks);
     }
@@ -1704,7 +1721,7 @@ protected void doEagerCheck(
                 var callTarget = member.getCallTarget();
                 memberValue = callTarget.call(object, owner, memberKey);
               }
-              object.setCachedValue(memberKey, memberValue, member);
+              object.setCachedValue(memberKey, memberValue);
             }
             valueTypeNode.executeEagerly(frame, memberValue);
           }
@@ -2391,14 +2408,14 @@ public VmList getTypeArgumentMirrors() {
     public Object execute(VirtualFrame frame, Object value) {
       var prevOwner = VmUtils.getOwner(frame);
       var prevReceiver = VmUtils.getReceiver(frame);
-      VmUtils.setOwner(frame, VmUtils.getOwner(typeAlias.getEnclosingFrame()));
-      VmUtils.setReceiver(frame, VmUtils.getReceiver(typeAlias.getEnclosingFrame()));
+      setOwner(frame, VmUtils.getOwner(typeAlias.getEnclosingFrame()));
+      setReceiver(frame, VmUtils.getReceiver(typeAlias.getEnclosingFrame()));
 
       try {
         return aliasedTypeNode.execute(frame, value);
       } finally {
-        VmUtils.setOwner(frame, prevOwner);
-        VmUtils.setReceiver(frame, prevReceiver);
+        setOwner(frame, prevOwner);
+        setReceiver(frame, prevReceiver);
       }
     }
 
@@ -2407,14 +2424,14 @@ public Object execute(VirtualFrame frame, Object value) {
     public Object executeAndSet(VirtualFrame frame, Object value) {
       var prevOwner = VmUtils.getOwner(frame);
       var prevReceiver = VmUtils.getReceiver(frame);
-      VmUtils.setOwner(frame, VmUtils.getOwner(typeAlias.getEnclosingFrame()));
-      VmUtils.setReceiver(frame, VmUtils.getReceiver(typeAlias.getEnclosingFrame()));
+      setOwner(frame, VmUtils.getOwner(typeAlias.getEnclosingFrame()));
+      setReceiver(frame, VmUtils.getReceiver(typeAlias.getEnclosingFrame()));
 
       try {
         return aliasedTypeNode.executeAndSet(frame, value);
       } finally {
-        VmUtils.setOwner(frame, prevOwner);
-        VmUtils.setReceiver(frame, prevReceiver);
+        setOwner(frame, prevOwner);
+        setReceiver(frame, prevReceiver);
       }
     }
 
@@ -2423,14 +2440,14 @@ public Object executeAndSet(VirtualFrame frame, Object value) {
     public Object executeEagerlyAndSet(VirtualFrame frame, Object value) {
       var prevOwner = VmUtils.getOwner(frame);
       var prevReceiver = VmUtils.getReceiver(frame);
-      VmUtils.setOwner(frame, VmUtils.getOwner(typeAlias.getEnclosingFrame()));
-      VmUtils.setReceiver(frame, VmUtils.getReceiver(typeAlias.getEnclosingFrame()));
+      setOwner(frame, VmUtils.getOwner(typeAlias.getEnclosingFrame()));
+      setReceiver(frame, VmUtils.getReceiver(typeAlias.getEnclosingFrame()));
 
       try {
         return aliasedTypeNode.executeEagerlyAndSet(frame, value);
       } finally {
-        VmUtils.setOwner(frame, prevOwner);
-        VmUtils.setReceiver(frame, prevReceiver);
+        setOwner(frame, prevOwner);
+        setReceiver(frame, prevReceiver);
       }
     }
 
@@ -2502,6 +2519,22 @@ protected boolean acceptTypeNode(TypeNodeConsumer consumer) {
     protected boolean isParametric() {
       return typeArgumentNodes.length > 0;
     }
+
+    // Note that mutating a frame's receiver and owner argument is very risky
+    // because any VmObject instantiated within the same root node execution
+    // holds a reference to (not immutable snapshot of) the frame
+    // via VmObjectLike.enclosingFrame.
+    // *Maybe* this works out for TypeAliasTypeNode because an object instantiated
+    // within a type constraint doesn't escape the constraint expression.
+    // If mutating receiver and owner can't be avoided, it would be safer
+    // to have VmObjectLike store them directly instead of storing enclosingFrame.
+    private static void setReceiver(Frame frame, Object receiver) {
+      frame.getArguments()[0] = receiver;
+    }
+
+    private static void setOwner(Frame frame, VmObjectLike owner) {
+      frame.getArguments()[1] = owner;
+    }
   }
 
   public static final class ConstrainedTypeNode extends TypeNode {