Correctly handle module types as typealias parameters

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

@@ -24,6 +24,7 @@
 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.MaterializedFrame;
 import com.oracle.truffle.api.frame.VirtualFrame;
 import com.oracle.truffle.api.nodes.ExplodeLoop;
 import com.oracle.truffle.api.nodes.LoopNode;
@@ -66,6 +67,10 @@ public interface ClassTypeNode {
     VmClass getVmClass();
   }
 
+  private interface StatefulTypeNode {
+    void invalidate();
+  }
+
   protected TypeNode(SourceSection sourceSection) {
     super(sourceSection);
   }
@@ -465,9 +470,42 @@ protected boolean acceptTypeNode(boolean visitTypeArguments, TypeNodeConsumer co
     }
   }
 
+  public abstract static class SelfTypeNode extends ObjectSlotTypeNode {
+    private @Nullable VmTypeAlias originalAnchor;
+
+    public SelfTypeNode(SourceSection sourceSection) {
+      super(sourceSection);
+    }
+
+    public void setOriginalAnchor(VmTypeAlias originalAnchor) {
+      if (this.originalAnchor == null) {
+        this.originalAnchor = originalAnchor;
+      }
+    }
+
+    protected final VirtualFrame getEffectiveFrame(VirtualFrame frame) {
+      if (originalAnchor == null) return frame;
+
+      var levelsUp = -1;
+      for (var node = getParent(); node != null; node = node.getParent()) {
+        if (node instanceof TypeAliasTypeNode typeAliasTypeNode) {
+          levelsUp++;
+          if (typeAliasTypeNode.getTypeAlias() == originalAnchor) {
+            break;
+          }
+        }
+      }
+      if (levelsUp == -1) {
+        return frame;
+      }
+
+      var localContext = VmLanguage.get(this).localContext.get();
+      return localContext.getRealTypeAliasFrame(levelsUp);
+    }
+  }
+
   /** The `module` type for an open module. */
