Fix name resolution of types in constraints (#1826)

This fixes an issue where type names are resolved incorrectly
in typealias constraints.

Co-authored-by: Islon Scherer <islonscherer@gmail.com>

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

@@ -129,6 +129,7 @@ import org.pkl.core.ast.expression.primary.ExecuteCustomThisWithRootNode;
 import org.pkl.core.ast.expression.primary.GetEnclosingReceiverNode;
 import org.pkl.core.ast.expression.primary.GetMemberKeyNode;
 import org.pkl.core.ast.expression.primary.GetModuleNode;
+import org.pkl.core.ast.expression.primary.GetModuleOwnerNode;
 import org.pkl.core.ast.expression.primary.GetOwnerNode;
 import org.pkl.core.ast.expression.primary.GetReceiverNode;
 import org.pkl.core.ast.expression.primary.GetTypeAliasModuleNode;
@@ -2225,13 +2226,21 @@ public class AstBuilder extends AbstractAstBuilder<Object> {
 
   private ResolveDeclaredTypeNode doVisitTypeName(QualifiedIdentifier ctx) {
     var identifiers = ctx.getIdentifiers();
+    var getModuleNode =
+        isBaseModule
+            ? new ConstantValueNode(BaseModule.getModule())
+            : symbolTable.isInTypeAliasScope
+                ? new GetTypeAliasModuleNode(VmUtils.unavailableSourceSection())
+                : new GetModuleOwnerNode(VmUtils.unavailableSourceSection());
     return switch (identifiers.size()) {
       case 1 -> {
         var identifier = identifiers.get(0);
+        var sourceSection = createSourceSection(identifier);
         yield new ResolveSimpleDeclaredTypeNode(
-            createSourceSection(identifier),
+            sourceSection,
             org.pkl.core.runtime.Identifier.get(identifier.getValue()),
-            isBaseModule);
+            isBaseModule,
+            getModuleNode);
       }
       case 2 -> {
         var identifier1 = identifiers.get(0);
@@ -2241,7 +2250,8 @@ public class AstBuilder extends AbstractAstBuilder<Object> {
             createSourceSection(identifier1),
             createSourceSection(identifier2),
             org.pkl.core.runtime.Identifier.localProperty(identifier1.getValue()),
-            org.pkl.core.runtime.Identifier.get(identifier2.getValue()));
+            org.pkl.core.runtime.Identifier.get(identifier2.getValue()),
+            getModuleNode);
       }
       default ->
           throw exceptionBuilder()

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

@@ -41,6 +41,10 @@ public final class ExecuteCustomThisWithRootNode extends ExpressionNode {
       new CustomThisNode(VmUtils.unavailableSourceSection());
   private @Child DirectCallNode callNode;
 
+  // shouldn't be marked `@Child` because this node is actually the child of the SimpleRootNode
+  // created in the constructor.
+  private final ExpressionNode expressionNode;
+
   public ExecuteCustomThisWithRootNode(
       SourceSection sourceSection,
       ExpressionNode expressionNode,
@@ -49,6 +53,7 @@ public final class ExecuteCustomThisWithRootNode extends ExpressionNode {
       int[] forGeneratorSlots,
       int[] parameterSlots) {
     super(sourceSection);
+    this.expressionNode = expressionNode;
     frameDescriptor.findOrAddAuxiliarySlot(CustomThisScope.FRAME_SLOT_ID);
     var rootNode =
         new SimpleRootNode(
@@ -60,6 +65,10 @@ public final class ExecuteCustomThisWithRootNode extends ExpressionNode {
     this.callNode = DirectCallNode.create(rootNode.getCallTarget());
   }
 
+  public ExpressionNode getExpressionNode() {
+    return expressionNode;
+  }
+
   @Override
   public Object executeGeneric(VirtualFrame frame) {
     var customThis = customThisNode.executeGeneric(frame);

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

@@ -0,0 +1,46 @@
+/*
+ * Copyright © 2026 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.expression.primary;
+
+import com.oracle.truffle.api.CompilerDirectives;
+import com.oracle.truffle.api.frame.VirtualFrame;
+import com.oracle.truffle.api.source.SourceSection;
+import org.pkl.core.ast.ExpressionNode;
+import org.pkl.core.runtime.VmUtils;
+
+public final class GetModuleOwnerNode extends ExpressionNode {
+
+  public GetModuleOwnerNode(SourceSection sourceSection) {
+    super(sourceSection);
+  }
+
+  @Override
+  public Object executeGeneric(VirtualFrame frame) {
+    CompilerDirectives.transferToInterpreter();
+
+    var levelsUp = -1;
+    for (var current = VmUtils.getOwner(frame);
+        current != null;
+        current = current.getEnclosingOwner()) {
+      if (!current.isParseTimeInvisibleScope()) {
+        levelsUp += 1;
+      }
+    }
+
+    return replace(levelsUp == 0 ? new GetOwnerNode() : new GetEnclosingOwnerNode(levelsUp))
+        .executeGeneric(frame);
+  }
+}

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

@@ -18,12 +18,14 @@ package org.pkl.core.ast.member;
 import com.oracle.truffle.api.CompilerDirectives;
 import com.oracle.truffle.api.CompilerDirectives.CompilationFinal;
 import com.oracle.truffle.api.frame.VirtualFrame;
+import com.oracle.truffle.api.nodes.Node;
 import com.oracle.truffle.api.source.SourceSection;
 import java.util.ArrayList;
 import java.util.List;
 import org.jspecify.annotations.Nullable;
 import org.pkl.core.TypeParameter;
 import org.pkl.core.ast.ExpressionNode;
+import org.pkl.core.ast.expression.primary.ExecuteCustomThisWithRootNode;
 import org.pkl.core.ast.expression.primary.GetTypeAliasModuleNode;
 import org.pkl.core.ast.type.UnresolvedTypeNode;
 import org.pkl.core.runtime.VmTypeAlias;
@@ -87,16 +89,26 @@ public final class TypeAliasNode extends ExpressionNode {
             frame.materialize());
 
     VmUtils.evaluateAnnotations(frame, annotationNodes, annotations);
+    initTypeAliasModule(typeAnnotationNode, module);
     var bodyTypeNode = typeAnnotationNode.execute(frame);
-    bodyTypeNode.accept(
-        node -> {
-          if (node instanceof GetTypeAliasModuleNode getTypeAliasModuleNode) {
-            getTypeAliasModuleNode.lateInitModule(module);
-          }
-          return true;
-        });
     cachedTypeAlias.initTypeCheckNode(bodyTypeNode);
 
     return cachedTypeAlias;
   }
+
+  private static void initTypeAliasModule(Node node, VmTyped module) {
+    node.accept(
+        n -> {
+          if (n instanceof GetTypeAliasModuleNode getTypeAliasModuleNode) {
+            getTypeAliasModuleNode.lateInitModule(module);
+            // `node.accept` won't walk past root nodes, but this can happen if we have a
+            // let expression in a constraint.
+            //
+            // If we encounter `ExecuteCustomThisWithRootNode`; walk its expression node directly.
+          } else if (n instanceof ExecuteCustomThisWithRootNode executeCustomThisWithRootNode) {
+            initTypeAliasModule(executeCustomThisWithRootNode.getExpressionNode(), module);
+          }
+          return true;
+        });
+  }
 }

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

@@ -1,5 +1,5 @@
 /*
- * Copyright © 2024 Apple Inc. and the Pkl project authors. All rights reserved.
+ * Copyright © 2024-2026 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.
@@ -18,6 +18,7 @@ package org.pkl.core.ast.type;
 import com.oracle.truffle.api.CompilerDirectives;
 import com.oracle.truffle.api.frame.VirtualFrame;
 import com.oracle.truffle.api.source.SourceSection;
+import org.pkl.core.ast.ExpressionNode;
 import org.pkl.core.runtime.*;
 
 public final class ResolveQualifiedDeclaredTypeNode extends ResolveDeclaredTypeNode {
@@ -25,19 +26,22 @@ public final class ResolveQualifiedDeclaredTypeNode extends ResolveDeclaredTypeN
   private final SourceSection typeNameSection;
   private final Identifier moduleName;
   private final Identifier typeName;
+  @Child private ExpressionNode getModuleNode;
 
   public ResolveQualifiedDeclaredTypeNode(
       SourceSection sourceSection,
       SourceSection moduleNameSection,
       SourceSection typeNameSection,
       Identifier moduleName,
-      Identifier typeName) {
+      Identifier typeName,
+      ExpressionNode getModuleNode) {
 
     super(sourceSection);
     this.moduleNameSection = moduleNameSection;
     this.typeNameSection = typeNameSection;
     this.moduleName = moduleName;
     this.typeName = typeName;
+    this.getModuleNode = getModuleNode;
 
     assert moduleName.isLocalProp();
     assert typeName.isRegular();
@@ -47,7 +51,7 @@ public final class ResolveQualifiedDeclaredTypeNode extends ResolveDeclaredTypeN
   public Object executeGeneric(VirtualFrame frame) {
     CompilerDirectives.transferToInterpreter();
 
-    var enclosingModule = getEnclosingModule(VmUtils.getOwner(frame));
+    var enclosingModule = (VmTyped) getModuleNode.executeGeneric(frame);
     var importedModule = getImport(enclosingModule, moduleName, moduleNameSection);
 
     // search module hierarchy

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

@@ -1,5 +1,5 @@
 /*
- * Copyright © 2024 Apple Inc. and the Pkl project authors. All rights reserved.
+ * Copyright © 2024-2026 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.
@@ -18,18 +18,24 @@ package org.pkl.core.ast.type;
 import com.oracle.truffle.api.CompilerDirectives;
 import com.oracle.truffle.api.frame.VirtualFrame;
 import com.oracle.truffle.api.source.SourceSection;
+import org.pkl.core.ast.ExpressionNode;
 import org.pkl.core.runtime.*;
 
 public final class ResolveSimpleDeclaredTypeNode extends ResolveDeclaredTypeNode {
   private final Identifier typeName;
   private final boolean isBaseModule;
+  @Child private ExpressionNode getModuleNode;
 
   public ResolveSimpleDeclaredTypeNode(
-      SourceSection sourceSection, Identifier typeName, boolean isBaseModule) {
+      SourceSection sourceSection,
+      Identifier typeName,
+      boolean isBaseModule,
+      ExpressionNode getModuleNode) {
 
     super(sourceSection);
     this.typeName = typeName;
     this.isBaseModule = isBaseModule;
+    this.getModuleNode = getModuleNode;
   }
 
   @Override
@@ -37,7 +43,7 @@ public final class ResolveSimpleDeclaredTypeNode extends ResolveDeclaredTypeNode
     CompilerDirectives.transferToInterpreter();
 
     var localTypeName = typeName.toLocalProperty();
-    var enclosingModule = getEnclosingModule(VmUtils.getOwner(frame));
+    var enclosingModule = (VmTyped) getModuleNode.executeGeneric(frame);
 
     // search enclosing module for local class/type alias or module import
     var result = getType(enclosingModule, localTypeName, sourceSection);

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

@@ -1,3 +1,8 @@
 class MyClass
 
 typealias MyTypeAlias = MyClass(getClass() == MyClass)
+
+typealias MyTypeAlias2 = MyClass(this is MyClass)
+
+// a let expression turns the constraint node into its own root node
+typealias MyTypeAlias3 = MyClass(let (_ = true) this is MyClass)

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

@@ -0,0 +1,3 @@
+import "typeAliasConstraint2.pkl"
+
+typealias MyTypeAlias = Any(this is typeAliasConstraint2.MyClass)

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

@@ -0,0 +1,8 @@
+import "../../input-helper/types/typeAliasConstraint2.pkl"
+import "../../input-helper/types/typeAliasConstraint3.pkl"
+
+res: typeAliasConstraint2.MyTypeAlias2 = new typeAliasConstraint2.MyClass {}
+
+res2: typeAliasConstraint3.MyTypeAlias = new typeAliasConstraint2.MyClass {}
+
+res3: typeAliasConstraint2.MyTypeAlias3 = new typeAliasConstraint2.MyClass {}

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

@@ -0,0 +1,6 @@
+import "../../input-helper/types/typeAliasAmendModule.pkl"
+
+local const function isValid(s: String): Boolean = s.startsWith("Alice")
+
+res = typeAliasAmendModule.name
+

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

@@ -0,0 +1,3 @@
+res {}
+res2 {}
+res3 {}

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

@@ -0,0 +1 @@
+res = "Bob Marley"