Restore late-binding for constraints within typealiases

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

@@ -19,6 +19,7 @@
 import com.oracle.truffle.api.frame.VirtualFrame;
 import com.oracle.truffle.api.source.SourceSection;
 import org.pkl.core.runtime.*;
+import org.pkl.core.util.Pair;
 
 public final class ResolveQualifiedDeclaredTypeNode extends ResolveDeclaredTypeNode {
   private final SourceSection moduleNameSection;
@@ -54,7 +55,7 @@ public Object executeGeneric(VirtualFrame frame) {
     // (type declared in base module is accessible through extending and amending modules)
     for (var currModule = importedModule; currModule != null; currModule = currModule.getParent()) {
       var result = getType(currModule, typeName, sourceSection);
-      if (result != null) return result;
+      if (result != null) return Pair.of(importedModule, result);
     }
 
     throw exceptionBuilder()

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

@@ -1563,11 +1563,16 @@ public VmTyped getMirror() {
   public static final class TypeAliasTypeNode extends TypeNode {
     private final VmTypeAlias typeAlias;
     private final TypeNode[] typeArgumentNodes;
+    private final VmTyped receiver;
     @Child private TypeNode aliasedTypeNode;
 
     public TypeAliasTypeNode(
-        SourceSection sourceSection, VmTypeAlias typeAlias, TypeNode[] typeArgumentNodes) {
+        SourceSection sourceSection,
+        VmTypeAlias typeAlias,
+        TypeNode[] typeArgumentNodes,
+        VmTyped receiver) {
       super(sourceSection);
+      this.receiver = receiver;
 
       if (!typeAlias.isInitialized()) {
         CompilerDirectives.transferToInterpreter();
@@ -1621,9 +1626,11 @@ public VmList getTypeArgumentMirrors() {
     public void 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()));
-
+      var owner = (VmTyped) VmUtils.getOwner(typeAlias.getEnclosingFrame());
+      VmUtils.setOwner(frame, owner);
+      if (!(prevReceiver instanceof VmTyped vmTyped && vmTyped.isAmending(owner))) {
+        VmUtils.setReceiver(frame, receiver);
+      }
       try {
         aliasedTypeNode.execute(frame, value);
       } finally {

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

@@ -27,6 +27,7 @@
 import org.pkl.core.ast.type.TypeNode.*;
 import org.pkl.core.ast.type.TypeNodeFactory.*;
 import org.pkl.core.runtime.*;
+import org.pkl.core.util.Pair;
 
 public abstract class UnresolvedTypeNode extends PklNode {
   protected UnresolvedTypeNode(SourceSection sourceSection) {
@@ -132,6 +133,11 @@ public TypeNode execute(VirtualFrame frame) {
       CompilerDirectives.transferToInterpreter();
 
       var type = resolveTypeNode.executeGeneric(frame);
+      Object receiver = null;
+      if (type instanceof Pair<?, ?> pair) {
+        receiver = pair.getFirst();
+        type = pair.getSecond();
+      }
 
       if (type instanceof VmClass clazz) {
         // Note: FinalClassTypeNode and NonFinalClassTypeNode assume that
@@ -183,7 +189,10 @@ public TypeNode execute(VirtualFrame frame) {
           }
         }
 
-        return new TypeAliasTypeNode(sourceSection, alias, new TypeNode[0]);
+        if (receiver == null) {
+          receiver = VmUtils.getReceiver(alias.getEnclosingFrame());
+        }
+        return new TypeAliasTypeNode(sourceSection, alias, new TypeNode[0], (VmTyped) receiver);
       }
 
       var module = (VmTyped) type;
@@ -214,6 +223,11 @@ public TypeNode execute(VirtualFrame frame) {
       CompilerDirectives.transferToInterpreter();
 
       var baseType = resolveTypeNode.executeGeneric(frame);
+      Object receiver = null;
+      if (baseType instanceof Pair<?, ?> pair) {
+        receiver = pair.getFirst();
+        baseType = pair.getSecond();
+      }
 
       if (baseType instanceof VmClass clazz) {
         checkNumberOfTypeArguments(clazz);
@@ -290,7 +304,11 @@ public TypeNode execute(VirtualFrame frame) {
         for (var i = 0; i < argLength; i++) {
           resolvedTypeArgumentNodes[i] = typeArgumentNodes[i].execute(frame);
         }
-        return new TypeAliasTypeNode(sourceSection, typeAlias, resolvedTypeArgumentNodes);
+        if (receiver == null) {
+          receiver = VmUtils.getReceiver(typeAlias.getEnclosingFrame());
+        }
+        return new TypeAliasTypeNode(
+            sourceSection, typeAlias, resolvedTypeArgumentNodes, (VmTyped) receiver);
       }
 
       var module = (VmTyped) baseType;

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

@@ -64,6 +64,16 @@ public VmClass getVmClass() {
     return (VmTyped) parent;
   }
 
+  public boolean isAmending(VmTyped other) {
+    if (this == other) {
+      return true;
+    }
+    if (parent == null) {
+      return false;
+    }
+    return ((VmTyped) parent).isAmending(other);
+  }
+
   @Override
   public boolean isPrototype() {
     return this == getPrototype();

pkl-core/src/test/files/LanguageSnippetTests/input/types/helpers/someModule2.pkl

@@ -0,0 +1,3 @@
+typealias IsValid = Any(isValid)
+
+isValid: Boolean = false

pkl-core/src/test/files/LanguageSnippetTests/input/types/helpers/someModule3.pkl

@@ -0,0 +1,3 @@
+amends "someModule2.pkl"
+
+isValid = true

pkl-core/src/test/files/LanguageSnippetTests/input/types/typeAliasConstraint3.pkl

@@ -0,0 +1,25 @@
+typealias IsValid = Any(isValid)
+
+hidden isValid = false
+
+prop: Boolean = 5 is IsValid
+
+hidden obj1 = (module) {
+  isValid = false
+}
+
+hidden obj2 = (module) {
+  isValid = true
+}
+
+hidden obj3 = (module) {
+  local isValid = throw("nuh uh")
+}
+
+output {
+  value = new Dynamic {
+    prop1 = obj1
+    prop2 = obj2
+    prop3 = obj3
+  }
+}

pkl-core/src/test/files/LanguageSnippetTests/input/types/typeAliasConstraint4.pkl

@@ -0,0 +1,5 @@
+import "helpers/someModule2.pkl"
+import "helpers/someModule3.pkl"
+
+res1 = 5 is someModule2.IsValid
+res2 = 5 is someModule3.IsValid

pkl-core/src/test/files/LanguageSnippetTests/output/types/helpers/someModule2.pcf

@@ -0,0 +1 @@
+isValid = false

pkl-core/src/test/files/LanguageSnippetTests/output/types/helpers/someModule3.pcf

@@ -0,0 +1 @@
+isValid = true

pkl-core/src/test/files/LanguageSnippetTests/output/types/typeAliasConstraint3.pcf

@@ -0,0 +1,9 @@
+prop1 {
+  prop = false
+}
+prop2 {
+  prop = true
+}
+prop3 {
+  prop = false
+}

pkl-core/src/test/files/LanguageSnippetTests/output/types/typeAliasConstraint4.pcf

@@ -0,0 +1,2 @@
+res1 = false
+res2 = true