-  public static final class NonFinalModuleTypeNode extends ObjectSlotTypeNode
-      implements ClassTypeNode {
+  public static final class NonFinalModuleTypeNode extends SelfTypeNode implements ClassTypeNode {
     private final VmClass moduleClass; // only used by getVmClass()
     @Child private ExpressionNode getModuleNode;
 
@@ -479,7 +517,8 @@ public NonFinalModuleTypeNode(SourceSection sourceSection, VmClass moduleClass)
 
     @Override
     protected Object executeLazily(VirtualFrame frame, Object value) {
-      var moduleClass = ((VmTyped) getModuleNode.executeGeneric(frame)).getVmClass();
+      var moduleClass =
+          ((VmTyped) getModuleNode.executeGeneric(getEffectiveFrame(frame))).getVmClass();
 
       if (value instanceof VmTyped typed) {
         var valueClass = typed.getVmClass();
@@ -523,7 +562,8 @@ protected boolean acceptTypeNode(boolean visitTypeArguments, TypeNodeConsumer co
         VmLanguage language,
         SourceSection headerSection,
         String qualifiedName) {
-      var moduleClass = ((VmTyped) getModuleNode.executeGeneric(frame)).getVmClass();
+      var moduleClass =
+          ((VmTyped) getModuleNode.executeGeneric(getEffectiveFrame(frame))).getVmClass();
       return TypeNode.createDefaultValue(moduleClass);
     }
   }
@@ -847,9 +887,9 @@ protected boolean acceptTypeNode(boolean visitTypeArguments, TypeNodeConsumer co
     }
   }
 
-  public static class UnionTypeNode extends WriteFrameSlotTypeNode {
+  public static class UnionTypeNode extends WriteFrameSlotTypeNode implements StatefulTypeNode {
     @Children final TypeNode[] elementTypeNodes;
-    private final boolean skipElementTypeChecks;
+    private boolean skipElementTypeChecks;
     private final int defaultIndex;
 
     public UnionTypeNode(
@@ -864,6 +904,15 @@ public UnionTypeNode(
       this.skipElementTypeChecks = skipElementTypeChecks;
     }
 
+    @Override
+    public void invalidate() {
+      var skipElementTypeChecks = true;
+      for (var i = 0; i < elementTypeNodes.length; i++) {
+        skipElementTypeChecks &= elementTypeNodes[i].isNoopTypeCheck();
+      }
+      this.skipElementTypeChecks = skipElementTypeChecks;
+    }
+
     @Override
     public VmTyped getMirror() {
       return MirrorFactories.unionTypeFactory.create(this);
@@ -1162,7 +1211,6 @@ public static final class CollectionTypeNode extends ObjectSlotTypeNode {
     @Child private TypeNode elementTypeNode;
 
     public CollectionTypeNode(SourceSection sourceSection, TypeNode elementTypeNode) {
-
       super(sourceSection);
       this.elementTypeNode = elementTypeNode;
     }
@@ -1270,16 +1318,21 @@ protected boolean isParametric() {
     }
   }
 
-  public static final class ListTypeNode extends ObjectSlotTypeNode {
+  public static final class ListTypeNode extends ObjectSlotTypeNode implements StatefulTypeNode {
     @Child private TypeNode elementTypeNode;
-    private final boolean skipElementTypeChecks;
+    private boolean skipElementTypeChecks;
 
     public ListTypeNode(SourceSection sourceSection, TypeNode elementTypeNode) {
       super(sourceSection);
       this.elementTypeNode = elementTypeNode;
       skipElementTypeChecks = elementTypeNode.isNoopTypeCheck();
     }
 
+    @Override
+    public void invalidate() {
+      skipElementTypeChecks = elementTypeNode.isNoopTypeCheck();
+    }
+
     @Override
     protected boolean acceptTypeNode(boolean visitTypeArguments, TypeNodeConsumer consumer) {
       if (visitTypeArguments) {
@@ -1369,16 +1422,21 @@ protected boolean isParametric() {
     }
   }
 
-  public abstract static class SetTypeNode extends ObjectSlotTypeNode {
+  public abstract static class SetTypeNode extends ObjectSlotTypeNode implements StatefulTypeNode {
     @Child private TypeNode elementTypeNode;
-    private final boolean skipElementTypeChecks;
+    private boolean skipElementTypeChecks;
 
     protected SetTypeNode(SourceSection sourceSection, TypeNode elementTypeNode) {
       super(sourceSection);
       this.elementTypeNode = elementTypeNode;
       skipElementTypeChecks = elementTypeNode.isNoopTypeCheck();
     }
 
+    @Override
+    public void invalidate() {
+      skipElementTypeChecks = elementTypeNode.isNoopTypeCheck();
+    }
+
     @Override
     public final Object createDefaultValue(
         VirtualFrame frame,
@@ -1448,19 +1506,23 @@ protected boolean isParametric() {
     }
   }
 
-  public static final class MapTypeNode extends ObjectSlotTypeNode {
+  public static final class MapTypeNode extends ObjectSlotTypeNode implements StatefulTypeNode {
     @Child private TypeNode keyTypeNode;
     @Child private TypeNode valueTypeNode;
-    private final boolean skipEntryTypeChecks;
+    private boolean skipEntryTypeChecks;
 
     public MapTypeNode(SourceSection sourceSection, TypeNode keyTypeNode, TypeNode valueTypeNode) {
-
       super(sourceSection);
       this.keyTypeNode = keyTypeNode;
       this.valueTypeNode = valueTypeNode;
       skipEntryTypeChecks = keyTypeNode.isNoopTypeCheck() && valueTypeNode.isNoopTypeCheck();
     }
 
+    @Override
+    public void invalidate() {
+      skipEntryTypeChecks = keyTypeNode.isNoopTypeCheck() && valueTypeNode.isNoopTypeCheck();
+    }
+
     @Override
     protected Object executeLazily(VirtualFrame frame, Object value) {
       if (value instanceof VmMap vmMap) {
@@ -1706,14 +1768,15 @@ protected boolean acceptTypeNode(boolean visitTypeArguments, TypeNodeConsumer co
     }
   }
 
-  public abstract static class ListingOrMappingTypeNode extends ObjectSlotTypeNode {
+  public abstract static class ListingOrMappingTypeNode extends ObjectSlotTypeNode
+      implements StatefulTypeNode {
     private final VmLanguage language;
     @Child protected @Nullable TypeNode keyTypeNode;
     @Child protected TypeNode valueTypeNode;
     @Child @Nullable protected ListingOrMappingTypeCastNode valueTypeCastNode;
 
-    private final boolean skipKeyTypeChecks;
-    private final boolean skipValueTypeChecks;
+    private boolean skipKeyTypeChecks;
+    private boolean skipValueTypeChecks;
 
     protected ListingOrMappingTypeNode(
         SourceSection sourceSection,
@@ -1730,6 +1793,12 @@ protected ListingOrMappingTypeNode(
       skipValueTypeChecks = valueTypeNode.isNoopTypeCheck();
     }
 
+    @Override
+    public void invalidate() {
+      skipKeyTypeChecks = keyTypeNode == null || keyTypeNode.isNoopTypeCheck();
+      skipValueTypeChecks = valueTypeNode.isNoopTypeCheck();
+    }
+
     private boolean isListing() {
       return keyTypeNode == null;
     }
@@ -2762,6 +2831,13 @@ public TypeAliasTypeNode(
       this.typeAlias = typeAlias;
       this.typeArgumentNodes = typeArgumentNodes;
       aliasedTypeNode = typeAlias.instantiate(typeArgumentNodes);
+      aliasedTypeNode.accept(
+          node -> {
+            if (node instanceof StatefulTypeNode typeNode) {
+              typeNode.invalidate();
+            }
+            return true;
+          });
       checkReferentConstraints(typeAlias);
     }
 
@@ -2826,32 +2902,80 @@ public VmList getTypeArgumentMirrors() {
      * where the typealias was declared, so that we preserve its original scope.
      */
     protected Object executeLazily(VirtualFrame frame, Object value) {
+      var localContext = VmLanguage.get(this).localContext.get();
       var prevOwner = VmUtils.getOwner(frame);
       var prevReceiver = VmUtils.getReceiver(frame);
       setOwner(frame, VmUtils.getOwner(typeAlias.getEnclosingFrame()));
       setReceiver(frame, VmUtils.getReceiver(typeAlias.getEnclosingFrame()));
+      localContext.pushRealTypeAliasFrame(new FakeFrame(prevReceiver, prevOwner));
 
       try {
         return aliasedTypeNode.executeLazily(frame, value);
       } finally {
         setOwner(frame, prevOwner);
         setReceiver(frame, prevReceiver);
+        localContext.popRealTypeAliasFrame();
+      }
+    }
+
+    @Override
+    public Object executeEagerly(VirtualFrame frame, Object value) {
+      var localContext = VmLanguage.get(this).localContext.get();
+      var prevOwner = VmUtils.getOwner(frame);
+      var prevReceiver = VmUtils.getReceiver(frame);
+      setOwner(frame, VmUtils.getOwner(typeAlias.getEnclosingFrame()));
+      setReceiver(frame, VmUtils.getReceiver(typeAlias.getEnclosingFrame()));
+      localContext.pushRealTypeAliasFrame(new FakeFrame(prevReceiver, prevOwner));
+
+      try {
+        return aliasedTypeNode.executeEagerly(frame, value);
+      } finally {
+        setOwner(frame, prevOwner);
+        setReceiver(frame, prevReceiver);
+        localContext.popRealTypeAliasFrame();
       }
     }
 
     /** See docstring on {@link TypeAliasTypeNode#executeLazily}. */
     @Override
     public Object executeAndSet(VirtualFrame frame, Object value) {
+      var localContext = VmLanguage.get(this).localContext.get();
       var prevOwner = VmUtils.getOwner(frame);
       var prevReceiver = VmUtils.getReceiver(frame);
       setOwner(frame, VmUtils.getOwner(typeAlias.getEnclosingFrame()));
       setReceiver(frame, VmUtils.getReceiver(typeAlias.getEnclosingFrame()));
+      localContext.pushRealTypeAliasFrame(new FakeFrame(prevReceiver, prevOwner));
 
       try {
         return aliasedTypeNode.executeAndSet(frame, value);
       } finally {
         setOwner(frame, prevOwner);
         setReceiver(frame, prevReceiver);
+        localContext.popRealTypeAliasFrame();
+      }
+    }
+
+    private static final class FakeFrame implements VirtualFrame, MaterializedFrame {
+
+      private final Object[] args;
+
+      public FakeFrame(Object receiver, Object owner) {
+        this.args = new Object[] {receiver, owner};
+      }
+
+      @Override
+      public FrameDescriptor getFrameDescriptor() {
+        throw PklBugException.unreachableCode();
+      }
+
+      @Override
+      public Object[] getArguments() {
+        return args;
+      }
+
+      @Override
+      public MaterializedFrame materialize() {
+        return this;
       }
     }
 

pkl-core/src/main/java/org/pkl/core/runtime/VmLocalContext.java

@@ -15,6 +15,10 @@
  */
 package org.pkl.core.runtime;
 
+import com.oracle.truffle.api.frame.VirtualFrame;
+import java.util.ArrayList;
+import java.util.List;
+
 /** A per-context thread-local value that can be used to influence execution. */
 public class VmLocalContext {
   private boolean shouldEagerTypecheck = false;
@@ -30,6 +34,8 @@ public class VmLocalContext {
 
   private boolean instrumentationEverUsed = false;
 
+  private List<VirtualFrame> realTypeAliasFrames = new ArrayList<>();
+
   public VmLocalContext() {}
 
   public void shouldEagerTypecheck(boolean shouldEagerTypecheck) {
@@ -64,4 +70,16 @@ public boolean hasActiveTracker() {
   public boolean isInstrumentationEverUsed() {
     return instrumentationEverUsed;
   }
+
+  public VirtualFrame getRealTypeAliasFrame(int index) {
+    return realTypeAliasFrames.get(realTypeAliasFrames.size() - 1 - index);
+  }
+
+  public void pushRealTypeAliasFrame(VirtualFrame frame) {
+    realTypeAliasFrames.add(frame);
+  }
+
+  public void popRealTypeAliasFrame() {
+    realTypeAliasFrames.remove(realTypeAliasFrames.size() - 1);
+  }
 }

pkl-core/src/main/java/org/pkl/core/runtime/VmTypeAlias.java

@@ -32,6 +32,7 @@
 import org.pkl.core.ast.VmModifier;
 import org.pkl.core.ast.type.TypeNode;
 import org.pkl.core.ast.type.TypeNode.ConstrainedTypeNode;
+import org.pkl.core.ast.type.TypeNode.SelfTypeNode;
 import org.pkl.core.ast.type.TypeNode.TypeVariableNode;
 import org.pkl.core.ast.type.TypeNode.UnknownTypeNode;
 
@@ -192,10 +193,18 @@ public TypeNode instantiate(TypeNode[] typeArgumentNodes) {
             int index = typeVarNode.getTypeParameterIndex();
             // should not need to clone type argument node because it is not used by its original
             // root node
-            node.replace(
-                typeArgumentNodes.length == 0
-                    ? new UnknownTypeNode(sourceSection)
-                    : typeArgumentNodes[index]);
+            if (typeArgumentNodes.length == 0) {
+              node.replace(new UnknownTypeNode(sourceSection));
+            } else {
+              node.replace(typeArgumentNodes[index])
+                  .accept(
+                      n -> {
+                        if (n instanceof SelfTypeNode selfTypeNode) {
+                          selfTypeNode.setOriginalAnchor(this);
+                        }
+                        return true;
+                      });
+            }
           }
           return true;
         });

pkl-core/src/test/files/LanguageSnippetTests/input-helper/types/moduleUnion1.pkl

@@ -0,0 +1,5 @@
+open module moduleUnion1
+
+import "moduleUnion2.pkl"
+
+typealias ModuleUnion<T> = moduleUnion2.ModuleUnion<T, module>

pkl-core/src/test/files/LanguageSnippetTests/input-helper/types/moduleUnion2.pkl

@@ -0,0 +1,3 @@
+open module moduleUnion2
+
+typealias ModuleUnion<T, U> = T | U | module