Implement subtyping for function types

3 files changed, 65 insertions, 18 deletions

diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 2991c344..b115f529 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -18,6 +18,7 @@ import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
+import static org.shedlang.compiler.typechecker.TypeCheckerSubTyping.isSubEffect;
 import static org.shedlang.compiler.typechecker.TypeCheckerSubTyping.isSubType;
 
 public class TypeChecker {
@@ -39,10 +40,6 @@ public class TypeChecker {
         }
     }
 
-    private static boolean isSubEffect(Effect subEffect, Effect superEffect, TypeCheckerContext context) {
-        return subEffect == Effects.NONE || subEffect.equals(superEffect);
-    }
-
     private static Type compileTimeLookup(String name, Source source, TypeCheckerContext context) {
         var variable = context.compileTimeLookup(name)
             .orElseThrow(() -> new UnknownVariableError(name, source));
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerSubTyping.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerSubTyping.java
index e0d6fe78..ea86f964 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerSubTyping.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerSubTyping.java
@@ -1,33 +1,81 @@
 package org.shedlang.compiler.typechecker;
 
 import org.shedlang.compiler.sources.Sources;
-import org.shedlang.compiler.types.SumType;
-import org.shedlang.compiler.types.Type;
-import org.shedlang.compiler.types.Types;
+import org.shedlang.compiler.types.*;
+
+import java.util.stream.Collectors;
 
 public class TypeCheckerSubTyping {
     private TypeCheckerSubTyping() {
     }
 
-    static boolean isSubType(Type subType, Type superType, TypeCheckerContext context) {
-        if (superType == Types.TOP) {
+    static boolean isSubType(Type subtype, Type supertype, TypeCheckerContext context) {
+        if (supertype == Types.TOP) {
             return true;
         }
 
-        if (subType == Types.BOTTOM) {
+        if (subtype == Types.BOTTOM) {
             return true;
         }
 
-        if (subType.equals(superType)) {
+        if (subtype.equals(supertype)) {
             return true;
         }
 
-        if (superType instanceof SumType superTypeSum) {
-            var superTypeInfo = context.sumTypeInfo(superTypeSum, Sources.NULL);
+        if (supertype instanceof SumType supertypeSum) {
+            var superTypeInfo = context.sumTypeInfo(supertypeSum, Sources.NULL);
             return superTypeInfo.variants().stream()
-                .anyMatch(superTypeVariant -> isSubType(subType, superTypeVariant, context));
+                .anyMatch(supertypeVariant -> isSubType(subtype, supertypeVariant, context));
+        }
+
+        if (
+            supertype instanceof FunctionType supertypeFunction &&
+            subtype instanceof FunctionType subtypeFunction
+        ) {
+            if (subtypeFunction.positionalParams().size() != supertypeFunction.positionalParams().size()) {
+                return false;
+            }
+
+            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 subtypeNamedParams = subtypeFunction.namedParams().stream()
+                .collect(Collectors.toMap(x -> x.name(), x -> x));
+            var supertypeNamedParams = supertypeFunction.namedParams().stream()
+                .collect(Collectors.toMap(x -> x.name(), x -> x));
+
+            if (!subtypeNamedParams.keySet().equals(supertypeNamedParams.keySet())) {
+                return false;
+            }
+
+            for (var paramName : subtypeNamedParams.keySet()) {
+                var subtypeParamType = subtypeNamedParams.get(paramName).type();
+                var supertypeParamType = supertypeNamedParams.get(paramName).type();
+                if (!isSubType(supertypeParamType, subtypeParamType, context)) {
+                    return false;
+                }
+            }
+
+            if  (!isSubEffect(subtypeFunction.effect(), supertypeFunction.effect(), context)) {
+                return false;
+            }
+
+            if  (!isSubType(subtypeFunction.returnType(), supertypeFunction.returnType(), context)) {
+                return false;
+            }
+
+            return true;
         }
 
         return false;
     }
+
+    static boolean isSubEffect(Effect subEffect, Effect superEffect, TypeCheckerContext context) {
+        return subEffect == Effects.NONE || subEffect.equals(superEffect);
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/types/FunctionType.java b/src/main/java/org/shedlang/compiler/types/FunctionType.java
index c52de92e..285dadf8 100644
--- a/src/main/java/org/shedlang/compiler/types/FunctionType.java
+++ b/src/main/java/org/shedlang/compiler/types/FunctionType.java
@@ -4,6 +4,8 @@ import org.pcollections.Empty;
 import org.pcollections.PVector;
 
 import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
 
 import static org.shedlang.compiler.util.Iterables.forEachInterspersed;
 
@@ -13,10 +15,6 @@ public record FunctionType(
     Effect effect,
     Type returnType
 ) implements Type {
-    public static Builder builder() {
-        return new Builder(Empty.vector(), Empty.vector(), Effects.NONE, Types.UNIT);
-    }
-
     @Override
     public String describe() {
         var description = new StringBuilder();
@@ -56,6 +54,10 @@ public record FunctionType(
         return description.toString();
     }
 
+    public static Builder builder() {
+        return new Builder(Empty.vector(), Empty.vector(), Effects.NONE, Types.UNIT);
+    }
+
     public record Builder(
         PVector<Type> positionalParams,
         PVector<NamedParam> namedParams,