Make typechecking more PE friendly (#1871)

This makes various improvements to help truffle do partial
evaluation of code related to typechecking.

* Split `UnionTypeNode` into `UnionTypeNodeExploded` and
 `UnionTypeNodeLooped` to avoid gigantic graphs. This especially helps
  in the case of typealiases, because they are inlined into their usage
  sites.

* Move constraint-failure exception construction in
  `TypeConstraintNode` behind `@TruffleBoundary` helper methods
  instead of `transferToInterpreterAndInvalidate`, since building the
  exception is a rare, off-hot-path operation.

* Move `ClassNode`'s one-time class-initialization logic behind a
  `@TruffleBoundary` method (`initClass`), store `VmLanguage` as a
  field instead of recomputing it via `VmLanguage.get(this)`, and use
  `transferToInterpreterAndInvalidate` since class init never needs to
  re-run on the compiled path.

* Pass `VmLanguage` explicitly into `TypeTestNode` and `ClassNode`
  (via `AstBuilder`) rather than resolving it with
  `VmLanguage.get(this)` on a hot path.

* Fix various issues where the partial evaluator determines that a frame
  should have been materialized but was not.

* Pre-initialize call targets in `ObjectMember.setMemberNode()` and
  `FunctionLiteralNode` so the first invocation doesn't trigger a
  compilation invalidation.

* Fix truffle node tree issues (don't mark `@RootNode` as child,
  deepCopy instead of passing an already adopted node).

* Remove unnecessary interpreter transfers when throwing type mismatch exceptions

pkl-core/src/main/java/org/pkl/core/PType.java

@@ -15,6 +15,7 @@
  */
 package org.pkl.core;
 
+import com.oracle.truffle.api.CompilerDirectives.TruffleBoundary;
 import java.io.Serial;
 import java.io.Serializable;
 import java.util.*;
@@ -115,14 +116,17 @@ public abstract class PType implements Serializable {
       this.typeArguments = typeArguments;
     }
 
+    @TruffleBoundary
     public Class(PClass pClass) {
       this(pClass, List.of());
     }
 
+    @TruffleBoundary
     public Class(PClass pClass, PType typeArgument1) {
       this(pClass, List.of(typeArgument1));
     }
 
+    @TruffleBoundary
     public Class(PClass pClass, PType typeArgument1, PType typeArgument2) {
       this(pClass, List.of(typeArgument1, typeArgument2));
     }

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

@@ -1268,7 +1268,7 @@ public class AstBuilder extends AbstractAstBuilder<Object> {
   @Override
   public ExpressionNode visitTypeCheckExpr(TypeCheckExpr expr) {
     return new TypeTestNode(
-        createSourceSection(expr), visitExpr(expr.getExpr()), visitType(expr.getType()));
+        createSourceSection(expr), language, visitExpr(expr.getExpr()), visitType(expr.getType()));
   }
 
   @Override
@@ -1697,7 +1697,8 @@ public class AstBuilder extends AbstractAstBuilder<Object> {
               supertypeNode,
               moduleProperties,
               unresolvedPropertyNodes,
-              doVisitMethodDefs(mod.getMethods()));
+              doVisitMethodDefs(mod.getMethods()),
+              language);
 
       return new ModuleNode(
           language, moduleInfo.getSourceSection(), moduleInfo.getModuleName(), classNode);
