Add minimal type checking for function definition

10 files changed, 163 insertions, 5 deletions

diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
new file mode 100644
index 00000000..cdfde02f
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionDefinitionNode.java
@@ -0,0 +1,19 @@
+package org.shedlang.compiler.ast.typed;
+
+import org.shedlang.compiler.sources.Source;
+import org.shedlang.compiler.types.Effect;
+import org.shedlang.compiler.types.Type;
+
+import java.util.List;
+import java.util.Optional;
+
+public record TypedFunctionDefinitionNode(
+    String name,
+    TypedParamsNode params,
+    Optional<TypedTypeLevelExpressionNode<Effect>> effect,
+    TypedTypeLevelExpressionNode<Type> returnType,
+    List<TypedFunctionStatementNode> body,
+    Source source
+) implements TypedNamespaceStatementNode {
+
+}
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedNamespaceStatementNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedNamespaceStatementNode.java
new file mode 100644
index 00000000..6b9e63d9
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedNamespaceStatementNode.java
@@ -0,0 +1,4 @@
+package org.shedlang.compiler.ast.typed;
+
+public sealed interface TypedNamespaceStatementNode extends TypedNode permits TypedFunctionDefinitionNode {
+}
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedParamNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedParamNode.java
new file mode 100644
index 00000000..6c6b4765
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedParamNode.java
@@ -0,0 +1,11 @@
+package org.shedlang.compiler.ast.typed;
+
+import org.shedlang.compiler.sources.Source;
+import org.shedlang.compiler.types.Type;
+
+public record TypedParamNode(
+    String name,
+    TypedTypeLevelExpressionNode<Type> type,
+    Source source
+) implements TypedNode {
+}
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedParamsNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedParamsNode.java
new file mode 100644
index 00000000..2c864f90
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedParamsNode.java
@@ -0,0 +1,12 @@
+package org.shedlang.compiler.ast.typed;
+
+import org.shedlang.compiler.sources.Source;
+
+import java.util.List;
+
+public record TypedParamsNode(
+    List<TypedParamNode> positional,
+    List<TypedParamNode> named,
+    Source source
+) implements TypedNode {
+}
diff --git a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionDefinitionNode.java b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionDefinitionNode.java
index 61e724d5..7cdaca92 100644
--- a/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionDefinitionNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/untyped/UntypedFunctionDefinitionNode.java
@@ -1,6 +1,9 @@
 package org.shedlang.compiler.ast.untyped;
 
+import org.pcollections.Empty;
+import org.pcollections.PVector;
 import org.shedlang.compiler.sources.Source;
+import org.shedlang.compiler.sources.Sources;
 
 import java.util.List;
 import java.util.Optional;
@@ -13,4 +16,56 @@ public record UntypedFunctionDefinitionNode(
     List<UntypedFunctionStatementNode> body,
     Source source
 ) implements UntypedNamespaceStatementNode {
+    public static Builder builder() {
+        return new Builder(
+            "f",
+            Empty.vector(),
+            Empty.vector(),
+            Optional.empty(),
+            new UntypedTypeLevelReferenceNode("Unit", Sources.NULL),
+            Empty.vector()
+        );
+    }
+
+    public record Builder(
+        String name,
+        PVector<UntypedParamNode> positionalParams,
+        PVector<UntypedParamNode> namedParams,
+        Optional<UntypedTypeLevelExpressionNode> effect,
+        UntypedTypeLevelExpressionNode returnType,
+        PVector<UntypedFunctionStatementNode> body
+    ) {
+        public UntypedFunctionDefinitionNode build() {
+            return new UntypedFunctionDefinitionNode(
+                name,
+                new UntypedParamsNode(positionalParams, namedParams, Sources.NULL),
+                effect,
+                returnType,
+                body,
+                Sources.NULL
+            );
+        }
+
+        public Builder withName(String name) {
+            return new Builder(
+                name,
+                positionalParams,
+                namedParams,
+                effect,
+                returnType,
+                body
+            );
+        }
+
+        public Builder withReturnType(UntypedTypeLevelExpressionNode returnType) {
+            return new Builder(
+                name,
+                positionalParams,
+                namedParams,
+                effect,
+                returnType,
+                body
+            );
+        }
+    }
 }
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 63699416..8bb6ca23 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -11,6 +11,8 @@ import org.shedlang.compiler.types.TypeLevelValue;
 import org.shedlang.compiler.types.TypeLevelValueType;
 
 import java.util.HashSet;
