Correctly type check Class<T>

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

@@ -34,6 +34,7 @@
 import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.Set;
 import java.util.stream.Collectors;
 import org.jspecify.annotations.NonNull;
@@ -62,10 +63,18 @@
 
 public abstract class TypeNode extends PklNode {
 
-  public interface ClassTypeNode {
+  /**
+   * Type node that corresponds to a simple, unparameterized {@link VmClass}.
+   *
+   * <p>This includes generic classes written without any type arguments like {@code List}.
+   */
+  public interface SimpleClassTypeNode {
     VmClass getVmClass();
   }
 
+  /** Type node that corresponds to a user-defined class (or module class). */
+  public interface UserClassTypeNode extends SimpleClassTypeNode {}
+
   protected TypeNode(SourceSection sourceSection) {
     super(sourceSection);
   }
@@ -412,7 +421,7 @@ protected boolean acceptTypeNode(boolean visitTypeArguments, TypeNodeConsumer co
 
   /** The `module` type for a final module. */
   public static final class FinalModuleTypeNode extends ObjectSlotTypeNode
-      implements ClassTypeNode {
+      implements UserClassTypeNode {
     private final VmClass moduleClass;
 
     public FinalModuleTypeNode(SourceSection sourceSection, VmClass moduleClass) {
@@ -467,7 +476,7 @@ protected boolean acceptTypeNode(boolean visitTypeArguments, TypeNodeConsumer co
 
   /** The `module` type for an open module. */
   public static final class NonFinalModuleTypeNode extends ObjectSlotTypeNode
-      implements ClassTypeNode {
+      implements UserClassTypeNode {
     private final VmClass moduleClass; // only used by getVmClass()
     @Child private ExpressionNode getModuleNode;
 
@@ -581,7 +590,8 @@ protected boolean acceptTypeNode(boolean visitTypeArguments, TypeNodeConsumer co
     }
   }
 
-  public static final class TypedTypeNode extends ObjectSlotTypeNode {
+  public static final class TypedTypeNode extends ObjectSlotTypeNode
+      implements SimpleClassTypeNode {
     public TypedTypeNode(SourceSection sourceSection) {
       super(sourceSection);
     }
@@ -609,7 +619,8 @@ protected boolean acceptTypeNode(boolean visitTypeArguments, TypeNodeConsumer co
     }
   }
 
-  public static final class DynamicTypeNode extends ObjectSlotTypeNode {
+  public static final class DynamicTypeNode extends ObjectSlotTypeNode
+      implements SimpleClassTypeNode {
     public DynamicTypeNode(SourceSection sourceSection) {
       super(sourceSection);
     }
@@ -652,7 +663,8 @@ protected boolean acceptTypeNode(boolean visitTypeArguments, TypeNodeConsumer co
    * String/Boolean/Int/Float and their supertypes, only `VmValue`s can possibly pass its type
    * check.
    */
-  public static final class FinalClassTypeNode extends ObjectSlotTypeNode implements ClassTypeNode {
+  public static final class FinalClassTypeNode extends ObjectSlotTypeNode
+      implements UserClassTypeNode {
     private final VmClass clazz;
 
     public FinalClassTypeNode(SourceSection sourceSection, VmClass clazz) {
@@ -709,7 +721,7 @@ protected boolean acceptTypeNode(boolean visitTypeArguments, TypeNodeConsumer co
    * check.
    */
   public abstract static class NonFinalClassTypeNode extends ObjectSlotTypeNode
-      implements ClassTypeNode {
+      implements UserClassTypeNode {
     protected final VmClass clazz;
 
     public NonFinalClassTypeNode(SourceSection sourceSection, VmClass clazz) {
@@ -2994,7 +3006,8 @@ public VmTyped getMirror() {
     }
   }
 
-  public static final class AnyTypeNode extends WriteFrameSlotTypeNode {
+  public static final class AnyTypeNode extends WriteFrameSlotTypeNode
+      implements SimpleClassTypeNode {
     public AnyTypeNode(SourceSection sourceSection) {
       super(sourceSection);
     }
@@ -3026,7 +3039,8 @@ protected boolean acceptTypeNode(boolean visitTypeArguments, TypeNodeConsumer co
     }
   }
 
-  public static final class StringTypeNode extends ObjectSlotTypeNode {
+  public static final class StringTypeNode extends ObjectSlotTypeNode
+      implements SimpleClassTypeNode {
     public StringTypeNode(SourceSection sourceSection) {
       super(sourceSection);
     }
@@ -3054,7 +3068,8 @@ protected boolean acceptTypeNode(boolean visitTypeArguments, TypeNodeConsumer co
     }
   }
 
-  public static final class NumberTypeNode extends FrameSlotTypeNode {
+  public static final class NumberTypeNode extends FrameSlotTypeNode
+      implements SimpleClassTypeNode {
     public NumberTypeNode(SourceSection sourceSection) {
       super(sourceSection);
     }
@@ -3113,7 +3128,7 @@ protected boolean acceptTypeNode(boolean visitTypeArguments, TypeNodeConsumer co
     }
   }
 
-  public static final class IntTypeNode extends IntSlotTypeNode {
+  public static final class IntTypeNode extends IntSlotTypeNode implements SimpleClassTypeNode {
     public IntTypeNode(SourceSection sourceSection) {
       super(sourceSection);
     }
@@ -3141,7 +3156,7 @@ protected boolean acceptTypeNode(boolean visitTypeArguments, TypeNodeConsumer co
     }
   }
 
-  public static final class FloatTypeNode extends FrameSlotTypeNode {
+  public static final class FloatTypeNode extends FrameSlotTypeNode implements SimpleClassTypeNode {
     public FloatTypeNode(SourceSection sourceSection) {
       super(sourceSection);
     }
@@ -3181,7 +3196,8 @@ protected boolean acceptTypeNode(boolean visitTypeArguments, TypeNodeConsumer co
     }
   }
 
-  public static final class BooleanTypeNode extends FrameSlotTypeNode {
+  public static final class BooleanTypeNode extends FrameSlotTypeNode
+      implements SimpleClassTypeNode {
     public BooleanTypeNode(SourceSection sourceSection) {
       super(sourceSection);
     }
@@ -3221,6 +3237,89 @@ protected boolean acceptTypeNode(boolean visitTypeArguments, TypeNodeConsumer co
     }
   }
 
+  public abstract static class ClassClassTypeNode extends ValidatingObjectSlotTypeNode {
+
+    @Child private TypeNode typeNode;
+    private @Nullable VmClass clazz;
+
+    public ClassClassTypeNode(SourceSection sourceSection, TypeNode typeNode) {
+      super(sourceSection);
+      this.typeNode = typeNode;
+      validate();
+    }
+
+    @Override
+    public String getValidationErrorKey() {
+      return "notAClassType";
+    }
+
+    @Override
+    public @Nullable Node getViolatingNode() {
+      if (typeNode instanceof SimpleClassTypeNode simpleClassTypeNode) {
+        clazz = simpleClassTypeNode.getVmClass();
+        return null;
+      }
+      if (typeNode instanceof UnknownTypeNode || typeNode instanceof TypeVariableNode) {
+        clazz = null;
+        return null;
+      }
+      return typeNode;
+    }
+
+    @Override
+    protected boolean isIncludedInTrace(Node node) {
+      return node instanceof ClassClassTypeNode;
+    }
+
+    public @Nullable VmClass getClazz() {
+      return clazz;
+    }
+
+    @Override
+    public VmClass getVmClass() {
+      return BaseModule.getClassClass();
+    }
+
+    @Specialization
+    protected Object eval(VmClass value) {
+      // clazz will be null iff the type arg is a type variable.
+      // in this case, skip the subclass check and behave like a bare `Class` type annotation.
+      if (clazz == null) {
+        return value;
+      }
+
+      if (!value.isSubclassOf(clazz)) {
+        throw new VmTypeMismatchException.Class(sourceSection, value, clazz);
+      }
+
+      return value;
+    }
+
+    @Fallback
+    protected Object fallback(Object value) {
+      throw typeMismatch(value, BaseModule.getClassClass());
+    }
+
+    @Override
+    protected boolean acceptTypeNode(boolean visitTypeArguments, TypeNodeConsumer consumer) {
+      return consumer.accept(this);
+    }
+
+    @Override
+    protected boolean doIsEquivalentTo(TypeNode other) {
+      if (!(other instanceof ClassClassTypeNode classClassTypeNode)) {
+        return false;
+      }
+
+      return Objects.equals(clazz, classClassTypeNode.getClazz());
+    }
+
+    @Override
+    public VmList getTypeArgumentMirrors() {
+      return VmList.of(typeNode.getMirror());
+    }
+  }
+
   public abstract static class ValidatingObjectSlotTypeNode extends ObjectSlotTypeNode {
 
     protected ValidatingObjectSlotTypeNode(SourceSection sourceSection) {

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

@@ -278,9 +278,8 @@ public TypeNode execute(VirtualFrame frame) {
           return FunctionNClassTypeNodeGen.create(sourceSection, resolvedTypeArgumentNodes);
         }
 
-        // erase `x: Class<Foo>` to `x: Class` for now (cf. function types)
         if (clazz.isClassClass()) {
-          return new FinalClassTypeNode(sourceSection, clazz);
+          return ClassClassTypeNodeGen.create(sourceSection, typeArgumentNodes[0].execute(frame));
         }
 
         if (clazz.isVarArgsClass()) {

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

@@ -169,6 +169,54 @@ protected Boolean hasHint() {
     }
   }
 
+  public static final class Class extends VmTypeMismatchException {
+
+    private final VmClass expectedClass;
+
+    public Class(SourceSection sourceSection, VmClass actualClass, VmClass expectedClass) {
+      super(sourceSection, actualClass);
+      this.expectedClass = expectedClass;
+    }
+
+    @Override
+    @TruffleBoundary
+    public void buildMessage(
+        AnsiStringBuilder builder, String indent, boolean withPowerAssertions) {
+      var actualClass = (VmClass) actualValue;
+      var renderedExpectedClass = "Class<" + expectedClass + ">";
+      var renderedActualClass = "Class<" + actualClass + ">";
+
+      // give better error than "expected Class<foo.Bar>, but got Class<foo.Bar>" in case of naming
+      // conflict
+      if (actualClass.getQualifiedName().equals(expectedClass.getQualifiedName())) {
+        var actualModuleUri = actualClass.getModule().getModuleInfo().getModuleKey().getUri();
+        var expectedModuleUri = expectedClass.getModule().getModuleInfo().getModuleKey().getUri();
+
+        builder
+            .append(
+                ErrorMessages.createIndented(
+                    actualClass.getPClassInfo().isModuleClass()
+                        ? "typeMismatchVersionConflict1"
+                        : "typeMismatchVersionConflict2",
+                    indent,
+                    renderedExpectedClass,
+                    expectedModuleUri,
+                    actualModuleUri))
+            .append("\n");
+        return;
+      }
+
+      builder.append(
+          ErrorMessages.createIndented(
+              "typeMismatch", indent, renderedExpectedClass, renderedActualClass));
+    }
+
+    @Override
+    protected Boolean hasHint() {
+      return false;
+    }
+  }
+
   public static final class Constraint extends VmTypeMismatchException {
 
     private final SourceSection constraintBodySourceSection;

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

@@ -173,7 +173,7 @@ private VmClass getOptionsClass(VmTyped command) {
     if (optionsTypeNode instanceof TypeNode.TypedTypeNode) {
       return BaseModule.getTypedClass();
     }
-    if (!(optionsTypeNode instanceof TypeNode.ClassTypeNode node)) {
+    if (!(optionsTypeNode instanceof TypeNode.UserClassTypeNode node)) {
       throw exceptionBuilder()
           .withSourceSection(optionsTypeNode.getSourceSection())
           .evalError(

pkl-core/src/main/resources/org/pkl/core/errorMessages.properties

@@ -200,6 +200,9 @@ Duplicate type parameter `{0}`.
 notAParameterizableClass=\
 Class `{0}` does not have type parameters.
 
+notAClassType=\
+`Class` type arguments may only be class types, module types, `unknown`, or `module`.
+
 notASubclassOfTyped=\
 Class `{0}` is not a subtype of `Typed`.
 

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

@@ -0,0 +1,2 @@
+module classType
+class Foo

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

@@ -0,0 +1,2 @@
+module classType
+class Foo

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

@@ -0,0 +1,27 @@
+open module classType
+
+import ".../input-helper/types/classTypeA.pkl"
+import ".../input-helper/types/classTypeB.pkl"
+
+open class A
+open class B
+class C extends A
+class D extends module
+typealias E = A
+
+res0 = C is Class<C>
+res1 = C is Class<A>
+
+res3 = C is Class
+res4 = C is Class<unknown>
+res5 = C is Class<Any>
+
+res15 = D is Class<module>
+res16 = D is Class<Object>
+res17 = D is Class<Typed>
+res18 = D is Class<Dynamic>
+res19 = D is Class<Int>
+
+typealias F<T> = List<Class<T>>
+res20 = List(A, C) is F<A>
+res21 = List(new A {}, new B {}, new C {}).filterIsInstance(A).length

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

@@ -0,0 +1,3 @@
+extends "classType.pkl"
+
+res2 = C as Class<B>

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

@@ -0,0 +1,3 @@
+extends "classType.pkl"
+
+res14 = C as Class<module>

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

@@ -0,0 +1,3 @@
+extends "classType.pkl"
+
+res21 = List(A, C) is F<A | B>

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

@@ -0,0 +1,5 @@
+extends "classType.pkl"
+
+typealias G<T> = F<T>
+
+res21 = List(A, C) is G<A | B>

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

@@ -0,0 +1,5 @@
+extends "classType.pkl"
+
+typealias G<T> = F<T | B>
+
+res21 = List(A, C) is G<A>

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

@@ -0,0 +1,6 @@
+extends "classType.pkl"
+
+import ".../input-helper/types/classTypeA.pkl"
+import ".../input-helper/types/classTypeB.pkl"
+
+res6 = classTypeA.getClass() as Class<classTypeB>

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

@@ -0,0 +1,6 @@
+extends "classType.pkl"
+
+import ".../input-helper/types/classTypeA.pkl"
+import ".../input-helper/types/classTypeB.pkl"
+
+res7 = classTypeA.Foo as Class<classTypeB.Foo>

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

@@ -0,0 +1,3 @@
+extends "classType.pkl"
+
+res8 = C as Class<A | B>

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

@@ -0,0 +1,3 @@
+extends "classType.pkl"
+
+res9 = C as Class<A?>