Add basic inference of type level args

4 files changed, 245 insertions, 42 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/ReceiverType.java b/src/main/java/org/shedlang/compiler/typechecker/ReceiverType.java
new file mode 100644
index 00000000..5bd3645b
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/typechecker/ReceiverType.java
@@ -0,0 +1,34 @@
+package org.shedlang.compiler.typechecker;
+
+import org.shedlang.compiler.types.*;
+
+import java.util.List;
+
+public record ReceiverType(List<TypeParam> typeLevelParams, FunctionType type) {
+    public static ReceiverType function(FunctionType type) {
+        return new ReceiverType(List.of(), type);
+    }
+
+    public static ReceiverType genericFunction(
+        List<TypeParam> typeLevelParams,
+        FunctionType type
+    ) {
+        return new ReceiverType(typeLevelParams, type);
+    }
+
+    public List<Type> positionalParams() {
+        return type.positionalParams();
+    }
+
+    public List<NamedParam> namedParams() {
+        return type.namedParams();
+    }
+
+    public Effect effect() {
+        return type.effect();
+    }
+
+    public Type returnType() {
+        return type.returnType();
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 03e2e021..39ab0e28 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -24,7 +24,6 @@ import static org.shedlang.compiler.typechecker.TypeCheckerSubTyping.isSubEffect
 import static org.shedlang.compiler.typechecker.TypeCheckerSubTyping.isSubType;
 
 public class TypeChecker {
-
     private TypeChecker() {
     }
 
@@ -61,11 +60,21 @@ public class TypeChecker {
         });
     }
 
-    private static TypeCheckerResult<TypedArgsNode> typeCheckArgs(
+    private record TypeCheckArgsValues(
+        TypedArgsNode typedArgsNode,
+        TypeSubstitutionMap typeSubstitutionMap
+    ) {}
+
+    private static TypeCheckerResult<TypeCheckArgsValues> typeCheckArgs(
         UntypedArgsNode args,
         TypeCheckerContext context,
-        FunctionType receiverType
+        ReceiverType receiverType
     ) {
+        var constraints = TypeCheckerSubTyping.constraints(
+            receiverType.typeLevelParams(),
+            context
+        );
+
         if (args.positional().size() != receiverType.positionalParams().size()) {
             throw new WrongNumberOfPositionalArgsError(
                 receiverType.positionalParams().size(),
@@ -87,7 +96,12 @@ public class TypeChecker {
             var typedPositionalArg = typedPositionalArgResult.value();
             context = typedPositionalArgResult.context();
 
-            checkExpressionType(typedPositionalArg, paramType, context);
+            constraints = TypeCheckerSubTyping.addSubtypeConstraint(
+                typedPositionalArg.type(),
+                paramType,
+                typedPositionalArg.source(),
+                constraints
+            );
 
             typedPositionalArgs.add(typedPositionalArg);
         }
@@ -108,7 +122,12 @@ public class TypeChecker {
             var typedExpression = typedExpressionResult.value();
             context = typedExpressionResult.context();
 
-            checkExpressionType(typedExpression, param.type(), context);
+            constraints = TypeCheckerSubTyping.addSubtypeConstraint(
+                typedExpression.type(),
+                param.type(),
+                typedExpression.source(),
+                constraints
+            );
 
             if (seenNamedArgNames.contains(untypedNamedArgNode.name())) {
                 throw new CannotPassNamedArgTwiceError(untypedNamedArgNode.name(), untypedNamedArgNode.source());
@@ -128,7 +147,8 @@ public class TypeChecker {
             }
         }
 
-        return context.toResult(new TypedArgsNode(typedPositionalArgs, typedNamedArgs, args.source()));
+        var typedArgsNode = new TypedArgsNode(typedPositionalArgs, typedNamedArgs, args.source());
+        return context.toResult(new TypeCheckArgsValues(typedArgsNode, constraints.solve()));
     }
 
     private static TypeCheckSteps<TypedAssignNode> typeCheckStepsAssign(
@@ -339,7 +359,7 @@ public class TypeChecker {
                     initialReceiverType.value() instanceof TypeConstructor<?> typeConstructor
                 ) {
                     // TODO: check args
-                    return receiverTypeToFunctionType(
+                    return typeCheckCallReceiverType(
                         Types.metaType(Types.constructedType(typeConstructor, typeLevelArgs)),
                         typedReceiver.source(),
                         contextForReceiver
@@ -348,10 +368,11 @@ public class TypeChecker {
                     throw new IllegalArgumentException("TODO");
                 }
             })
-            .orElseGet(() -> receiverTypeToFunctionType(typedReceiver, contextForReceiver));
+            .orElseGet(() -> typeCheckCallReceiverType(typedReceiver, contextForReceiver));
 
         var argsResult = typeCheckArgs(untyped.args(), context, receiverType);
-        var typedArgs = argsResult.value();
+        var typedArgs = argsResult.value().typedArgsNode;
+        var typeSubstitutionMap = argsResult.value().typeSubstitutionMap;
         context = argsResult.context();
 
         var receiverHasEffects = !receiverType.effect().equals(Effects.NONE);
@@ -376,29 +397,29 @@ public class TypeChecker {
             typedReceiver,
             typedArgs,
             untyped.hasEffects(),
-            receiverType.returnType(),
+            TypeSubstitution.substitute(receiverType.returnType(), typeSubstitutionMap),
             untyped.source()
         ));
     }
 
-    private static FunctionType receiverTypeToFunctionType(
+    private static ReceiverType typeCheckCallReceiverType(
         TypedExpressionNode typedReceiver,
         TypeCheckerContext context
     ) {
-        return receiverTypeToFunctionType(
+        return typeCheckCallReceiverType(
             typedReceiver.type(),
             typedReceiver.source(),
             context
         );
     }
 
-    private static FunctionType receiverTypeToFunctionType(
+    private static ReceiverType typeCheckCallReceiverType(
         Type receiverType,
         Source receiverSource,
         TypeCheckerContext context
     ) {
         if (receiverType instanceof FunctionType receiverTypeFunction) {
-            return receiverTypeFunction;
+            return ReceiverType.function(receiverTypeFunction);
         }
 
         if (
@@ -406,7 +427,10 @@ public class TypeChecker {
             typeLevelValueType.value() instanceof StructType structType
         ) {
             var structTypeInfo = context.structTypeInfo(structType, receiverSource);
-            return receiverStructTypeInfoToFunctionType(structType, structTypeInfo);
+            return ReceiverType.function(receiverStructTypeInfoToFunctionType(
+                structType,
+                structTypeInfo
+            ));
         }
 
         if (
@@ -416,7 +440,29 @@ public class TypeChecker {
         ) {
             // TODO: make specific constructed types to avoid generic casting?
             var structTypeInfo = context.structTypeInfo((ConstructedType<StructType>) constructedType, receiverSource);
-            return receiverStructTypeInfoToFunctionType(constructedType, structTypeInfo);
+            return ReceiverType.function(receiverStructTypeInfoToFunctionType(
+                constructedType,
+                structTypeInfo
+            ));
+        }
+
+        if (
+            receiverType instanceof TypeLevelValueType typeLevelValueType &&
+            typeLevelValueType.value() instanceof TypeConstructor<?> typeConstructor &&
+            typeConstructor.genericType() instanceof StructType structType
+        ) {
+            var structTypeInfo = context.structTypeInfo(
+                (StructType) typeConstructor.genericType(),
+                receiverSource
+            );
+            var constructedType = Types.constructedType(
+                typeConstructor,
+                typeConstructor.params()
+            );
+            return ReceiverType.genericFunction(
+                typeConstructor.params(),
+                receiverStructTypeInfoToFunctionType(constructedType, structTypeInfo)
+            );
         }
 
         throw new UnexpectedTypeError(
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerSubTyping.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerSubTyping.java
index 8e0fb017..e96d433f 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerSubTyping.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerSubTyping.java
@@ -1,41 +1,141 @@
 package org.shedlang.compiler.typechecker;
 
+import org.pcollections.PMap;
+import org.pcollections.PVector;
+import org.shedlang.compiler.sources.Source;
 import org.shedlang.compiler.sources.Sources;
+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 java.util.List;
+import java.util.Map;
+import java.util.Optional;
 import java.util.stream.Collectors;
 
 public class TypeCheckerSubTyping {
-    private TypeCheckerSubTyping() {
-    }
+    private TypeCheckerSubTyping() {}
 
     static boolean isSubType(Type subtype, Type supertype, TypeCheckerContext context) {
+        var constraints = new Constraints(List.of(), P.map(), context);
+        return tryAddSubtypeConstraint(subtype, supertype, constraints)
+            .isPresent();
+    }
+
+    record TypeBounds(PVector<Type> lower, PVector<Type> upper) {
+        private static final TypeBounds EMPTY = new TypeBounds(P.vector(), P.vector());
+
+        public TypeBounds addLowerBound(Type lowerBound) {
+            return new TypeBounds(
+                this.lower.plus(lowerBound),
+                this.upper
+            );
+        }
+    }
+
+    static Constraints constraints(List<TypeParam> typeLevelParams, TypeCheckerContext context) {
+        return new Constraints(typeLevelParams, P.map(), context);
+    }
+
+    record Constraints(
+        // TODO: typeLevelParams are redundant if we make typeBounds always have those params as keys?
+        List<TypeParam> typeLevelParams,
+        PMap<TypeParam, TypeBounds> typeBounds,
+        // TODO: replace with type info
+        TypeCheckerContext context
+    ) {
+        public Constraints addLowerBound(TypeParam typeParam, Type lowerBound) {
+            var bounds = this.typeBounds.getOrDefault(typeParam, TypeBounds.EMPTY)
+                .addLowerBound(lowerBound);
+            return new Constraints(
+                this.typeLevelParams,
+                this.typeBounds.plus(typeParam, bounds),
+                this.context
+            );
+        }
+
+        public TypeSubstitutionMap solve() {
+            // TODO: handle upper bounds
+            // TODO: handle type params that are constrained by other type params
+            // TODO: handle unconstrained types
+            // TODO: do we need any special handling for mutual recursion (i.e. do we need separate type params for the parent call and the child call?)
+
+            var map = this.typeLevelParams.stream()
+                .map(typeParam -> {
+                    var typeParamBounds = this.typeBounds.getOrDefault(typeParam, TypeBounds.EMPTY);
+                    var bound = Types.BOTTOM;
+
+                    for (var lower : typeParamBounds.lower) {
+                        bound = Types.commonSupertype(List.of(bound, lower));
+                    }
+
+                    return Map.entry(typeParam, bound);
+                })
+                .collect(Collectors.toMap(x -> x.getKey(), x -> x.getValue()));
+
+            return new TypeSubstitutionMap(map);
+        }
+    }
+
+    static Constraints addSubtypeConstraint(
+        Type subtype,
+        Type supertype,
+        Source source,
+        Constraints constraints
+    ) {
+        var result = tryAddSubtypeConstraint(subtype, supertype, constraints);
+        if (result.isEmpty()) {
+            throw new UnexpectedTypeError(
+                TypeSets.type(supertype),
+                subtype,
+                source
+            );
+        } else {
+            return result.get();
+        }
+    }
+
+    private static Optional<Constraints> tryAddSubtypeConstraint(
+        Type subtype,
+        Type supertype,
+        Constraints constraints
+    ) {
         if (supertype == Types.TOP) {
-            return true;
+            return Optional.of(constraints);
         }
 
         if (subtype == Types.BOTTOM) {
-            return true;
+            return Optional.of(constraints);
         }
 
         if (subtype.equals(supertype)) {
-            return true;
+            return Optional.of(constraints);
+        }
+
+        if (
+            supertype instanceof TypeParam supertypeTypeParam &&
+            constraints.typeLevelParams.contains(supertypeTypeParam)
+        ) {
+            constraints = constraints.addLowerBound(supertypeTypeParam, subtype);
+
+            return Optional.of(constraints);
         }
 
         if (supertype instanceof SumType supertypeSum) {
-            var superTypeInfo = context.sumTypeInfo(supertypeSum, Sources.NULL);
-            return isVariantSubtype(subtype, superTypeInfo, context);
+            var superTypeInfo = constraints.context.sumTypeInfo(supertypeSum, Sources.NULL);
+            return tryAddVariantSubtypeConstraint(subtype, superTypeInfo, constraints);
         }
 
         if (
             supertype instanceof ConstructedType<?> supertypeConstructedType &&
             supertypeConstructedType.genericType() instanceof SumType
         ) {
-            var superTypeInfo = context.sumTypeInfo(
+            var superTypeInfo = constraints.context.sumTypeInfo(
                 (ConstructedType<SumType>) supertypeConstructedType,
                 Sources.NULL
             );
-            return isVariantSubtype(subtype, superTypeInfo, context);
+            return tryAddVariantSubtypeConstraint(subtype, superTypeInfo, constraints);
         }
 
         if (
@@ -43,14 +143,18 @@ public class TypeCheckerSubTyping {
             subtype instanceof FunctionType subtypeFunction
         ) {
             if (subtypeFunction.positionalParams().size() != supertypeFunction.positionalParams().size()) {
-                return false;
+                return Optional.empty();
             }
 
             for (var paramIndex = 0; paramIndex < subtypeFunction.positionalParams().size(); paramIndex++) {
                 var subtypeParamType = subtypeFunction.positionalParams().get(paramIndex);
                 var supertypeParamType = supertypeFunction.positionalParams().get(paramIndex);
-                if (!isSubType(supertypeParamType, subtypeParamType, context)) {
-                    return false;
+                var result = tryAddSubtypeConstraint(supertypeParamType, subtypeParamType, constraints);
+                if (result.isPresent()) {
+                    // TODO: test assignment
+                    constraints = result.get();
+                } else {
+                    return Optional.empty();
                 }
             }
 
@@ -60,38 +164,57 @@ public class TypeCheckerSubTyping {
                 .collect(Collectors.toMap(x -> x.name(), x -> x));
 
             if (!subtypeNamedParams.keySet().equals(supertypeNamedParams.keySet())) {
-                return false;
+                return Optional.empty();
             }
 
             for (var paramName : subtypeNamedParams.keySet()) {
                 var subtypeParamType = subtypeNamedParams.get(paramName).type();
                 var supertypeParamType = supertypeNamedParams.get(paramName).type();
-                if (!isSubType(supertypeParamType, subtypeParamType, context)) {
-                    return false;
+                var result = tryAddSubtypeConstraint(supertypeParamType, subtypeParamType, constraints);
+                if (result.isPresent()) {
+                    // TODO: test assignment
+                    constraints = result.get();
+                } else {
+                    return Optional.empty();
                 }
             }
 
-            if  (!isSubEffect(subtypeFunction.effect(), supertypeFunction.effect(), context)) {
-                return false;
+            if  (!isSubEffect(subtypeFunction.effect(), supertypeFunction.effect(), constraints.context)) {
+                return Optional.empty();
             }
 
-            if  (!isSubType(subtypeFunction.returnType(), supertypeFunction.returnType(), context)) {
-                return false;
+            var returnTypeResult = tryAddSubtypeConstraint(
+                subtypeFunction.returnType(),
+                supertypeFunction.returnType(),
+                constraints
+            );
+            if (returnTypeResult.isPresent()) {
+                // TODO: test assignment
+                constraints = returnTypeResult.get();
+            } else {
+                return Optional.empty();
             }
 
-            return true;
+            return Optional.of(constraints);
         }
 
-        return false;
+        return Optional.empty();
     }
 
-    private static boolean isVariantSubtype(
+    private static Optional<Constraints> tryAddVariantSubtypeConstraint(
         Type subtype,
         SumTypeInfo superTypeInfo,
-        TypeCheckerContext context
+        Constraints constraints
     ) {
-        return superTypeInfo.variants().stream()
-            .anyMatch(supertypeVariant -> isSubType(subtype, supertypeVariant.type(), context));
+        for (var supertypeVariant : superTypeInfo.variants()) {
+            var result = tryAddSubtypeConstraint(subtype, supertypeVariant.type(), constraints);
+            // TODO: is it possible for multiple variants to match? If so, do we need to handle that?
+            if (result.isPresent()) {
+                return result;
+            }
+        }
+
+        return Optional.empty();
     }
 
     static boolean isSubEffect(Effect subEffect, Effect superEffect, TypeCheckerContext context) {
diff --git a/src/main/java/org/shedlang/compiler/types/ConstructedType.java b/src/main/java/org/shedlang/compiler/types/ConstructedType.java
index 09cfcc83..d4f688dc 100644
--- a/src/main/java/org/shedlang/compiler/types/ConstructedType.java
+++ b/src/main/java/org/shedlang/compiler/types/ConstructedType.java
@@ -6,7 +6,7 @@ import java.util.stream.Collectors;
 
 public record ConstructedType<T extends StructuredType>(
     TypeConstructor<T> constructor,
-    List<Type> args
+    List<? extends Type> args
 ) implements Type {
     @Override
     public String describe() {