Typecheck return values in the entire function body

1 files changed, 81 insertions, 8 deletions

diff --git a/src/main.c b/src/main.c
index e93a538..acc2b46 100644
--- a/src/main.c
+++ b/src/main.c
@@ -201,6 +201,70 @@ emit_semantic_error(Analyzer *a, Node *n, Str msg) {
     eprintln("%s:%d:%d: error: %s", a->file_name, n->line, n->col, msg);
 }
 
+bool
+typecheck_returns(Node *node, Str expected) {
+    if (!node) {
+        return true;
+    }
+
+    // Traverse the tree again.
+    switch (node->kind) {
+        case NODE_COND:
+        case NODE_MATCH: {
+            for (sz i = 0; i < array_size(node->match_cases); i++) {
+                Node *next = node->match_cases[i];
+                if (!typecheck_returns(next, expected)) {
+                    return false;
+                }
+            }
+        } break;
+        case NODE_RETURN: {
+            return str_eq(node->type, expected);
+        } break;
+        case NODE_BLOCK: {
+            for (sz i = 0; i < array_size(node->elements); i++) {
+                Node *next = node->elements[i];
+                if (!typecheck_returns(next, expected)) {
+                    return false;
+                }
+            }
+        } break;
+        case NODE_IF: {
+            if (node->cond_expr) {
+                if (!typecheck_returns(node->cond_expr, expected)) {
+                    return false;
+                }
+            }
+            if (node->cond_else) {
+                if (!typecheck_returns(node->cond_else, expected)) {
+                    return false;
+                }
+            }
+        } break;
+        case NODE_SET:
+        case NODE_LET: {
+            if (node->var_val) {
+                if (!typecheck_returns(node->var_val, expected)) {
+                    return false;
+                }
+            }
+        } break;
+        default: {
+            if (node->left) {
+                if (!typecheck_returns(node->left, expected)) {
+                    return false;
+                }
+            }
+            if (node->right) {
+                if (!typecheck_returns(node->right, expected)) {
+                    return false;
+                }
+            }
+        } break;
+    }
+    return true;
+}
+
 Type
 type_inference(Analyzer *a, Node *node, TypeScope *scope) {
     assert(a);
@@ -554,21 +618,30 @@ type_inference(Analyzer *a, Node *node, TypeScope *scope) {
             ret_type = str_concat(ret_type, cstr(")"), a->storage);
             node->fun_return = ret_type;
 
-            node->type = cstr("nil");
             if (node->func_body->kind == NODE_BLOCK) {
-                node = node->func_body;
                 Type type;
-                for (sz i = 0; i < array_size(node->elements); i++) {
-                    Node *expr = node->elements[i];
+                for (sz i = 0; i < array_size(node->func_body->elements); i++) {
+                    Node *expr = node->func_body->elements[i];
                     type = type_inference(a, expr, scope);
                 }
-                node->type = type;
-                // TODO: ensure ALL return statements match the function
-                // prototype.
+                node->func_body->type = type;
             } else {
                 type_inference(a, node->func_body, scope);
             }
-            return cstr("nil");
+
+            // Ensure main body return matches the prototype.
+            if (!str_eq(node->func_body->type, ret_type)) {
+                emit_semantic_error(a, node, cstr("mismatched return types"));
+            }
+
+            // Ensure ALL return statements match the function prototype.
+            if (!typecheck_returns(node->func_body, ret_type)) {
+                emit_semantic_error(a, node, cstr("mismatched return types"));
+            }
+
+            // TODO: insert this into a functions map in the current scope? need
+            // the arity and return parameters for funcalls.
+            return node->type;
         } break;
         default: {
             emit_semantic_error(a, node,