added return

5 files changed, 107 insertions, 17 deletions

diff --git a/abbrevs.txt b/abbrevs.txt
index 77176ff..8efb3e3 100644
--- a/abbrevs.txt
+++ b/abbrevs.txt
@@ -9,6 +9,7 @@ kw - keyword
 num - number
 op - operator
 ptr - pointer
+ret - return
 stmt - statement
 str - string
 tokr - tokenizer
diff --git a/parse.c b/parse.c
index 63a8637..84066ad 100644
--- a/parse.c
+++ b/parse.c
@@ -162,13 +162,20 @@ typedef struct FnExpr {
 	Declaration params;	/* declaration of the parameters to this function */
 	Type ret_type;
 	Block body;
-} FnExpr; /* an expression such as fn(x: int) int {return 2 * x;} */
+} FnExpr; /* an expression such as fn(x: int) int { 2 * x } */
 
 typedef enum {
 			  STMT_DECL,
-			  STMT_EXPR
+			  STMT_EXPR,
+			  STMT_RET
 } StatementKind;
 
+#define RET_FLAG_EXPR 0x01
+typedef struct {
+	uint16_t flags;
+	Expression expr;
+} Return;
+
 #define STMT_FLAG_VOIDED_EXPR 0x01 /* the "4;" in fn () { 4; } is a voided expression, but the "4" in fn () int { 4 } is not */
 typedef struct {
 	Location where;
@@ -177,6 +184,7 @@ typedef struct {
 	union {
 		Declaration decl;
 		Expression expr;
+		Return ret;
 	};
 } Statement;
 
@@ -1311,6 +1319,30 @@ static bool parse_stmt(Parser *p, Statement *s) {
 	/* 
 	   TODO: statements such as 3, 5; will not work.
 	*/
+	if (token_is_kw(t->token, KW_RETURN)) {
+		s->kind = STMT_RET;
+		t->token++;
+		s->ret.flags = 0;
+		if (token_is_kw(t->token, KW_SEMICOLON)) {
+			/* return with no expr */
+			t->token++;
+			return true;
+		}
+		s->ret.flags |= RET_FLAG_EXPR;
+		Token *end = expr_find_end(p, 0, NULL);
+		if (!end) {
+			while (t->token->kind != TOKEN_EOF) t->token++; /* move to end of file */
+			return false;
+		}
+		if (!token_is_kw(end, KW_SEMICOLON)) {
+			err_print(end->where, "Expected ';' at end of return statement.");
+			t->token = end->kind == TOKEN_EOF ? end : end + 1;
+			return false;
+		}
+	    bool success = parse_expr(p, &s->ret.expr, end);
+		t->token = end + 1;
+		return success;
+	}
 	if (token_is_kw(t->token + 1, KW_COLON) || token_is_kw(t->token + 1, KW_COMMA)
 		|| token_is_kw(t->token + 1, KW_AT)) {
 		s->kind = STMT_DECL;
@@ -1498,6 +1530,7 @@ static void fprint_stmt(FILE *out, Statement *s) {
 		fprintf(out, "(void)");
 	
 	switch (s->kind) {
+
 	case STMT_DECL:
 		fprint_decl(out, &s->decl);
 		fprintf(out, ";\n");
@@ -1506,6 +1539,12 @@ static void fprint_stmt(FILE *out, Statement *s) {
 		fprint_expr(out, &s->expr);
 		fprintf(out, ";\n");
 		break;
+	case STMT_RET:
+		fprintf(out, "return ");
+		if (s->ret.flags & RET_FLAG_EXPR)
+			fprint_expr(out, &s->ret.expr);
+		fprintf(out, ";\n");
+		break;
 	}
 }
 
diff --git a/test.toc b/test.toc
index 8746614..af97878 100644
--- a/test.toc
+++ b/test.toc
@@ -1,8 +1,4 @@
 main @= fn() {
-	i := 1;
-	asdf := while (fn(i: i64) i64 { i - 1024 })(i) {
-			 i = i * 2;
-			 i
-	};
+	 foo @= fn(x: int)  { return; };
 };
 
diff --git a/tokenizer.c b/tokenizer.c
index e62b4fe..1e095a3 100644
--- a/tokenizer.c
+++ b/tokenizer.c
@@ -28,9 +28,12 @@ typedef enum {
 			  KW_EQEQ,
 			  KW_LT,
 			  KW_LE,
+			  KW_GT,
+			  KW_GE,
 			  KW_PLUS,
 			  KW_MINUS,
 			  KW_ASTERISK,
+			  KW_EXCLAMATION,
 			  KW_AMPERSAND,
 			  KW_SLASH,
 			  KW_LAST_SYMBOL = KW_SLASH, /* last one entirely consisting of symbols */
@@ -38,6 +41,7 @@ typedef enum {
 			  KW_ELIF,
 			  KW_ELSE,
 			  KW_WHILE,
+			  KW_RETURN,
 			  KW_FN,
 			  KW_INT,
 			  KW_I8,
@@ -55,9 +59,10 @@ typedef enum {
 } Keyword;
 
 static const char *keywords[KW_COUNT] =
-	{";", "=", ":", "@", ",", "(", ")", "{", "}", "[", "]", "==", "<", "<=", "+", "-", "*",
-	 "&", "/", "if", "elif", "else", "while", "fn", "int", "i8", "i16", "i32", "i64", "u8",
-	 "u16", "u32", "u64", "float", "f32", "f64"};
+	{";", "=", ":", "@", ",", "(", ")", "{", "}", "[", "]", "==", "<", "<=", ">", ">=",
+	 "+", "-", "*", "!", "&", "/", "if", "elif", "else", "while", "return", "fn", "int",
+	 "i8", "i16", "i32", "i64",
+	 "u8", "u16", "u32", "u64", "float", "f32", "f64"};
 
 static const char *directives[DIRECT_COUNT] =
 	{"C"};
diff --git a/types.c b/types.c
index cdb4b89..2040f4d 100644
--- a/types.c
+++ b/types.c
@@ -1,6 +1,7 @@
 typedef struct {
 	Array in_decls;	/* array of declarations we are currently inside */
 	Block *block;
+	Type *ret_type;	/* the return type of the function we're currently parsing. NULL for none. */
 } Typer;
 
 static bool types_stmt(Typer *tr, Statement *s);
@@ -307,40 +308,69 @@ static bool type_can_be_truthy(Type *t) {
 
 static bool types_expr(Typer *tr, Expression *e) {
 	if (e->flags & EXPR_FLAG_FOUND_TYPE) return true;
+	Type *prev_ret_type = tr->ret_type;
+	
 	Type *t = &e->type;
 	t->flags = 0;
 	t->kind = TYPE_UNKNOWN; /* default to unknown type (in the case of an error) */
 	e->flags |= EXPR_FLAG_FOUND_TYPE; /* even if failed, pretend we found the type */
+	bool success = true;
 	switch (e->kind) {
 	case EXPR_FN: {
 		FnExpr *f = e->fn;
-		if (!type_of_fn(tr, f, t)) return false;
+		if (!type_of_fn(tr, f, t)) {
+			success = false;
+			goto fn_ret;
+		}
+		tr->ret_type = t->fn.types.data;
 		add_ident_decls(&f->body, &f->params);
-		bool success = true;
-		success = types_block(tr, &e->fn->body);
+		bool block_success = true;
+		block_success = types_block(tr, &e->fn->body);
 		remove_ident_decls(&f->body, &f->params);
-		if (!success) return false;
+		if (!block_success) {
+			success = false;
+			goto fn_ret;
+		}
 		Expression *ret_expr = f->body.ret_expr;
 		assert(t->kind == TYPE_FN);
 		Type *ret_type = t->fn.types.data;
 		if (ret_expr) {
-			if (!types_expr(tr, ret_expr)) return false;
+			if (!types_expr(tr, ret_expr)) {
+				success = false;
+				goto fn_ret;
+			}
 			if (!type_eq(ret_type, &ret_expr->type)) {
 				char *got = type_to_str(&ret_expr->type);
 				char *expected = type_to_str(ret_type);
 				err_print(ret_expr->where, "Returning type %s, but function returns type %s.", got, expected);
 				info_print(e->where, "Function declaration is here.");
 				free(got); free(expected);
-				return false;
+				success = false;
+				goto fn_ret;
 			}
 		} else if (ret_type->kind != TYPE_VOID) {
+			Array stmts = e->fn->body.stmts;
+			if (stmts.len) {
+				Statement *last_stmt = (Statement *)stmts.data + (stmts.len - 1);
+				if (last_stmt->kind == STMT_RET) {
+					/*
+					  last statement is a return, so it doesn't matter that the function has no return value
+					ideally this would handle if foo { return 5; } else { return 6; } */
+					success = true;
+					goto fn_ret;
+				}
+			}
 			/* TODO: this should really be at the closing brace, and not the function declaration */
 			char *expected = type_to_str(ret_type);
 			err_print(f->body.end, "No return value in function which returns %s.", expected);
 			free(expected);
 			info_print(e->where, "Function was declared here:");
-			return false;
+			success = false;
+			goto fn_ret;
 		}
+		fn_ret:
+		tr->ret_type = prev_ret_type;
+		if (!success) return false;
 	} break;
 	case EXPR_LITERAL_INT:
 	    t->kind = TYPE_BUILTIN;
@@ -701,6 +731,25 @@ static bool types_stmt(Typer *tr, Statement *s) {
 		if (!types_decl(tr, &s->decl))
 			return false;
 		break;
+	case STMT_RET:
+		if (!tr->ret_type) {
+			err_print(s->where, "return outside of a function.");
+			return false;
+		}
+		if (s->ret.flags & RET_FLAG_EXPR) {
+			if (tr->ret_type->kind == TYPE_VOID) {
+				err_print(s->where, "Return value in void function.");
+				return false;
+			}
+			if (!types_expr(tr, &s->ret.expr))
+				return false;
+		} else {
+			if (tr->ret_type->kind != TYPE_VOID) {
+				err_print(s->where, "No return value in non-void function.");
+				return false;
+			}
+		}
+		break;
 	}
 	return true;
 }