@@ -1930,7 +1931,8 @@ public class AstBuilder extends AbstractAstBuilder<Object> {
                   supertypeNode,
                   EconomicMaps.create(),
                   doVisitClassProperties(properties, propertyNames),
-                  doVisitMethodDefs(methods));
+                  doVisitMethodDefs(methods),
+                  language);
 
           var result =
               new ObjectMember(

pkl-core/src/main/java/org/pkl/core/ast/expression/literal/FunctionLiteralNode.java

@@ -46,6 +46,9 @@ public final class FunctionLiteralNode extends ExpressionNode {
     if (functionNode == null) {
       CompilerDirectives.transferToInterpreterAndInvalidate();
       functionNode = unresolvedFunctionNode.execute(frame);
+      // pre-initialize the function node's call target so that the first call to this function
+      // won't invalidate again.
+      functionNode.getCallTarget();
       if (isCustomThisScope) {
         customThisSlot = VmUtils.findCustomThisSlot(frame);
       }

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

@@ -71,6 +71,9 @@ public final class ExecuteCustomThisWithRootNode extends ExpressionNode {
   @Override
   public Object executeGeneric(VirtualFrame frame) {
     var customThis = customThisNode.executeGeneric(frame);
-    return callNode.call(VmUtils.getReceiver(frame), VmUtils.getOwner(frame), customThis, frame);
+    // can't pass a non-materialized frame as an argument into a call node (truffle compiler will
+    // error)
+    return callNode.call(
+        VmUtils.getReceiver(frame), VmUtils.getOwner(frame), customThis, frame.materialize());
   }
 }

pkl-core/src/main/java/org/pkl/core/ast/expression/unary/ImportGlobNode.java

@@ -40,7 +40,8 @@ import org.pkl.core.util.GlobResolver.InvalidGlobPatternException;
 @NodeInfo(shortName = "import*")
 public class ImportGlobNode extends AbstractImportNode {
   private final String globPattern;
-  @Child private @Nullable SharedMemberNode memberNode;
+  // Not a @Child: a RootNode must never be adopted as a child (it always has a null parent).
+  private @Nullable SharedMemberNode memberNode;
   @CompilationFinal private @Nullable VmMapping cachedResult;
 
   public ImportGlobNode(

pkl-core/src/main/java/org/pkl/core/ast/expression/unary/ReadGlobNode.java

@@ -41,7 +41,7 @@ import org.pkl.core.util.IoUtils;
 @NodeInfo(shortName = "read*")
 public abstract class ReadGlobNode extends AbstractReadNode {
   private final EconomicMap<String, VmMapping> cachedResults = EconomicMap.create();
-  @Child private @Nullable SharedMemberNode memberNode;
+  private @Nullable SharedMemberNode memberNode;
 
   protected ReadGlobNode(SourceSection sourceSection, ModuleKey currentModule) {
     super(sourceSection, currentModule);

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

@@ -48,6 +48,7 @@ public final class ClassNode extends ExpressionNode {
   @Children private final UnresolvedMethodNode[] unresolvedMethodNodes;
 
   @CompilationFinal @LateInit private VmClass cachedClass;
+  private final VmLanguage language;
 
   public ClassNode(
       SourceSection section,
@@ -61,7 +62,8 @@ public final class ClassNode extends ExpressionNode {
       @Nullable UnresolvedTypeNode unresolvedSupertypeNode,
       EconomicMap<Object, ObjectMember> prototypeMembers,
       UnresolvedPropertyNode[] unresolvedPropertyNodes,
-      UnresolvedMethodNode[] unresolvedMethodNodes) {
+      UnresolvedMethodNode[] unresolvedMethodNodes,
+      VmLanguage language) {
 
     super(section);
     this.headerSection = headerSection;
@@ -75,6 +77,7 @@ public final class ClassNode extends ExpressionNode {
     this.prototypeMembers = prototypeMembers;
     this.unresolvedPropertyNodes = unresolvedPropertyNodes;
     this.unresolvedMethodNodes = unresolvedMethodNodes;
+    this.language = language;
   }
 
   @Override
@@ -88,7 +91,6 @@ public final class ClassNode extends ExpressionNode {
     if (cachedClass != null) return cachedClass;
 
     CompilerDirectives.transferToInterpreter();
-
     var module = VmUtils.getTypedObjectReceiver(frame);
 
     VmTyped prototype;
@@ -123,12 +125,12 @@ public final class ClassNode extends ExpressionNode {
             typeParameters,
             prototype);
 
-    var localContext = VmLanguage.get(this).localContext.get();
+    var localContext = language.localContext.get();
     localContext.beginClassInit(cachedClass);
 
     try {
       if (unresolvedSupertypeNode != null) {
-        var supertypeNode = unresolvedSupertypeNode.execute(frame);
+        var supertypeNode = insert(unresolvedSupertypeNode.execute(frame));
         var superclass = supertypeNode.getType().getVmClass();
 
         checkSupertype(supertypeNode, superclass);

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

@@ -47,9 +47,12 @@ public final class LocalTypedPropertyNode extends RegularMemberNode {
   public @Nullable Object getDefaultValue(VirtualFrame frame) {
     if (!defaultValueInitialized) {
       assert typeNode != null;
+      // typeNode.createDefaultValue() is a regular polymorphic call rather than a child node.
+      // passing a non-materialized frame will cause GraalVM to emit diagnostics that it failed
+      // to compile this call.
       defaultValue =
           typeNode.createDefaultValue(
-              frame, language, member.getHeaderSection(), member.getQualifiedName());
+              frame.materialize(), language, member.getHeaderSection(), member.getQualifiedName());
       defaultValueInitialized = true;
     }
     return defaultValue;

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

@@ -58,6 +58,9 @@ public final class ObjectMember extends Member {
     assert memberNode == null;
 
     memberNode = node;
+    // pre-initialize the call target so that later calls to `getCallTarget()` does not invalidate
+    // the caller's compiled code
+    node.getCallTarget();
   }
 
   /**

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

@@ -28,6 +28,7 @@ import org.pkl.core.runtime.*;
 
 public final class PropertyTypeNode extends PklRootNode {
   private final String qualifiedPropertyName;
+  private final VmLanguage language;
   @Child private TypeNode typeNode;
 
   private @Nullable Object defaultValue;
@@ -43,6 +44,7 @@ public final class PropertyTypeNode extends PklRootNode {
     super(language, descriptor);
     this.qualifiedPropertyName = qualifiedPropertyName;
     this.typeNode = childNode;
+    this.language = language;
   }
 
   public TypeNode getTypeNode() {
@@ -65,15 +67,18 @@ public final class PropertyTypeNode extends PklRootNode {
   }
 
   public @Nullable Object getDefaultValue(VirtualFrame frame) {
-    if (defaultValueInitialized) return defaultValue;
-
-    defaultValue =
-        typeNode.createDefaultValue(
-            frame, VmLanguage.get(this), getSourceSection(), qualifiedPropertyName);
-    // can't cache default value for `module` type in a non-final module because it's a self-type
-    // (the default value changes when inherited).
-    if (typeNode.isFinalType()) {
-      defaultValueInitialized = true;
+    if (!defaultValueInitialized) {
+      // typeNode.createDefaultValue() is a regular polymorphic call rather than a child node.
+      // passing a non-materialized frame will cause GraalVM to emit diagnostics that it failed
+      // to compile this call.
+      defaultValue =
+          typeNode.createDefaultValue(
+              frame.materialize(), language, getSourceSection(), qualifiedPropertyName);
+      // can't cache default value for `module` type in a non-final module because it's a self-type
+      // (the default value changes when inherited).
+      if (typeNode.isFinalType()) {
+        defaultValueInitialized = true;
+      }
     }
     return defaultValue;
   }

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

@@ -17,10 +17,12 @@ package org.pkl.core.ast.type;
 
 import com.oracle.truffle.api.CompilerDirectives;
 import com.oracle.truffle.api.CompilerDirectives.CompilationFinal;
+import com.oracle.truffle.api.CompilerDirectives.TruffleBoundary;
 import com.oracle.truffle.api.dsl.Cached;
 import com.oracle.truffle.api.dsl.Fallback;
 import com.oracle.truffle.api.dsl.NodeChild;
 import com.oracle.truffle.api.dsl.Specialization;
+import com.oracle.truffle.api.frame.MaterializedFrame;
 import com.oracle.truffle.api.frame.VirtualFrame;
 import com.oracle.truffle.api.source.SourceSection;
 import org.pkl.core.ast.ExpressionNode;
@@ -54,28 +56,32 @@ public abstract class TypeConstraintNode extends PklNode {
     initConstraintSlot(frame);
 
     if (!result) {
-      CompilerDirectives.transferToInterpreterAndInvalidate();
-      var vmContext = VmContext.get(this);
-      var localContext = VmLanguage.get(this).localContext.get();
-      // Use power assertions if enabled and not in type test or already instrumenting.
-      // This prevents `is` checks from triggering instrumentation, but allows them to
-      // participate if instrumentation is already active.
-      var usePowerAssertions =
-          vmContext.getPowerAssertionsEnabled()
-              && (!localContext.isInTypeTest() || localContext.hasActiveTracker());
-      if (usePowerAssertions) {
-        try (var valueTracker = vmContext.getValueTrackerFactory().create()) {
-          getBodyNode().executeGeneric(frame);
-          throw new VmTypeMismatchException.Constraint(
-              sourceSection,
-              frame.getAuxiliarySlot(customThisSlot),
-              sourceSection,
-              valueTracker.values());
-        }
-      } else {
-        throw new VmTypeMismatchException.Constraint(
-            sourceSection, frame.getAuxiliarySlot(customThisSlot), sourceSection, null);
+      throw typeMismatchException(frame.materialize());
+    }
+  }
+
+  @TruffleBoundary
+  private VmTypeMismatchException typeMismatchException(MaterializedFrame frame) {
+    var vmContext = VmContext.get(this);
+    var localContext = VmLanguage.get(this).localContext.get();
+    // Use power assertions if enabled and not in type test or already instrumenting.
+    // This prevents `is` checks from triggering instrumentation, but allows them to
+    // participate if instrumentation is already active.
+    var usePowerAssertions =
+        vmContext.getPowerAssertionsEnabled()
+            && (!localContext.isInTypeTest() || localContext.hasActiveTracker());
+    if (usePowerAssertions) {
+      try (var valueTracker = vmContext.getValueTrackerFactory().create()) {
+        getBodyNode().executeGeneric(frame);
+        return new VmTypeMismatchException.Constraint(
+            sourceSection,
+            frame.getAuxiliarySlot(customThisSlot),
+            sourceSection,
+            valueTracker.values());
       }
+    } else {
+      return new VmTypeMismatchException.Constraint(
+          sourceSection, frame.getAuxiliarySlot(customThisSlot), sourceSection, null);
     }
   }
 
@@ -89,28 +95,30 @@ public abstract class TypeConstraintNode extends PklNode {
     var value = frame.getAuxiliarySlot(customThisSlot);
     var result = applyNode.executeBoolean(function, value);
     if (!result) {
-      CompilerDirectives.transferToInterpreterAndInvalidate();
-      var vmContext = VmContext.get(this);
-      var localContext = VmLanguage.get(this).localContext.get();
-      // Use power assertions if enabled and not in type test or already instrumenting.
-      // This prevents `is` checks from triggering instrumentation, but allows them to
-      // participate if instrumentation is already active.
-      var usePowerAssertions =
-          vmContext.getPowerAssertionsEnabled()
-              && (!localContext.isInTypeTest() || localContext.hasActiveTracker());
-      if (usePowerAssertions) {
-        try (var valueTracker = vmContext.getValueTrackerFactory().create()) {
-          applyNode.executeBoolean(function, value);
-          throw new VmTypeMismatchException.Constraint(
-              sourceSection,
-              value,
-              function.getRootNode().getSourceSection(),
-              valueTracker.values());
-        }
-      } else {
-        throw new VmTypeMismatchException.Constraint(
-            sourceSection, value, function.getRootNode().getSourceSection(), null);
+      throw typeMismatchException(function, value, applyNode);
+    }
+  }
+
+  @TruffleBoundary
+  private VmTypeMismatchException typeMismatchException(
+      VmFunction function, Object value, ApplyVmFunction1Node applyNode) {
+    var vmContext = VmContext.get(this);
+    var localContext = VmLanguage.get(this).localContext.get();
+    // Use power assertions if enabled and not in type test or already instrumenting.
+    // This prevents `is` checks from triggering instrumentation, but allows them to
+    // participate if instrumentation is already active.
+    var usePowerAssertions =
+        vmContext.getPowerAssertionsEnabled()
+            && (!localContext.isInTypeTest() || localContext.hasActiveTracker());
+    if (usePowerAssertions) {
+      try (var valueTracker = vmContext.getValueTrackerFactory().create()) {
+        applyNode.executeBoolean(function, value);
+        return new VmTypeMismatchException.Constraint(
+            sourceSection, value, function.getRootNode().getSourceSection(), valueTracker.values());
       }
+    } else {
+      return new VmTypeMismatchException.Constraint(
+          sourceSection, value, function.getRootNode().getSourceSection(), null);
     }
   }
 

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

@@ -24,11 +24,13 @@ import com.oracle.truffle.api.dsl.Idempotent;
 import com.oracle.truffle.api.dsl.Specialization;
 import com.oracle.truffle.api.frame.FrameDescriptor;
 import com.oracle.truffle.api.frame.FrameSlotKind;
+import com.oracle.truffle.api.frame.MaterializedFrame;
 import com.oracle.truffle.api.frame.VirtualFrame;
 import com.oracle.truffle.api.nodes.ExplodeLoop;
 import com.oracle.truffle.api.nodes.LoopNode;
 import com.oracle.truffle.api.nodes.Node;
 import com.oracle.truffle.api.nodes.NodeUtil;
+import com.oracle.truffle.api.profiles.LoopConditionProfile;
 import com.oracle.truffle.api.source.SourceSection;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -200,6 +202,7 @@ public abstract class TypeNode extends PklNode {
   /** Visit child type nodes of this type. */
   protected abstract boolean acceptTypeNode(boolean visitTypeArguments, TypeNodeConsumer consumer);
 
+  @TruffleBoundary
   protected VmTypeMismatchException constraintException(Object value, SourceSection sourceSection) {
     throw new VmTypeMismatchException.Constraint(
         sourceSection,
@@ -376,7 +379,6 @@ public abstract class TypeNode extends PklNode {
 
     @Override
     protected Object executeLazily(VirtualFrame frame, Object value) {
-      CompilerDirectives.transferToInterpreter();
       throw new VmTypeMismatchException.Nothing(sourceSection, value);
     }
 
@@ -570,6 +572,10 @@ public abstract class TypeNode extends PklNode {
         String qualifiedName) {
       return literal;
     }
+
+    public String getLiteral() {
+      return literal;
+    }
   }
 
   public static final class TypedTypeNode extends ObjectSlotTypeNode {
@@ -798,7 +804,7 @@ public abstract class TypeNode extends PklNode {
     }
   }
 
-  public static class UnionTypeNode extends WriteFrameSlotTypeNode {
+  public abstract static class UnionTypeNode extends WriteFrameSlotTypeNode {
     @Children final TypeNode[] elementTypeNodes;
     private final int defaultIndex;
 
@@ -820,16 +826,16 @@ public abstract class TypeNode extends PklNode {
       return MirrorFactories.unionTypeFactory.create(this);
     }
 
-    public VmList getElementTypeMirrors() {
+    public final VmList getElementTypeMirrors() {
       return getMirrors(elementTypeNodes);
     }
 
-    public TypeNode[] getElementTypeNodes() {
+    public final TypeNode[] getElementTypeNodes() {
       return elementTypeNodes;
     }
 
     @Override
-    public boolean isNoopTypeCheck() {
+    public final boolean isNoopTypeCheck() {
       for (var element : elementTypeNodes) {
         if (!element.isNoopTypeCheck()) return false;
       }
@@ -837,7 +843,7 @@ public abstract class TypeNode extends PklNode {
     }
 
     @Override
-    public @Nullable Object createDefaultValue(
+    public final @Nullable Object createDefaultValue(
         VirtualFrame frame,
         VmLanguage language,
         SourceSection headerSection,
@@ -849,27 +855,6 @@ public abstract class TypeNode extends PklNode {
               frame, language, headerSection, qualifiedName);
     }
 
-    @Override
-    @ExplodeLoop
-    protected boolean acceptTypeNode(boolean visitTypeArguments, TypeNodeConsumer consumer) {
-      if (!consumer.accept(this)) {
-        return false;
-      }
-      var ret = true;
-      // don't break early to ensure constant number of iterations
-      //noinspection ForLoopReplaceableByForEach
-      for (var i = 0; i < elementTypeNodes.length; i++) {
-        if (!ret) {
-          continue;
-        }
-        if (!elementTypeNodes[i].acceptTypeNode(visitTypeArguments, consumer)) {
-          ret = false;
-        }
-      }
-      LoopNode.reportLoopCount(this, elementTypeNodes.length);
-      return ret;
-    }
-
     /**
      * Tells if the union type should be eagerly checked or not (shallow-force members of
      * Listing/Mapping).
@@ -878,33 +863,120 @@ public abstract class TypeNode extends PklNode {
      * type; e.g. {@code Listing<Person>|Listing<Animal>}
      */
     @TruffleBoundary
-    private boolean shouldEagerCheck() {
+    protected final boolean shouldEagerCheck() {
       var seenParameterizedClasses = EconomicSets.<VmClass>create();
       var ret = new MutableBoolean(false);
-      acceptTypeNode(
+      this.acceptTypeNode(
           false,
           (typeNode) -> {
             if (!typeNode.getType().isParametric()) {
               return true;
             }
-            var typeClass = typeNode.getType().getVmClass();
-            if (typeClass == null) {
+            var typeNodeClass = typeNode.getType().getVmClass();
+            if (typeNodeClass == null) {
               return true;
             }
-            if (seenParameterizedClasses.contains(typeClass)) {
+            if (seenParameterizedClasses.contains(typeNodeClass)) {
               ret.set(true);
               return false;
             } else {
-              EconomicSets.add(seenParameterizedClasses, typeClass);
+              EconomicSets.add(seenParameterizedClasses, typeNodeClass);
               return true;
             }
           });
       return ret.get();
     }
 
-    @Fallback
+    // All members failed to type check; if enabled, re-execute type checks to generate power
+    // assertions.
+    //
+    // Intentionally no `@ExplodeLoop` because this is the error path (prevent error path code from
+    // exploding the compiled IR node count).
+    //
+    // Needs a `MaterializedFrame` (rather than VirtualFrame): this method is called from within an
+    // `@ExplodeLoop` caller, and passing a VirtualFrame into a separately-compiled method called
+    // from a loop makes Graal's escape analysis unable to prove the frame can stay virtual across
+    // the call, causing a permanent "must not let virtual object escape" bailout.
+    //
+    // Materializing here is fine; this only runs once all union members have already failed to type
+    // check.
+    protected final VmTypeMismatchException.Union computeLazyUnionMismatch(
+        MaterializedFrame frame,
+        Object value,
+        VmTypeMismatchException[] typeMismatches,
+        boolean shouldEagerCheck,
+        VmLocalContext localContext,
+        boolean wasInTypeTest) {
+      localContext.setInTypeTest(wasInTypeTest);
+      if (VmContext.get(this).getPowerAssertionsEnabled()
+          && (!wasInTypeTest || localContext.hasActiveTracker())) {
+        for (var i = 0; i < elementTypeNodes.length; i++) {
+          var elementTypeNode = elementTypeNodes[i];
+          try {
+            if (shouldEagerCheck) {
+              elementTypeNode.executeEagerly(frame, value);
+            } else {
+              elementTypeNode.executeLazily(frame, value);
+            }
+          } catch (VmTypeMismatchException e) {
+            typeMismatches[i] = e;
+          }
+        }
+      }
+      return new VmTypeMismatchException.Union(sourceSection, value, this, typeMismatches);
+    }
+
+    // All members failed to type check; if enabled, re-execute type checks to generate power
+    // assertions.
+    //
+    // Intentionally no `@ExplodeLoop` because this is the error path (prevent error path code from
+    // exploding the compiled IR node count).
+    //
+    // Needs a `MaterializedFrame`, see comment on `computeLazyUnionMismatch` for details.
+    protected final VmTypeMismatchException.Union computeEagerUnionMismatch(
+        MaterializedFrame frame,
+        Object value,
+        VmTypeMismatchException[] typeMismatches,
+        VmLocalContext localContext,
+        boolean wasInTypeTest) {
+      localContext.setInTypeTest(wasInTypeTest);
+      if (VmContext.get(this).getPowerAssertionsEnabled()
+          && (!wasInTypeTest || localContext.hasActiveTracker())) {
+        for (var i = 0; i < elementTypeNodes.length; i++) {
+          try {
+            elementTypeNodes[i].executeEagerly(frame, value);
+          } catch (VmTypeMismatchException e) {
+            typeMismatches[i] = e;
+          }
+        }
+      }
+      return new VmTypeMismatchException.Union(sourceSection, value, this, typeMismatches);
+    }
+  }
+
+  public static class UnionTypeNodeExploded extends UnionTypeNode {
+    public UnionTypeNodeExploded(
+        SourceSection sourceSection, int defaultIndex, TypeNode[] elementTypeNodes) {
+      super(sourceSection, defaultIndex, elementTypeNodes);
+    }
+
+    @Override
     @ExplodeLoop
-    protected Object executeLazily(VirtualFrame frame, Object value) {
+    protected boolean acceptTypeNode(boolean visitTypeArguments, TypeNodeConsumer consumer) {
+      if (!consumer.accept(this)) {
+        return false;
+      }
+      for (var elementTypeNode : elementTypeNodes) {
+        if (!elementTypeNode.acceptTypeNode(visitTypeArguments, consumer)) {
+          return false;
+        }
+      }
+      return true;
+    }
+
+    @SuppressWarnings("DuplicatedCode")
+    @ExplodeLoop
+    protected final Object executeLazily(VirtualFrame frame, Object value) {
       // escape analysis should remove this allocation in compiled code
       var typeMismatches = new VmTypeMismatchException[elementTypeNodes.length];
 
@@ -931,30 +1003,17 @@ public abstract class TypeNode extends PklNode {
         }
       }
 
-      // all members failed to type check
-      // if enabled, re-execute type checks to generate power assertions
-      localContext.setInTypeTest(wasInTypeTest);
-      if (VmContext.get(this).getPowerAssertionsEnabled()
-          && (!wasInTypeTest || localContext.hasActiveTracker())) {
-        for (var i = 0; i < elementTypeNodes.length; i++) {
-          var elementTypeNode = elementTypeNodes[i];
-          try {
-            if (shouldEagerCheck) {
-              elementTypeNode.executeEagerly(frame, value);
-            } else {
-              elementTypeNode.executeLazily(frame, value);
-            }
-          } catch (VmTypeMismatchException e) {
-            typeMismatches[i] = e;
-          }
-        }
-      }
-
-      throw new VmTypeMismatchException.Union(sourceSection, value, this, typeMismatches);
+      throw computeLazyUnionMismatch(
+          frame.materialize(),
+          value,
+          typeMismatches,
+          shouldEagerCheck,
+          localContext,
+          wasInTypeTest);
     }
 
-    @Override
-    public Object executeEagerly(VirtualFrame frame, Object value) {
+    @ExplodeLoop
+    public final Object executeEagerly(VirtualFrame frame, Object value) {
       // escape analysis should remove this allocation in compiled code
       var typeMismatches = new VmTypeMismatchException[elementTypeNodes.length];
 
@@ -974,21 +1033,91 @@ public abstract class TypeNode extends PklNode {
         }
       }
 
-      // all members failed to type check
-      // if enabled, re-execute type checks to generate power assertions
-      localContext.setInTypeTest(wasInTypeTest);
-      if (VmContext.get(this).getPowerAssertionsEnabled()
-          && (!wasInTypeTest || localContext.hasActiveTracker())) {
-        for (var i = 0; i < elementTypeNodes.length; i++) {
-          try {
-            elementTypeNodes[i].executeEagerly(frame, value);
-          } catch (VmTypeMismatchException e) {
-            typeMismatches[i] = e;
-          }
+      throw computeEagerUnionMismatch(
+          frame.materialize(), value, typeMismatches, localContext, wasInTypeTest);
+    }
+  }
+
+  public static class UnionTypeNodeLooped extends UnionTypeNode {
+    private final LoopConditionProfile loopConditionProfile = LoopConditionProfile.create();
+
+    public UnionTypeNodeLooped(
+        SourceSection sourceSection, int defaultIndex, TypeNode[] elementTypeNodes) {
+      super(sourceSection, defaultIndex, elementTypeNodes);
+    }
+
+    // keep in sync with UnionTypeNodeExploded.acceptTypeNode
+    @Override
+    protected boolean acceptTypeNode(boolean visitTypeArguments, TypeNodeConsumer consumer) {
+      if (!consumer.accept(this)) {
+        return false;
+      }
+      for (var i = 0; loopConditionProfile.inject(i < elementTypeNodes.length); i++) {
+        if (!elementTypeNodes[i].acceptTypeNode(visitTypeArguments, consumer)) {
+          return false;
+        }
+      }
+      return true;
+    }
+
+    @SuppressWarnings("DuplicatedCode")
+    @Override
+    protected Object executeLazily(VirtualFrame frame, Object value) {
+      // escape analysis should remove this allocation in compiled code
+      var typeMismatches = new VmTypeMismatchException[elementTypeNodes.length];
+      var materializedFrame = frame.materialize();
+
+      // disallow power assertions from triggering in case one union member checks successfully
+      var localContext = VmLanguage.get(this).localContext.get();
+      var wasInTypeTest = localContext.isInTypeTest();
+      localContext.setInTypeTest(true);
+
+      // Do eager checks (shallow-force) if there are two listings or two mappings represented.
+      // (we can't know that `new Listing { 0; "hi" }[0]` fails for `Listing<Int>|Listing<String>`
+      // without checking both index 0 and index 1).
+      var shouldEagerCheck = shouldEagerCheck();
+      for (var i = 0; loopConditionProfile.inject(i < elementTypeNodes.length); i++) {
+        var elementTypeNode = elementTypeNodes[i];
+        try {
+          var result =
+              shouldEagerCheck
+                  ? elementTypeNode.executeEagerly(materializedFrame, value)
+                  : elementTypeNode.executeLazily(materializedFrame, value);
+          localContext.setInTypeTest(wasInTypeTest);
+          return result;
+        } catch (VmTypeMismatchException e) {
+          typeMismatches[i] = e;
         }
       }
 
-      throw new VmTypeMismatchException.Union(sourceSection, value, this, typeMismatches);
+      throw computeLazyUnionMismatch(
+          materializedFrame, value, typeMismatches, shouldEagerCheck, localContext, wasInTypeTest);
+    }
+
+    @Override
+    public Object executeEagerly(VirtualFrame frame, Object value) {
+      // escape analysis should remove this allocation in compiled code
+      var typeMismatches = new VmTypeMismatchException[elementTypeNodes.length];
+      var materializedFrame = frame.materialize();
+
+      // disallow power assertions from triggering in case one union member checks successfully
+      var localContext = VmLanguage.get(this).localContext.get();
+      var wasInTypeTest = localContext.isInTypeTest();
+      localContext.setInTypeTest(true);
+
+      for (var i = 0; loopConditionProfile.inject(i < elementTypeNodes.length); i++) {
+        // eager checks
+        try {
+          var result = elementTypeNodes[i].executeEagerly(materializedFrame, value);
+          localContext.setInTypeTest(wasInTypeTest);
+          return result;
+        } catch (VmTypeMismatchException e) {
+          typeMismatches[i] = e;
+        }
+      }
+
+      throw computeEagerUnionMismatch(
+          materializedFrame, value, typeMismatches, localContext, wasInTypeTest);
     }
   }
 
@@ -1001,8 +1130,8 @@ public abstract class TypeNode extends PklNode {
         SourceSection sourceSection, int defaultIndex, Set<String> stringLiterals) {
       super(sourceSection);
       assert !stringLiterals.isEmpty();
-      this.stringLiterals = stringLiterals;
       this.defaultIndex = defaultIndex;
+      this.stringLiterals = stringLiterals;
       if (defaultIndex == -1) {
         unionDefault = null;
       } else {
@@ -1054,6 +1183,14 @@ public abstract class TypeNode extends PklNode {
         String qualifiedName) {
       return unionDefault;
     }
+
+    public int getDefaultIndex() {
+      return defaultIndex;
+    }
+
+    public Set<String> getStringLiterals() {
+      return stringLiterals;
+    }
   }
 
   public static final class CollectionTypeNode extends ObjectSlotTypeNode {
@@ -1207,7 +1344,6 @@ public abstract class TypeNode extends PklNode {
 
     @SuppressWarnings("DuplicatedCode")
     @Override
-    @ExplodeLoop
     protected Object executeLazily(VirtualFrame frame, Object value) {
       if (!(value instanceof VmList vmList)) {
         throw typeMismatch(value, BaseModule.getListClass());
@@ -1493,7 +1629,10 @@ public abstract class TypeNode extends PklNode {
     private final VmLanguage language;
     @Child protected @Nullable TypeNode keyTypeNode;
     @Child protected TypeNode valueTypeNode;
-    @Child @Nullable protected ListingOrMappingTypeCastNode valueTypeCastNode;
+    // Not a @Child: a RootNode must never be adopted as a child (it always has a null parent).
+    // It is only ever invoked indirectly, via an IndirectCallNode held by the VmListing/VmMapping
+    // values this type node produces (see ElementOrEntryNode).
+    @Nullable protected ListingOrMappingTypeCastNode valueTypeCastNode;
 
     protected ListingOrMappingTypeNode(
         SourceSection sourceSection,
@@ -1521,9 +1660,14 @@ public abstract class TypeNode extends PklNode {
     protected ListingOrMappingTypeCastNode getValueTypeCastNode(FrameDescriptor descriptor) {
       if (valueTypeCastNode == null) {
         CompilerDirectives.transferToInterpreterAndInvalidate();
+        // `valueTypeNode` is already an adopted @Child of `this`, and is also executed directly
+        // by `this` (see executeEagerly()/executeLazily() below). It must not be reused as a
+        // @Child of a second, unrelated root node -- a node can only ever have one parent, and
+        // doing so corrupts `valueTypeNode.getRootNode()`, which breaks the compiler's ability to
+        // constant-fold context/language lookups reached through it.
         valueTypeCastNode =
             new ListingOrMappingTypeCastNode(
-                language, descriptor, valueTypeNode, getRootNode().getName());
+                language, descriptor, (TypeNode) valueTypeNode.deepCopy(), getRootNode().getName());
       }
       return valueTypeCastNode;
     }
@@ -1656,10 +1800,7 @@ public abstract class TypeNode extends PklNode {
             try {
               keyTypeNode.executeEagerly(frame, memberKey);
             } catch (VmTypeMismatchException e) {
-              CompilerDirectives.transferToInterpreter();
-              e.putInsertedStackFrame(
-                  getRootNode().getCallTarget(),
-                  VmUtils.createStackFrame(member.getHeaderSection(), member.getQualifiedName()));
+              putStackFrame(member, e);
               throw e;
             }
           }
@@ -1681,6 +1822,13 @@ public abstract class TypeNode extends PklNode {
 
       LoopNode.reportLoopCount(this, loopCount);
     }
+
+    @TruffleBoundary
+    private void putStackFrame(ObjectMember member, VmTypeMismatchException e) {
+      e.putInsertedStackFrame(
+          getRootNode().getCallTarget(),
+          VmUtils.createStackFrame(member.getHeaderSection(), member.getQualifiedName()));
+    }
   }
 
   // A type such as `(Int, String) -> Duration`.
@@ -2057,7 +2205,6 @@ public abstract class TypeNode extends PklNode {
     @Override
     protected Object executeLazily(VirtualFrame frame, Object value) {
       if (value instanceof VmNull) {
-        CompilerDirectives.transferToInterpreterAndInvalidate();
         throw constraintException(value, BaseModule.getNonNullTypeAlias().getConstraintSection());
       }
       return value;
@@ -2095,7 +2242,6 @@ public abstract class TypeNode extends PklNode {
       if (value instanceof Long l) {
         if ((l & mask) == l) return value;
 
-        CompilerDirectives.transferToInterpreterAndInvalidate();
         var sourceSection = typeAlias.getConstraintSection();
         throw constraintException(value, sourceSection);
       }
@@ -2136,7 +2282,6 @@ public abstract class TypeNode extends PklNode {
       if (value instanceof Long l) {
         if (l == l.byteValue()) return value;
 
-        CompilerDirectives.transferToInterpreterAndInvalidate();
         var sourceSection = BaseModule.getInt8TypeAlias().getConstraintSection();
         throw constraintException(value, sourceSection);
       }
@@ -2171,7 +2316,6 @@ public abstract class TypeNode extends PklNode {
       if (value instanceof Long l) {
         if (l == l.shortValue()) return value;
 
-        CompilerDirectives.transferToInterpreterAndInvalidate();
         var sourceSection = BaseModule.getInt16TypeAlias().getConstraintSection();
         throw constraintException(value, sourceSection);
       }
@@ -2206,11 +2350,9 @@ public abstract class TypeNode extends PklNode {
       if (value instanceof Long l) {
         if (l == l.intValue()) return value;
 
-        CompilerDirectives.transferToInterpreterAndInvalidate();
         throw constraintException(value, BaseModule.getInt32TypeAlias().getConstraintSection());
       }
 
-      CompilerDirectives.transferToInterpreterAndInvalidate();
       throw new VmTypeMismatchException.Simple(
           BaseModule.getInt32TypeAlias().getBaseTypeSection(), value, BaseModule.getIntClass());
     }
@@ -2348,6 +2490,10 @@ public abstract class TypeNode extends PklNode {
     protected boolean acceptTypeNode(boolean visitTypeArguments, TypeNodeConsumer consumer) {
       return consumer.accept(this) && aliasedTypeNode.acceptTypeNode(visitTypeArguments, consumer);
     }
+
+    public TypeNode getAliasedTypeNode() {
+      return aliasedTypeNode;
+    }
   }
 
   public static final class ConstrainedTypeNode extends TypeNode {
@@ -2696,12 +2842,10 @@ public abstract class TypeNode extends PklNode {
 
       // clazz will be null iff the type arg is a not a valid class type
       if (clazz == null) {
-        CompilerDirectives.transferToInterpreter();
         throw new VmTypeMismatchException.ClassType(sourceSection, value, getType());
       }
 
       if (!value.isSubclassOf(clazz)) {
-        CompilerDirectives.transferToInterpreter();
         throw new VmTypeMismatchException.ClassType(sourceSection, value, clazz);
       }
 

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

@@ -28,12 +28,15 @@ public final class TypeTestNode extends ExpressionNode {
   @Child private ExpressionNode valueNode;
   @Child private @Nullable UnresolvedTypeNode unresolvedTypeNode;
   @Child private @Nullable TypeNode typeNode;
+  private final VmLanguage language;
 
   public TypeTestNode(
       SourceSection sourceSection,
+      VmLanguage language,
       ExpressionNode valueNode,
       UnresolvedTypeNode unresolvedTypeNode) {
     super(sourceSection);
+    this.language = language;
     this.valueNode = valueNode;
     this.unresolvedTypeNode = unresolvedTypeNode;
   }
@@ -57,8 +60,8 @@ public final class TypeTestNode extends ExpressionNode {
     // TODO: throw if typeNode is FunctionTypeNode (it's impossible to check)
     // https://github.com/apple/pkl/issues/639
     Object value = valueNode.executeGeneric(frame);
-    var localContext = VmLanguage.get(this).localContext.get();
-    boolean wasInTypeTest = localContext.isInTypeTest();
+    var localContext = language.localContext.get();
+    var wasInTypeTest = localContext.isInTypeTest();
     localContext.setInTypeTest(true);
     try {
       typeNode.executeEagerly(frame, value);

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

@@ -364,6 +364,8 @@ public abstract class UnresolvedTypeNode extends PklNode {
   }
 
   public static final class Union extends UnresolvedTypeNode {
+    private static final int MAX_EXPLODE_LOOP = 20;
+
     @Children private final UnresolvedTypeNode[] unresolvedElementTypeNodes;
     private final int defaultIndex;
 
@@ -386,7 +388,15 @@ public abstract class UnresolvedTypeNode extends PklNode {
         elementTypeNodes[i] = elementTypeNode;
       }
 
-      return new UnionTypeNode(sourceSection, defaultIndex, elementTypeNodes);
+      // Don't use `@ExplodeLoop` variant of union type of union is too large;
+      // this can easily turn into too many IR nodes; the compiler will spend too much time
+      // compiling these code paths.
+      //
+      // This is especially true for typealiases of unions, because they get inlined into their
+      // usage sites.
+      return elementTypeNodes.length > MAX_EXPLODE_LOOP
+          ? new UnionTypeNodeLooped(sourceSection, defaultIndex, elementTypeNodes)
+          : new UnionTypeNodeExploded(sourceSection, defaultIndex, elementTypeNodes);
     }
   }
 

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

@@ -175,12 +175,14 @@ public abstract class VmTypeMismatchException extends ControlFlowException {
     private final String renderedExpected;
     private final @Nullable VmClass expectedClass;
 
+    @TruffleBoundary
     public ClassType(SourceSection sourceSection, VmClass actualClass, VmClass expectedClass) {
       super(sourceSection, actualClass);
       this.expectedClass = expectedClass;
       renderedExpected = "Class<" + expectedClass.toPklString() + ">";
     }
 
+    @TruffleBoundary
     public ClassType(SourceSection sourceSection, VmClass actualClass, VmType expectedType) {
       super(sourceSection, actualClass);
       this.expectedClass = null;

pkl-core/src/main/java/org/pkl/core/stdlib/protobuf/RendererNodes.java

@@ -43,6 +43,7 @@ import org.pkl.core.ast.type.TypeNode.StringTypeNode;
 import org.pkl.core.ast.type.TypeNode.TypeAliasTypeNode;
 import org.pkl.core.ast.type.TypeNode.UnionOfStringLiteralsTypeNode;
 import org.pkl.core.ast.type.TypeNode.UnionTypeNode;
+import org.pkl.core.ast.type.TypeNode.UnionTypeNodeExploded;
 import org.pkl.core.ast.type.VmTypeMismatchException;
 import org.pkl.core.runtime.Identifier;
 import org.pkl.core.runtime.VmBytes;
@@ -648,7 +649,7 @@ public final class RendererNodes {
         type =
             hasString && elements.size() == 1
                 ? elements.get(0)
-                : new UnionTypeNode(
+                : new UnionTypeNodeExploded(
                     VmUtils.unavailableSourceSection(), -1, elements.toArray(new TypeNode[0]));
         return type;
       }