Check function returns correct type

3 files changed, 39 insertions, 0 deletions

diff --git a/src/main/java/org/shedlang/compiler/parser/Parser.java b/src/main/java/org/shedlang/compiler/parser/Parser.java
index 4dd59eb4..c7945d5a 100644
--- a/src/main/java/org/shedlang/compiler/parser/Parser.java
+++ b/src/main/java/org/shedlang/compiler/parser/Parser.java
@@ -193,6 +193,11 @@ public class Parser {
 
         var expression = parseTopLevelExpression(tokens);
 
+        // TODO: if we're not at the top level of a function (e.g. inside an
+        // if), then we're not going to return (from the function). What's
+        // better terminology for this? Java uses yield, but that has other
+        // meanings in other languages. Perhaps explicitly distinguish
+        // "block return" and "function return"?
         var isReturn = !tokens.trySkip(TokenType.SYMBOL_SEMICOLON);
 
         var end = tokens.position();
diff --git a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
index afe7f565..b938c66e 100644
--- a/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/shedlang/compiler/typechecker/TypeChecker.java
@@ -7,6 +7,7 @@ import org.shedlang.compiler.typechecker.errors.*;
 import org.shedlang.compiler.typechecker.typesets.TypeSetTypeLevelValue;
 import org.shedlang.compiler.typechecker.typesets.TypeSets;
 import org.shedlang.compiler.types.*;
+import org.shedlang.compiler.util.Lists;
 import org.shedlang.compiler.util.PCollectors;
 
 import java.util.*;
@@ -351,6 +352,30 @@ public class TypeChecker {
         var typedBody = bodyResult.value();
         bodyContext = bodyResult.context();
 
+        var lastStatement = Lists.last(typedBody);
+        var actualReturnType = lastStatement
+            .flatMap(statement -> {
+                return switch (statement) {
+                    case TypedExpressionStatementNode typedExpressionStatement ->
+                        Optional.empty();
+
+                    case TypedReturnNode typedReturn ->
+                        Optional.of(typedReturn.expression().type());
+
+                    case TypedValNode typedVal ->
+                        Optional.empty();
+                };
+            })
+            .orElse(Types.UNIT);
+
+        if (!isSubType(actualReturnType, typedReturnType.value(), context)) {
+            throw new UnexpectedTypeError(
+                TypeSets.type(typedReturnType.value()),
+                actualReturnType,
+                untyped.source()
+            );
+        }
+
         var typed = new TypedFunctionDefinitionNode(
             untyped.name(),
             typedParams,
diff --git a/src/main/java/org/shedlang/compiler/util/Lists.java b/src/main/java/org/shedlang/compiler/util/Lists.java
index 20d39455..f9d1a538 100644
--- a/src/main/java/org/shedlang/compiler/util/Lists.java
+++ b/src/main/java/org/shedlang/compiler/util/Lists.java
@@ -2,11 +2,20 @@ package org.shedlang.compiler.util;
 
 import java.util.Iterator;
 import java.util.List;
+import java.util.Optional;
 
 public class Lists {
     private Lists() {
     }
 
+    public static <T> Optional<T> last(List<T> list) {
+        if (list.isEmpty()) {
+            return Optional.empty();
+        } else {
+            return Optional.of(list.get(list.size() - 1));
+        }
+    }
+
     public static <T> Iterable<T> reversed(List<T> list) {
         return new Iterable<T>() {
             @Override