Improve cases where type nodes can turn into TypeNode.UnionOfStringLiterals (#1875)

This adds logic to UnresolvedTypeNode to create a `UnionOfStringLiterals`
type node if a member is a typealias who is a string or a string literal union.

Also, this fixes an issue where marking a duplicate string literal with the
default marker throws an index out of bounds exception.

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

@@ -25,14 +25,12 @@ import java.io.IOException;
 import java.net.URI;
 import java.net.URISyntaxException;
 import java.util.ArrayList;
-import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Locale;
 import java.util.Objects;
 import java.util.Set;
 import java.util.function.BiFunction;
 import java.util.function.Supplier;
-import java.util.stream.Collectors;
 import org.graalvm.collections.EconomicMap;
 import org.jspecify.annotations.Nullable;
 import org.pkl.core.PClassInfo;
@@ -578,9 +576,7 @@ public class AstBuilder extends AbstractAstBuilder<Object> {
       return new UnresolvedTypeNode.UnionOfStringLiterals(
           createSourceSection(type),
           defaultIndex,
-          elementTypes.stream()
-              .map(it -> ((StringConstantType) it).getStr().getString())
-              .collect(Collectors.toCollection(LinkedHashSet::new)));
+          elementTypes.stream().map(it -> ((StringConstantType) it).getStr().getString()).toList());
     }
 
     var elements = new UnresolvedTypeNode[elementTypes.size()];

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

