Update commonSupertype to consider sum types

4 files changed, 48 insertions, 7 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerSubTyping.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerSubTyping.java
index fd4ea432..dd12d06b 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerSubTyping.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerSubTyping.java
@@ -17,7 +17,7 @@ import java.util.stream.Collectors;
 public class TypeCheckerSubTyping {
     private TypeCheckerSubTyping() {}
 
-    static boolean isSubType(Type subtype, Type supertype, TypeInfos typeInfos) {
+    public static boolean isSubType(Type subtype, Type supertype, TypeInfos typeInfos) {
         var constraints = new Constraints(List.of(), P.map(), typeInfos);
         return tryAddSubtypeConstraint(subtype, supertype, constraints)
             .isPresent();
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeInfos.java b/src/main/java/org/shedlang/compiler/typechecker/TypeInfos.java
index 7812614f..ff2e0aad 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeInfos.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeInfos.java
@@ -1,8 +1,6 @@
 package org.shedlang.compiler.typechecker;
 
-import org.pcollections.Empty;
-import org.pcollections.OrderedPMap;
-import org.pcollections.PMap;
+import org.pcollections.*;
 import org.shedlang.compiler.types.*;
 import org.shedlang.compiler.util.P;
 import org.shedlang.compiler.util.PCollectors;
@@ -11,22 +9,31 @@ import java.util.Optional;
 import java.util.Set;
 
 public class TypeInfos {
-    public static final TypeInfos EMPTY = new TypeInfos(Empty.map(), P.orderedPMap(), Empty.map(), Empty.map());
+    public static final TypeInfos EMPTY = new TypeInfos(
+        Empty.map(),
+        P.orderedPMap(),
+        Empty.map(),
+        Empty.map(),
+        Empty.map()
+    );
 
     private final PMap<NamespaceName, NamespaceTypeInfo> namespaceTypeInfos;
     private final OrderedPMap<SimpleStructType, StructTypeInfo> structTypeInfos;
     private final PMap<SimpleSumType, SumTypeInfo> sumTypeInfos;
+    private final PMap<StructType, OrderedPSet<SumType>> structSupertypes;
     private final PMap<UserDefinedEffect, UserDefinedEffectInfo> userDefinedEffectInfos;
 
     public TypeInfos(
         PMap<NamespaceName, NamespaceTypeInfo> namespaceTypeInfos,
         OrderedPMap<SimpleStructType, StructTypeInfo> structTypeInfos,
         PMap<SimpleSumType, SumTypeInfo> sumTypeInfos,
+        PMap<StructType, OrderedPSet<SumType>> structSupertypes,
         PMap<UserDefinedEffect, UserDefinedEffectInfo> userDefinedEffectInfos
     ) {
         this.namespaceTypeInfos = namespaceTypeInfos;
         this.structTypeInfos = structTypeInfos;
         this.sumTypeInfos = sumTypeInfos;
+        this.structSupertypes = structSupertypes;
         this.userDefinedEffectInfos = userDefinedEffectInfos;
     }
 
@@ -38,6 +45,7 @@ public class TypeInfos {
             this.namespaceTypeInfos.plus(namespaceName, namespaceTypeInfo),
             this.structTypeInfos,
             this.sumTypeInfos,
+            this.structSupertypes,
             this.userDefinedEffectInfos
         );
     }
@@ -54,6 +62,7 @@ public class TypeInfos {
             this.namespaceTypeInfos,
             this.structTypeInfos.plus(type, structTypeInfo),
             this.sumTypeInfos,
+            this.structSupertypes,
             this.userDefinedEffectInfos
         );
     }
@@ -94,10 +103,20 @@ public class TypeInfos {
     }
 
     public TypeInfos addSumTypeInfo(SimpleSumType sumType, SumTypeInfo sumTypeInfo) {
+        var structSupertypes = this.structSupertypes;
+
+        for (var variant : sumTypeInfo.variants()) {
+            structSupertypes = structSupertypes.plus(
+                variant,
+                structSupertypes.getOrDefault(variant, P.orderedSet()).plus(sumType)
+            );
+        }
+
         return new TypeInfos(
             this.namespaceTypeInfos,
             this.structTypeInfos,
             this.sumTypeInfos.plus(sumType, sumTypeInfo),
+            structSupertypes,
             this.userDefinedEffectInfos
         );
     }
@@ -112,6 +131,14 @@ public class TypeInfos {
         };
     }
 
+    public Set<SumType> supertypes(Type type) {
+        if (type instanceof StructType structType) {
+            return this.structSupertypes.getOrDefault(structType, P.orderedSet());
+        } else {
+            return P.set();
+        }
+    }
+
     private Optional<SumTypeInfo> simpleSumTypeInfo(SimpleSumType type) {
         return Optional.ofNullable(sumTypeInfos.get(type));
     }
@@ -136,6 +163,7 @@ public class TypeInfos {
             this.namespaceTypeInfos,
             this.structTypeInfos,
             this.sumTypeInfos,
+            this.structSupertypes,
             this.userDefinedEffectInfos.plus(effect, effectInfo)
         );
     }
diff --git a/src/main/java/org/shedlang/compiler/types/Types.java b/src/main/java/org/shedlang/compiler/types/Types.java
index 474d3a26..2e5f2562 100644
--- a/src/main/java/org/shedlang/compiler/types/Types.java
+++ b/src/main/java/org/shedlang/compiler/types/Types.java
@@ -5,6 +5,8 @@ import org.shedlang.compiler.typechecker.TypeInfos;
 import java.util.List;
 import java.util.stream.Collectors;
 
+import static org.shedlang.compiler.typechecker.TypeCheckerSubTyping.isSubType;
+
 public class Types {
     private Types() {
     }
@@ -87,9 +89,15 @@ public class Types {
             return BOTTOM;
         } else if (uniqueSubtypes.size() == 1) {
             return uniqueSubtypes.iterator().next();
-        } else {
-            return TOP;
         }
+
+        for (var supertype : typeInfos.supertypes(uniqueSubtypes.iterator().next())) {
+            if (uniqueSubtypes.stream().allMatch(subtype -> isSubType(subtype, supertype, typeInfos))) {
+                return supertype;
+            }
+        }
+
+        return TOP;
     }
 
     public static Type commonSubtype(List<Type> supertypes) {
diff --git a/src/main/java/org/shedlang/compiler/util/P.java b/src/main/java/org/shedlang/compiler/util/P.java
index 5d40d596..deff0c9a 100644
--- a/src/main/java/org/shedlang/compiler/util/P.java
+++ b/src/main/java/org/shedlang/compiler/util/P.java
@@ -43,6 +43,11 @@ public class P {
     }
 
     @SafeVarargs
+    public static <T> PSet<T> set(T... values) {
+        return HashTreePSet.from(Arrays.asList(values));
+    }
+
+    @SafeVarargs
     public static <T> PStack<T> stack(T... values) {
         return ConsPStack.from(Arrays.asList(values));
     }