track PType instead of Set<PType>, fix subtype check for type args with unions

Author: HT154

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

@@ -21,6 +21,8 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
+import java.util.function.BiConsumer;
+import java.util.function.Supplier;
 import org.jspecify.annotations.Nullable;
 import org.pkl.core.Composite;
 import org.pkl.core.PClass;
@@ -38,7 +40,7 @@ public final class VmReference extends VmValue {
   private final ImRrbt<VmTyped> path;
   // candidate types can only be: PType.Class, PType.Alias (only preservedAliasTypes),
   // PType.StringLiteral, or PType.UNKNOWN
-  private final Set<PType> candidateTypes;
+  private final PType referentType;
 
   private boolean forced = false;
 
@@ -69,10 +71,10 @@ public VmReference(VmTyped domain, VmClass clazz, Object data) {
         normalizeTypes(new PType.Class(clazz.export()), clazz.getModule().getVmClass().export()));
   }
 
-  public VmReference(VmTyped domain, Object data, ImRrbt<VmTyped> path, Set<PType> candidateTypes) {
+  public VmReference(VmTyped domain, Object data, ImRrbt<VmTyped> path, PType referentType) {
     this.domain = domain;
     this.data = data;
-    this.candidateTypes = candidateTypes;
+    this.referentType = referentType;
     this.path = path;
   }
 
@@ -88,19 +90,26 @@ public List<VmTyped> getPath() {
     return path;
   }
 
+  public PType getReferentType() {
+    return referentType;
+  }
+
   // simplifies a type by:
   // * erasing constraints
   // * transforming T? into T|Null
   // * dereferencing aliases (except for well-known stdlib alias types)
   // * flattening unions
   // * when moduleClass is supplied, replace PType.MODULE with appropriate PType.Class
   // * drop PType.NOTHING, PType.Function, and PType.TypeVariable
-  private static Set<PType> normalizeTypes(PType type, PClass moduleClass) {
+  private static PType normalizeTypes(PType type, PClass moduleClass) {
     var types = new HashSet<PType>();
     normalizeTypes(type, moduleClass, types);
-    if (types.contains(PType.UNKNOWN)) return Set.of(PType.UNKNOWN);
-    if (containsClass(types, anyType.getPClass())) return Set.of(anyType);
-    return types;
+    if (types.size() == 1) return types.iterator().next();
+    if (types.contains(PType.UNKNOWN)) return PType.UNKNOWN;
+    if (containsClass(types, anyType.getPClass())) return anyType;
+    var typesList = new ArrayList<>(types);
+    typesList.sort(Comparator.comparing(Object::toString));
+    return new PType.Union(typesList);
   }
 
   private static void normalizeTypes(PType type, PClass moduleClass, Set<PType> result) {
@@ -119,8 +128,7 @@ private static void normalizeTypes(PType type, PClass moduleClass, Set<PType> re
       } else {
         var typeArgs = new ArrayList<PType>(clazz.getTypeArguments().size());
         for (var arg : clazz.getTypeArguments()) {
-          var tt = new ArrayList<>(normalizeTypes(arg, moduleClass));
-          typeArgs.add(tt.size() == 1 ? tt.get(0) : new PType.Union(tt));
+          typeArgs.add(normalizeTypes(arg, moduleClass));
         }
         result.add(new PType.Class(clazz.getPClass(), typeArgs));
       }
@@ -144,33 +152,50 @@ private static void normalizeTypes(PType type, PClass moduleClass, Set<PType> re
     }
   }
 
+  private static Iterable<PType> iterateTypes(PType t) {
+    if (t instanceof PType.Union union) return union.getElementTypes();
+    return Collections.singleton(t);
+  }
+
   public @Nullable VmReference withPropertyAccess(Identifier property) {
-    Set<PType> candidates = new HashSet<>();
-    for (var t : candidateTypes) {
-      getCandidatePropertyType(t, property.toString(), candidates);
-    }
-    if (candidates.isEmpty()) {
-      return null; // no valid property found
-    } else if (candidates.contains(PType.UNKNOWN)) {
-      // optimization: unknown allows all references, erase all candidates to only unknown
-      candidates = Set.of(PType.UNKNOWN);
-    }
-    return new VmReference(
-        domain, data, path.append(newAccess(property.toString(), null)), candidates);
+    var propString = property.toString();
+    return withAccess(
+        (t, candidates) -> getCandidatePropertyType(t, propString, candidates),
+        () -> newAccess(property.toString(), null));
   }
 
   public @Nullable VmReference withSubscriptAccess(Object key) {
+    return withAccess(
+        (t, candidates) -> getCandidateSubscriptType(t, key, candidates),
+        () -> newAccess(null, key));
+  }
+
+  private @Nullable VmReference withAccess(
+      BiConsumer<PType, Set<PType>> checkCandidate, Supplier<VmTyped> makeAccess) {
     Set<PType> candidates = new HashSet<>();
-    for (var t : candidateTypes) {
-      getCandidateSubscriptType(t, key, candidates);
+    for (var t : iterateTypes(referentType)) {
+      checkCandidate.accept(t, candidates);
     }
     if (candidates.isEmpty()) {
-      return null; // no valid subscript found
+      return null; // no valid access found
+    }
+
+    PType newReferent;
+    if (candidates.size() == 1) {
+      newReferent = candidates.iterator().next();
     } else if (candidates.contains(PType.UNKNOWN)) {
       // optimization: unknown allows all references, erase all candidates to only unknown
-      candidates = Set.of(PType.UNKNOWN);
+      newReferent = PType.UNKNOWN;
+    } else if (containsClass(candidates, anyType.getPClass())) {
+      // optimization: All allows all references, erase all candidates to only All
+      newReferent = anyType;
+    } else {
+      var types = new ArrayList<>(candidates);
+      types.sort(Comparator.comparing(Object::toString));
+      newReferent = new PType.Union(types);
     }
-    return new VmReference(domain, data, path.append(newAccess(null, key)), candidates);
+
+    return new VmReference(domain, data, path.append(makeAccess.get()), newReferent);
   }
 
   @SuppressWarnings("DuplicatedCode")
@@ -234,7 +259,7 @@ private static void getCandidateSubscriptType(PType type, Object key, Set<PType>
         || clazz.getPClass().getInfo() == PClassInfo.Map) {
       var typeArgs = clazz.getTypeArguments();
       var keyTypes = normalizeTypes(typeArgs.get(0), clazz.getPClass().getModuleClass());
-      for (var kt : keyTypes) {
+      for (var kt : iterateTypes(keyTypes)) {
         if (kt == PType.UNKNOWN
             || (kt instanceof PType.Class klazz
                 && klazz.getPClass().getInfo() == PClassInfo.forValue(VmValue.export(key)))
@@ -252,38 +277,34 @@ private static void getCandidateSubscriptType(PType type, Object key, Set<PType>
    */
   public boolean referentTypeIsSubtypeOf(PType type, PClass moduleClass) {
     // fast path: if referent is unknown it can match any type check
-    if (candidateTypes.contains(PType.UNKNOWN)) {
+    if (referentType == PType.UNKNOWN) {
       return true;
     }
 
-    var checkTypes = normalizeTypes(type, moduleClass);
+    var checkType = normalizeTypes(type, moduleClass);
     // fast path: short circuit if any referent is accepted
-    if (checkTypes.contains(PType.UNKNOWN) || containsClass(checkTypes, anyType.getPClass())) {
+    if (checkType == PType.UNKNOWN || isClass(checkType, anyType.getPClass())) {
       return true;
     }
     // fast path: short circuit if nothing is accepted
-    if (checkTypes.size() == 1 && checkTypes.contains(PType.NOTHING)) {
+    if (checkType == PType.NOTHING) {
       return false;
     }
 
-    // all candidate types must be subtypes of at least one target type
-    candidate:
-    for (var c : candidateTypes) {
-      for (var t : checkTypes) {
-        if (isSubtype(c, t)) continue candidate;
-      }
-      return false;
-    }
-    return true;
+    return isSubtype(referentType, checkType);
   }
 
   private static boolean containsClass(Set<PType> types, PClass pClass) {
     for (var t : types) {
-      if (t instanceof PType.Class clazz && clazz.getPClass() == pClass) return true;
+      if (isClass(t, pClass)) return true;
     }
     return false;
   }
 
+  private static boolean isClass(PType t, PClass pClass) {
+    return t instanceof PType.Class clazz && clazz.getPClass() == pClass;
+  }
+
   private static boolean isSubtype(PType a, PType b) {
     // checks if A is a subtype of B
     // cases (A -> B)
@@ -305,6 +326,8 @@ private static boolean isSubtype(PType a, PType b) {
     //     * invariant: A_i must be identical to B_i
     //     * covariant: A_i must be a subtype of B_i
     //     * contravariant: B_i must be a subtype of A_i
+    // * Union -> Union: Each elem of A must be a subtype of at least one elem of B
+    // * Non-union -> Union: A must be a subtype of at least one elem of B
     if (a == b) return true;
 
     if (a instanceof PType.StringLiteral aStr) {
@@ -366,6 +389,21 @@ private static boolean isSubtype(PType a, PType b) {
         }
       }
       return true;
+    } else if (b instanceof PType.Union bUnion) {
+      if (a instanceof PType.Union aUnion) {
+        a:
+        for (var aElem : aUnion.getElementTypes()) {
+          for (var bElem : bUnion.getElementTypes()) {
+            if (isSubtype(aElem, bElem)) continue a;
+          }
+          return false;
+        }
+        return true;
+      } else {
+        for (var bElem : bUnion.getElementTypes()) {
+          if (isSubtype(a, bElem)) return true;
+        }
+      }
     }
     return false;
   }
@@ -395,22 +433,14 @@ public Reference export() {
       pathList.add(elem.export());
     }
 
-    return new Reference(domain.export(), VmValue.export(data), pathList, exportReferentType());
-  }
-
-  public PType exportReferentType() {
-    if (candidateTypes.size() == 1) return candidateTypes.iterator().next();
-    var types = new ArrayList<>(candidateTypes);
-    // sort multiple candidate types to ensure stable output
-    types.sort(Comparator.comparing(Object::toString));
-    return new PType.Union(types);
+    return new Reference(domain.export(), VmValue.export(data), pathList, getReferentType());
   }
 
   public PType exportType() {
     return new PType.Class(
         RefModule.getReferenceClass().export(),
         new PType.Class(domain.getVmClass().export()),
-        exportReferentType());
+        getReferentType());
   }
 
   @Override
@@ -433,15 +463,15 @@ public boolean equals(@Nullable Object o) {
     return domain.equals(that.domain)
         && data.equals(that.data)
         && path.equals(that.path)
-        && candidateTypes.equals(that.candidateTypes);
+        && referentType.equals(that.referentType);
   }
 
   @Override
   public int hashCode() {
     int result = domain.hashCode();
     result = 31 * result + data.hashCode();
     result = 31 * result + path.hashCode();
-    result = 31 * result + candidateTypes.hashCode();
+    result = 31 * result + referentType.hashCode();
     return result;
   }
 }

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

@@ -277,7 +277,7 @@ public void visitReference(VmReference value) {
       append("Reference(");
       visit(value.getDomain());
       append(", ");
-      append(value.exportReferentType());
+      append(value.getReferentType());
       append(", ");
       visit(value.getData());
       append(")");

pkl-core/src/test/files/LanguageSnippetTests/input/api/reference.pkl

@@ -4,16 +4,14 @@ import "pkl:ref"
 class D extends ref.Domain {
   function renderReference(reference: ref.Reference<D, Any>): String =
     let (data = reference.getData())
-      if (data is Resource)
-        let (
-          path =
-            reference
-              .getPath()
-              .map((elem) -> if (elem.isProperty) ".\(elem.property)" else "[\(elem.key)]")
-        )
-          "${\(data.name)\(path.join(""))}"
-      else
-        throw("can only render references rooted to Resource instances")
+    let (root = if (data is Resource) data.name else data.toString())
+    let (
+      path =
+        reference
+          .getPath()
+          .map((elem) -> if (elem.isProperty) ".\(elem.property)" else "[\(elem.key)]")
+    )
+      "${\(root)\(path.join(""))}"
 }
 
 local const d: D = new {}
@@ -86,7 +84,7 @@ aRef: Ref<A> = a.$
 bRef: Ref<B> = b.$
 unknownRef: Ref<A | B> = aRef
 unknownRef2: Ref<A> | Ref<B> = aRef
-unknownRef3: Ref<A | B> = aOrB.$
+unknownRef3: Ref<B | A> = aOrB.$
 
 k: K = new {
   aId = aRef.id
@@ -121,6 +119,12 @@ refInterpolation = "\(aRef.outputs.someListing[1])"
 kInterpolation = "\(k)"
 aValuesJoined = k.aValues.join("\n").replaceAll(Regex("@[a-z0-9]+"), "@<addr>")
 
+// ensure that type arguments that are unions are handled correctly
+typeArgs = ref.Reference(d, TypeHolder, null).prop as Ref<Listing<Number | Boolean | String>>
+class TypeHolder {
+  prop: Listing<String | Boolean | Number>
+}
+
 output {
   renderer {
     converters {

pkl-core/src/test/files/LanguageSnippetTests/output/api/reference.pcf

@@ -57,3 +57,4 @@ aValuesJoined = """
   org.pkl.core.runtime.VmReference@<addr>
   org.pkl.core.runtime.VmReference@<addr>
   """
+typeArgs = "${null.prop}"