@@ -34,6 +34,7 @@ import com.oracle.truffle.api.profiles.LoopConditionProfile;
 import com.oracle.truffle.api.source.SourceSection;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -1122,37 +1123,45 @@ public abstract class TypeNode extends PklNode {
   }
 
   public static final class UnionOfStringLiteralsTypeNode extends ObjectSlotTypeNode {
-    private final Set<String> stringLiterals;
-    private final @Nullable String unionDefault;
     private final int defaultIndex;
+    private final List<String> declaredStringLiterals;
+    // intentionally not marked `@Child` because this only used for metadata, and never executed.
+    private final @Nullable TypeNode originalUnionTypeNode;
+    private final Set<String> stringLiterals;
 
     UnionOfStringLiteralsTypeNode(
-        SourceSection sourceSection, int defaultIndex, Set<String> stringLiterals) {
+        SourceSection sourceSection,
+        int defaultIndex,
+        List<String> declaredStringLiterals,
+        @Nullable TypeNode originalUnionTypeNode) {
       super(sourceSection);
-      assert !stringLiterals.isEmpty();
+      assert !declaredStringLiterals.isEmpty();
       this.defaultIndex = defaultIndex;
-      this.stringLiterals = stringLiterals;
-      if (defaultIndex == -1) {
-        unionDefault = null;
-      } else {
-        unionDefault = stringLiterals.toArray(new String[0])[defaultIndex];
-      }
+      this.declaredStringLiterals = declaredStringLiterals;
+      this.originalUnionTypeNode = originalUnionTypeNode;
+      this.stringLiterals = new HashSet<>(declaredStringLiterals);
     }
 
     @Override
     @TruffleBoundary
     protected VmType doGetType() {
-      return new VmType.UnionType(defaultIndex, stringLiterals.toArray(new String[0]));
+      if (originalUnionTypeNode != null) {
+        return originalUnionTypeNode.getType();
+      }
+      return new VmType.UnionType(defaultIndex, declaredStringLiterals.toArray(new String[0]));
     }
 
     @Override
     public VmTyped getMirror() {
+      if (originalUnionTypeNode != null) {
+        return originalUnionTypeNode.getMirror();
+      }
       return MirrorFactories.unionOfStringLiteralsTypeFactory.create(this);
     }
 
     public VmList getElementTypeMirrors() {
       var builder = VmList.EMPTY.builder();
-      for (var literal : stringLiterals) {
+      for (var literal : declaredStringLiterals) {
         builder.add(MirrorFactories.stringLiteralTypeFactory2.create(literal));
       }
       return builder.build();
@@ -1161,7 +1170,7 @@ public abstract class TypeNode extends PklNode {
     @Override
     protected Object executeLazily(VirtualFrame frame, Object value) {
       if (contains(value)) return value;
-      throw typeMismatch(value, stringLiterals);
+      throw typeMismatch(value, declaredStringLiterals);
     }
 
     @TruffleBoundary
@@ -1176,20 +1185,23 @@ public abstract class TypeNode extends PklNode {
     }
 
     @Override
-    public @Nullable Object createDefaultValue(
+    public @Nullable String createDefaultValue(
         VirtualFrame frame,
         VmLanguage language,
         SourceSection headerSection,
         String qualifiedName) {
-      return unionDefault;
+      if (defaultIndex == -1) {
+        return null;
+      }
+      return declaredStringLiterals.get(defaultIndex);
     }
 
     public int getDefaultIndex() {
       return defaultIndex;
     }
 
-    public Set<String> getStringLiterals() {
-      return stringLiterals;
+    public List<String> getDeclaredStringLiterals() {
+      return declaredStringLiterals;
     }
   }
 

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

@@ -19,7 +19,9 @@ import com.oracle.truffle.api.CompilerDirectives;
 import com.oracle.truffle.api.frame.VirtualFrame;
 import com.oracle.truffle.api.nodes.ExplodeLoop;
 import com.oracle.truffle.api.source.SourceSection;
-import java.util.Set;
+import java.util.ArrayList;
+import java.util.List;
+import org.pkl.core.PklBugException;
 import org.pkl.core.TypeParameter;
 import org.pkl.core.ast.ExpressionNode;
 import org.pkl.core.ast.PklNode;
@@ -382,30 +384,78 @@ public abstract class UnresolvedTypeNode extends PklNode {
       CompilerDirectives.transferToInterpreter();
 
       var elementTypeNodes = new TypeNode[unresolvedElementTypeNodes.length];
+      var isUnionOfStringLiterals = true;
 
       for (var i = 0; i < elementTypeNodes.length; i++) {
         var elementTypeNode = unresolvedElementTypeNodes[i].execute(frame);
+        if (!isStringOrUnionOfStringLiterals(elementTypeNode)) {
+          isUnionOfStringLiterals = false;
+        }
         elementTypeNodes[i] = elementTypeNode;
       }
-
       // Don't use `@ExplodeLoop` variant of union type of union is too large;
       // this can easily turn into too many IR nodes; the compiler will spend too much time
       // compiling these code paths.
       //
       // This is especially true for typealiases of unions, because they get inlined into their
       // usage sites.
-      return elementTypeNodes.length > MAX_EXPLODE_LOOP
-          ? new UnionTypeNodeLooped(sourceSection, defaultIndex, elementTypeNodes)
-          : new UnionTypeNodeExploded(sourceSection, defaultIndex, elementTypeNodes);
+      var ret =
+          elementTypeNodes.length > MAX_EXPLODE_LOOP
+              ? new UnionTypeNodeLooped(sourceSection, defaultIndex, elementTypeNodes)
+              : new UnionTypeNodeExploded(sourceSection, defaultIndex, elementTypeNodes);
+      return isUnionOfStringLiterals ? unionOfStringLiterals(elementTypeNodes, ret) : ret;
+    }
+
+    private boolean isStringOrUnionOfStringLiterals(TypeNode typeNode) {
+      return typeNode instanceof StringLiteralTypeNode
+          || typeNode instanceof UnionOfStringLiteralsTypeNode
+          || typeNode instanceof TypeAliasTypeNode typeAliasTypeNode
+              && isStringOrUnionOfStringLiterals(typeAliasTypeNode.getAliasedTypeNode());
+    }
+
+    /**
+     * Adds {@code typeNode}'s string literals to {@code literals}, and returns the corresponding
+     * default index.
+     */
+    private int collectStringLiteralAndGetDefaultIndex(TypeNode typeNode, List<String> literals) {
+      if (typeNode instanceof StringLiteralTypeNode stringTypeNode) {
+        literals.add(stringTypeNode.getLiteral());
+        return 0;
+      }
+      if (typeNode instanceof UnionOfStringLiteralsTypeNode unionOfStringLiteralsTypeNode) {
+        literals.addAll(unionOfStringLiteralsTypeNode.getDeclaredStringLiterals());
+        return unionOfStringLiteralsTypeNode.getDefaultIndex();
+      }
+      if (typeNode instanceof TypeAliasTypeNode typeAliasTypeNode) {
+        return collectStringLiteralAndGetDefaultIndex(
+            typeAliasTypeNode.getAliasedTypeNode(), literals);
+      }
+      throw PklBugException.unreachableCode();
+    }
+
+    private TypeNode unionOfStringLiterals(TypeNode[] elementTypeNodes, TypeNode originalTypeNode) {
+      var stringLiterals = new ArrayList<String>(elementTypeNodes.length);
+      var computedDefaultIdx = -1;
+      for (var i = 0; i < elementTypeNodes.length; i++) {
+        var offsetBeforeElement = stringLiterals.size();
+        var elementDefaultIdx =
+            collectStringLiteralAndGetDefaultIndex(elementTypeNodes[i], stringLiterals);
+        if (i == defaultIndex) {
+          computedDefaultIdx =
+              elementDefaultIdx == -1 ? -1 : offsetBeforeElement + elementDefaultIdx;
+        }
+      }
+      return new UnionOfStringLiteralsTypeNode(
+          sourceSection, computedDefaultIdx, stringLiterals, originalTypeNode);
     }
   }
 
   public static final class UnionOfStringLiterals extends UnresolvedTypeNode {
-    private final Set<String> stringLiterals;
+    private final List<String> stringLiterals;
     private final int defaultIndex;
 
     public UnionOfStringLiterals(
-        SourceSection sourceSection, int defaultIndex, Set<String> stringLiterals) {
+        SourceSection sourceSection, int defaultIndex, List<String> stringLiterals) {
       super(sourceSection);
       this.stringLiterals = stringLiterals;
       this.defaultIndex = defaultIndex;
@@ -415,7 +465,7 @@ public abstract class UnresolvedTypeNode extends PklNode {
     public TypeNode execute(VirtualFrame frame) {
       CompilerDirectives.transferToInterpreter();
 
-      return new UnionOfStringLiteralsTypeNode(sourceSection, defaultIndex, stringLiterals);
+      return new UnionOfStringLiteralsTypeNode(sourceSection, defaultIndex, stringLiterals, null);
     }
   }
 

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

@@ -96,7 +96,7 @@ public abstract class VmTypeMismatchException extends ControlFlowException {
           || expectedType instanceof VmTypeAlias
           || expectedType instanceof VmType
           || expectedType instanceof String // string literal type
-          || expectedType instanceof Set; // union of string literal types
+          || expectedType instanceof List; // union of string literal types
 
       this.expectedType = expectedType;
     }
@@ -110,10 +110,10 @@ public abstract class VmTypeMismatchException extends ControlFlowException {
       if (expectedType instanceof String string) {
         // string literal type
         renderedType = valueFormatter.formatStringValue(string, "");
-      } else if (expectedType instanceof Set) {
+      } else if (expectedType instanceof List) {
         // union of string literal types
         @SuppressWarnings("unchecked")
-        var stringLiterals = (Set<String>) expectedType;
+        var stringLiterals = (List<String>) expectedType;
         renderedType =
             stringLiterals.stream()
                 .map((l) -> valueFormatter.formatStringValue(l, ""))
@@ -123,7 +123,7 @@ public abstract class VmTypeMismatchException extends ControlFlowException {
       }
 
       if (actualValue instanceof VmNull
-          || actualValue instanceof String && expectedType instanceof Set) {
+          || actualValue instanceof String && expectedType instanceof List) {
         builder.append(
             ErrorMessages.createIndented(
                 "typeMismatchValue", indent, renderedType, new ProgramValue("", actualValue)));

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

@@ -0,0 +1,11 @@
+import "pkl:reflect"
+
+typealias X = *"b"|"c"
+
+hidden foo: "a"|*X
+
+local fooType = reflect.Module(module).moduleClass.properties["foo"].type
+
+res1 = fooType is reflect.UnionType
+res2 = (fooType as reflect.UnionType).members[1] is reflect.DeclaredType
+res3 = ((fooType as reflect.UnionType).members[1] as reflect.DeclaredType).referent is reflect.TypeAlias

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

@@ -0,0 +1,35 @@
+import "pkl:test"
+
+typealias A = "a"
+typealias B = "a" | "b"
+typealias C = E | B
+typealias D = B | E
+typealias E = *"b" | "c"
+typealias F = "c" | *G
+typealias G = "a" | *"g" | "c"
+
+res1: "a" | "b" = "a"
+res2: "a" | "b" | *"a"
+res3: "a" | *E
+res4: "b" | *E
+res5: "a" | *F
+res6: "c" | *A | E
+res7: B | *"a"
+res8: E | B | *"a"
+res9: "x" | *G | "a"
+res10: C | *"a"
+res11: "a" | *F | "a"
+res12: *F | "a"
+res13: D | *"q"
+
+class NoDefaults {
+  prop1: C | "z"
+  prop2: E | "z"
+  prop3: *C | "a"
+}
+
+local noDefaults: NoDefaults = new {}
+
+res14 = test.catch(() -> noDefaults.prop1)
+res15 = test.catch(() -> noDefaults.prop2)
+res16 = test.catch(() -> noDefaults.prop3)

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

@@ -0,0 +1,3 @@
+res1 = true
+res2 = true
+res3 = true

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

@@ -0,0 +1,16 @@
+res1 = "a"
+res2 = "a"
+res3 = "b"
+res4 = "b"
+res5 = "g"
+res6 = "a"
+res7 = "a"
+res8 = "a"
+res9 = "g"
+res10 = "a"
+res11 = "g"
+res12 = "g"
+res13 = "q"
+res14 = "Tried to read property `prop1` but its value is undefined."
+res15 = "Tried to read property `prop2` but its value is undefined."
+res16 = "Tried to read property `prop3` but its value is undefined."