Type check returns

9 files changed, 91 insertions, 13 deletions

diff --git a/src/main/java/org/shedlang/compiler/ShedCompiler.java b/src/main/java/org/shedlang/compiler/ShedCompiler.java
index e91f6655..ee6b3501 100644
--- a/src/main/java/org/shedlang/compiler/ShedCompiler.java
+++ b/src/main/java/org/shedlang/compiler/ShedCompiler.java
@@ -39,7 +39,7 @@ public class ShedCompiler {
             .addNamespaceFields(ShedBuiltins.NAMESPACE_NAME, P.orderedMap(
                 Map.entry(ShedBuiltins.PRINT_NAME, ShedBuiltins.PRINT_TYPE)
             ))
-            .enterBlockWithDefinedVariables(P.map(
+            .enterNamespaceWithDefinedVariables(P.map(
                 Map.entry("Io", Types.effectType(Effects.IO)),
                 Map.entry("Unit", Types.metaType(Types.UNIT)),
 
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionStatementNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionStatementNode.java
index 8dfbd7cb..ac46cde4 100644
--- a/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionStatementNode.java
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedFunctionStatementNode.java
@@ -1,4 +1,4 @@
 package org.shedlang.compiler.ast.typed;
 
-public sealed interface TypedFunctionStatementNode extends TypedNode permits TypedExpressionStatementNode, TypedValNode {
+public sealed interface TypedFunctionStatementNode extends TypedNode permits TypedExpressionStatementNode, TypedReturnNode, TypedValNode {
 }
diff --git a/src/main/java/org/shedlang/compiler/ast/typed/TypedReturnNode.java b/src/main/java/org/shedlang/compiler/ast/typed/TypedReturnNode.java
new file mode 100644
index 00000000..ad6d1f08
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/ast/typed/TypedReturnNode.java
@@ -0,0 +1,9 @@
+package org.shedlang.compiler.ast.typed;
+
+import org.shedlang.compiler.sources.Source;
+
+public record TypedReturnNode(
+    TypedExpressionNode expression,
+    Source source
+) implements TypedFunctionStatementNode {
+}
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index 4f8d051b..dcb18033 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -31,7 +31,7 @@ public class TypeChecker {
     }
 
     private static boolean isSubType(Type subType, Type superType, TypeCheckerContext context) {
-        return subType.equals(superType);
+        return superType == Types.TOP || subType.equals(superType);
     }
 
     private static Type compileTimeLookup(String name, Source source, TypeCheckerContext context) {
@@ -280,7 +280,11 @@ public class TypeChecker {
             }
         }
 
-        var bodyContext = context.enterFunction(bodyEnvironment.build(), typedEffect.map(e -> e.value()).orElse(Effects.NONE));
+        var bodyContext = context.enterFunction(
+            bodyEnvironment.build(),
+            typedEffect.map(e -> e.value()).orElse(Effects.NONE),
+            Types.UNIT
+        );
 
         var bodyResult = bodyContext.mapWithContext(
             untyped.body(),
@@ -311,7 +315,7 @@ public class TypeChecker {
                 typeCheckExpressionStatement(expressionStatement, context);
 
             case UntypedReturnNode return_ ->
-                throw new UnsupportedOperationException("TODO");
+                typeCheckReturn(return_, context);
 
             case UntypedValNode val ->
                 typeCheckVal(val, context);
@@ -414,6 +418,28 @@ public class TypeChecker {
         });
     }
 
+    private static TypeCheckerResult<? extends TypedFunctionStatementNode> typeCheckReturn(
+        UntypedReturnNode untyped,
+        TypeCheckerContext context
+    ) {
+        var expressionResult = typeCheckExpression(untyped.expression(), context);
+        var expression = expressionResult.value();
+        context = expressionResult.context();
+
+        if (context.returnType().isEmpty()) {
+            // This should be syntactically impossible, but since we only have
+            // one context type, we check for this possibility and throw a
+            // reasonable error.
+            throw new CannotReturnOutsideOfFunction(untyped.source());
+        }
+
+        checkExpressionType(expression, context.returnType().get(), context);
+
+        var typed = new TypedReturnNode(expression, untyped.source());
+
+        return context.toResult(typed);
+    }
+
     private static TypeCheckerResult<TypedStringLiteralNode> typeCheckStringLiteral(
         UntypedStringLiteralNode untyped,
         TypeCheckerContext context
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
index 2703e246..151f0a0b 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeCheckerContext.java
@@ -14,6 +14,7 @@ public class TypeCheckerContext {
     public static final TypeCheckerContext EMPTY = new TypeCheckerContext(
         Empty.map(),
         Effects.NONE,
+        Optional.empty(),
         Set.of(),
         FreeVariables.builder(),
         Empty.map()
@@ -21,6 +22,7 @@ public class TypeCheckerContext {
 
     private final PMap<String, Variable> environment;
     private final Effect effect;
+    private final Optional<Type> returnType;
     private final Set<String> boundVariables;
     private final FreeVariables.Builder freeVariables;
     private final PMap<NamespaceName, NamespaceTypeInfo> namespaceTypeInfo;
@@ -41,37 +43,48 @@ public class TypeCheckerContext {
     public TypeCheckerContext(
         PMap<String, Variable> environment,
         Effect effect,
+        Optional<Type> returnType,
         Set<String> boundVariables,
         FreeVariables.Builder freeVariables,
         PMap<NamespaceName, NamespaceTypeInfo> namespaceTypeInfo
     ) {
         this.environment = environment;
         this.effect = effect;
+        this.returnType = returnType;
         this.boundVariables = boundVariables;
         this.freeVariables = freeVariables;
         this.namespaceTypeInfo = namespaceTypeInfo;
     }
 
-    public TypeCheckerContext enterBlockWithDefinedVariables(PMap<String, Type> variables) {
+    public TypeCheckerContext enterNamespaceWithDefinedVariables(PMap<String, Type> variables) {
         var environment = variables.entrySet().stream()
             .map(x -> Map.entry(x.getKey(), TypeCheckerContext.defined(x.getValue())))
             .collect(PCollectors.toMap());
 
-        return enterBlock(environment);
+        return enterNamespace(environment);
     }
 
-    public TypeCheckerContext enterBlock(PMap<String, Variable> environment) {
-        return enter(environment, this.effect);
+    public TypeCheckerContext enterNamespace(PMap<String, Variable> environment) {
+        return enter(environment, this.effect, this.returnType);
     }
 
-    public TypeCheckerContext enterFunction(PMap<String, Variable> environment, Effect effect) {
-        return enter(environment, effect);
+    public TypeCheckerContext enterFunction(
+        PMap<String, Variable> environment,
+        Effect effect,
+        Type returnType
+    ) {
+        return enter(environment, effect, Optional.of(returnType));
     }
 
-    private TypeCheckerContext enter(PMap<String, Variable> environment, Effect effect) {
+    private TypeCheckerContext enter(
+        PMap<String, Variable> environment,
+        Effect effect,
+        Optional<Type> returnType
+    ) {
         return new TypeCheckerContext(
             this.environment.plusAll(environment),
             effect,
+            returnType,
             environment.keySet(),
             FreeVariables.builder(),
             this.namespaceTypeInfo
@@ -95,9 +108,12 @@ public class TypeCheckerContext {
         Type type
     ) {
         // TODO: check if already defined
+        // TODO: check is not defined
+        // TODO: check is defined in this scope
         return new TypeCheckerContext(
             this.environment.plus(name, new Variable.Defined(type)),
             this.effect,
+            this.returnType,
             this.boundVariables,
             this.freeVariables,
             this.namespaceTypeInfo
@@ -107,11 +123,16 @@ public class TypeCheckerContext {
     public Effect effect() {
         return effect;
     }
+
+    public Optional<Type> returnType() {
+        return returnType;
+    }
     
     private TypeCheckerContext withFreeVariable(String name) {
         return new TypeCheckerContext(
             environment,
             this.effect,
+            this.returnType,
             boundVariables,
             freeVariables.plus(name),
             this.namespaceTypeInfo
@@ -131,6 +152,7 @@ public class TypeCheckerContext {
         return new TypeCheckerContext(
             this.environment,
             this.effect,
+            this.returnType,
             this.boundVariables,
             this.freeVariables,
             this.namespaceTypeInfo.plus(namespaceName, namespaceTypeInfo)
diff --git a/src/main/java/org/shedlang/compiler/typechecker/errors/CannotReturnOutsideOfFunction.java b/src/main/java/org/shedlang/compiler/typechecker/errors/CannotReturnOutsideOfFunction.java
new file mode 100644
index 00000000..8b4a3c18
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/typechecker/errors/CannotReturnOutsideOfFunction.java
@@ -0,0 +1,10 @@
+package org.shedlang.compiler.typechecker.errors;
+
+import org.shedlang.compiler.errors.SourceError;
+import org.shedlang.compiler.sources.Source;
+
+public class CannotReturnOutsideOfFunction extends SourceError {
+    public CannotReturnOutsideOfFunction(Source source) {
+        super("Cannot return outside of function", source);
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/types/TopType.java b/src/main/java/org/shedlang/compiler/types/TopType.java
new file mode 100644
index 00000000..2356f0c5
--- /dev/null
+++ b/src/main/java/org/shedlang/compiler/types/TopType.java
@@ -0,0 +1,10 @@
+package org.shedlang.compiler.types;
+
+public record TopType() implements Type {
+    public static final TopType INSTANCE = new TopType();
+
+    @Override
+    public String describe() {
+        return "Top";
+    }
+}
diff --git a/src/main/java/org/shedlang/compiler/types/Type.java b/src/main/java/org/shedlang/compiler/types/Type.java
index 5cb781a3..72849de1 100644
--- a/src/main/java/org/shedlang/compiler/types/Type.java
+++ b/src/main/java/org/shedlang/compiler/types/Type.java
@@ -6,5 +6,5 @@ package org.shedlang.compiler.types;
  * found through auxiliary data.
  */
 public sealed interface Type extends TypeLevelValue
-    permits BoolType, FunctionType, IntType, NamespaceType, StringType, TypeLevelValueType, UnitType {
+    permits BoolType, FunctionType, IntType, NamespaceType, StringType, TopType, TypeLevelValueType, UnitType {
 }
diff --git a/src/main/java/org/shedlang/compiler/types/Types.java b/src/main/java/org/shedlang/compiler/types/Types.java
index ad9d0547..1cf17a0b 100644
--- a/src/main/java/org/shedlang/compiler/types/Types.java
+++ b/src/main/java/org/shedlang/compiler/types/Types.java
@@ -7,6 +7,7 @@ public class Types {
     public static final Type BOOL = BoolType.INSTANCE;
     public static final Type INT = IntType.INSTANCE;
     public static final Type STRING = StringType.INSTANCE;
+    public static final Type TOP = TopType.INSTANCE;
     public static final Type UNIT = UnitType.INSTANCE;
 
     public static Type metaType(Type type) {