+import java.util.List;
+import java.util.Optional;
 
 public class TypeChecker {
     private TypeChecker() {
@@ -161,6 +163,22 @@ public class TypeChecker {
         );
     }
 
+    private static TypedNamespaceStatementNode typeCheckFunctionDefinition(
+        UntypedFunctionDefinitionNode untyped,
+        TypeCheckerContext context
+    ) {
+        var typedReturnType = typeCheckTypeLevelExpression(Type.class, untyped.returnType(), context);
+
+        return new TypedFunctionDefinitionNode(
+            untyped.name(),
+            new TypedParamsNode(List.of(), List.of(), untyped.params().source()),
+            Optional.empty(),
+            typedReturnType,
+            List.of(),
+            untyped.source()
+        );
+    }
+
     public static TypedFunctionStatementNode typeCheckFunctionStatement(
         UntypedExpressionStatementNode untyped,
         TypeCheckerContext context
@@ -180,6 +198,16 @@ public class TypeChecker {
         );
     }
 
+    static TypedNamespaceStatementNode typeCheckNamespaceStatement(
+        UntypedNamespaceStatementNode untyped,
+        TypeCheckerContext context
+    ) {
+        return switch (untyped) {
+            case UntypedFunctionDefinitionNode functionDefinition ->
+                typeCheckFunctionDefinition(functionDefinition, context);
+        };
+    }
+
     private static TypedReferenceNode typeCheckReference(
         UntypedReferenceNode untyped,
         TypeCheckerContext context
@@ -202,17 +230,19 @@ public class TypeChecker {
         );
     }
 
-    static TypedTypeLevelExpressionNode<TypeLevelValue> typeCheckTypeLevelExpression(
+    static <T extends TypeLevelValue> TypedTypeLevelExpressionNode<T> typeCheckTypeLevelExpression(
+        Class<T> valueClass,
         UntypedTypeLevelExpressionNode untyped,
         TypeCheckerContext context
     ) {
         return switch (untyped) {
             case UntypedTypeLevelReferenceNode reference ->
-                typeCheckTypeLevelReference(reference, context);
+                typeCheckTypeLevelReference(valueClass, reference, context);
         };
     }
 
-    private static TypedTypeLevelExpressionNode<TypeLevelValue> typeCheckTypeLevelReference(
+    private static <T extends TypeLevelValue> TypedTypeLevelExpressionNode<T> typeCheckTypeLevelReference(
+        Class<T> valueClass,
         UntypedTypeLevelReferenceNode untyped,
         TypeCheckerContext context
     ) {
@@ -222,9 +252,11 @@ public class TypeChecker {
             throw new UnexpectedTypeError(TypeSets.TYPE_LEVEL_VALUE, type, untyped.source());
         }
 
+        // TODO: check instance and raise proper error
+
         return new TypedTypeLevelReferenceNode<>(
             untyped.name(),
-            typeLevelValueType.value(),
+            (T)typeLevelValueType.value(),
             untyped.source()
         );
     }
diff --git a/src/main/java/org/shedlang/compiler/types/Effect.java b/src/main/java/org/shedlang/compiler/types/Effect.java
new file mode 100644
index 00000000..257a3847
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/types/Effect.java
@@ -0,0 +1,4 @@
+package org.shedlang.compiler.types;
+
+public sealed interface Effect extends TypeLevelValue permits IoEffect {
+}
diff --git a/src/main/java/org/shedlang/compiler/types/Effects.java b/src/main/java/org/shedlang/compiler/types/Effects.java
new file mode 100644
index 00000000..820d31c5
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/types/Effects.java
@@ -0,0 +1,8 @@
+package org.shedlang.compiler.types;
+
+public class Effects {
+    private Effects() {
+    }
+
+    public static final Effect IO = IoEffect.INSTANCE;
+}
diff --git a/src/main/java/org/shedlang/compiler/types/IoEffect.java b/src/main/java/org/shedlang/compiler/types/IoEffect.java
new file mode 100644
index 00000000..7c12036e
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/types/IoEffect.java
@@ -0,0 +1,13 @@
+package org.shedlang.compiler.types;
+
+public final class IoEffect implements Effect {
+    private IoEffect() {
+    }
+
+    static final IoEffect INSTANCE = new IoEffect();
+
+    @Override
+    public String describe() {
+        return "Io";
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/types/TypeLevelValue.java b/src/main/java/org/shedlang/compiler/types/TypeLevelValue.java
index 3eaddb7f..f10232bb 100644
--- a/src/main/java/org/shedlang/compiler/types/TypeLevelValue.java
+++ b/src/main/java/org/shedlang/compiler/types/TypeLevelValue.java
@@ -1,5 +1,5 @@
 package org.shedlang.compiler.types;
 
-public sealed interface TypeLevelValue permits Type {
+public sealed interface TypeLevelValue permits Effect, Type {
     String describe();
 }