Remove typeLevelArgs field in Constraints

1 files changed, 17 insertions, 11 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerSubTyping.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerSubTyping.java
index 1a0c0509..5ff13d5b 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerSubTyping.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerSubTyping.java
@@ -1,12 +1,13 @@
 package org.shedlang.compiler.typechecker;
 
-import org.pcollections.PMap;
+import org.pcollections.OrderedPMap;
 import org.pcollections.PVector;
 import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.typechecker.errors.UnexpectedTypeError;
 import org.shedlang.compiler.typechecker.typesets.TypeSets;
 import org.shedlang.compiler.types.*;
 import org.shedlang.compiler.util.P;
+import org.shedlang.compiler.util.PCollectors;
 
 import java.util.List;
 import java.util.Map;
@@ -18,7 +19,7 @@ public class TypeCheckerSubTyping {
     private TypeCheckerSubTyping() {}
 
     public static boolean isSubType(Type subtype, Type supertype, TypeInfos typeInfos) {
-        var constraints = new Constraints(List.of(), P.map(), typeInfos);
+        var constraints = new Constraints(P.orderedPMap(), typeInfos);
         return tryAddSubtypeConstraint(subtype, supertype, constraints)
             .isPresent();
     }
@@ -42,15 +43,21 @@ public class TypeCheckerSubTyping {
     }
 
     static Constraints constraints(List<TypeArg> typeLevelArgs, TypeInfos typeInfos) {
-        return new Constraints(typeLevelArgs, P.map(), typeInfos);
+        var typeBounds = typeLevelArgs.stream()
+            .map(typeLevelArg -> Map.entry(typeLevelArg, TypeBounds.EMPTY))
+            .collect(PCollectors.toOrderedPMap());
+
+        return new Constraints(typeBounds, typeInfos);
     }
 
     record Constraints(
-        // TODO: typeLevelArgs are redundant if we make typeBounds always have those as keys?
-        List<TypeArg> typeLevelArgs,
-        PMap<TypeArg, TypeBounds> typeBounds,
+        OrderedPMap<TypeArg, TypeBounds> typeBounds,
         TypeInfos typeInfos
     ) {
+        public boolean isBoundTypeArg(TypeArg typeArg) {
+            return this.typeBounds.containsKey(typeArg);
+        }
+
         public Constraints addLowerBound(TypeArg typeArg, Type lowerBound) {
             return updateBounds(typeArg, bounds -> bounds.addLowerBound(lowerBound));
         }
@@ -63,9 +70,8 @@ public class TypeCheckerSubTyping {
             TypeArg typeArg,
             Function<TypeBounds, TypeBounds> update
         ) {
-            var bounds = update.apply(this.typeBounds.getOrDefault(typeArg, TypeBounds.EMPTY));
+            var bounds = update.apply(this.typeBounds.get(typeArg));
             return new Constraints(
-                this.typeLevelArgs,
                 this.typeBounds.plus(typeArg, bounds),
                 this.typeInfos
             );
@@ -75,7 +81,7 @@ public class TypeCheckerSubTyping {
             // TODO: handle type params that are constrained by other type params
             // TODO: handle unconstrained types
 
-            var map = this.typeLevelArgs.stream()
+            var map = this.typeBounds.keySet().stream()
                 .map(typeArg -> {
                     var typeParamBounds = this.typeBounds.getOrDefault(typeArg, TypeBounds.EMPTY);
 
@@ -149,7 +155,7 @@ public class TypeCheckerSubTyping {
 
         if (
             supertype instanceof TypeArg supertypeTypeArg &&
-            constraints.typeLevelArgs.contains(supertypeTypeArg)
+            constraints.isBoundTypeArg(supertypeTypeArg)
         ) {
             constraints = constraints.addLowerBound(supertypeTypeArg, subtype);
 
@@ -158,7 +164,7 @@ public class TypeCheckerSubTyping {
 
         if (
             subtype instanceof TypeArg subtypeTypeArg &&
-            constraints.typeLevelArgs.contains(subtypeTypeArg)
+            constraints.isBoundTypeArg(subtypeTypeArg)
         ) {
             constraints = constraints.addUpperBound(subtypeTypeArg, supertype);