returning functions now work

2 files changed, 11 insertions, 5 deletions

diff --git a/src/compiler/ast-tree.c b/src/compiler/ast-tree.c
index 77b93d6..4c6d36d 100644
--- a/src/compiler/ast-tree.c
+++ b/src/compiler/ast-tree.c
@@ -2897,7 +2897,11 @@ bool setTypesReturn(AstTree *tree, AstTreeSetTypesHelper _helper,
     };
     if (!setAllTypes(metadata->value, helper, NULL)) {
       return false;
-    } else if (!typeIsEqual(metadata->value->type, function->returnType)) {
+    }
+    if (helper.lookingType != function->returnType) {
+      astTreeDelete(helper.lookingType); // TODO: change plan
+    }
+    if (!typeIsEqual(metadata->value->type, function->returnType)) {
       printError(tree->str_begin, tree->str_end, "Type mismatch");
       return false;
     }
diff --git a/src/compiler/parser.c b/src/compiler/parser.c
index 12985e1..69a9c8a 100644
--- a/src/compiler/parser.c
+++ b/src/compiler/parser.c
@@ -1156,7 +1156,6 @@ ParserNode *parserFunction(LexerNode *node, LexerNode *begin, LexerNode *end,
                            ParserNode *parent) {
   LexerNode *paramsNode = node - 1;
   LexerNode *retTypeNode = node + 1;
-  LexerNode *bodyNode = node + 2;
   if (paramsNode < begin || paramsNode->parserNode == NULL) {
     printError(node->str_begin, node->str_end, "No params");
     return NULL;
@@ -1166,12 +1165,15 @@ ParserNode *parserFunction(LexerNode *node, LexerNode *begin, LexerNode *end,
   }
   ParserNode *params = getUntilCommonParent(paramsNode->parserNode, parent);
   ParserNode *retType = getUntilCommonParent(retTypeNode->parserNode, parent);
-  ParserNode *body;
 
-  if (bodyNode >= end || bodyNode->parserNode == NULL) {
+  LexerNode *bodyNode =
+      getNextLexerNodeUsingCommonParent(retTypeNode, end, parent);
+
+  ParserNode *body;
+  if (bodyNode == NULL || bodyNode >= end || bodyNode->parserNode == NULL) {
     body = NULL;
   } else {
-    body = getUntilCommonParent(bodyNode->parserNode, parent);
+    body = bodyNode->parserNode;
     if (body == NULL || body->token != PARSER_TOKEN_SYMBOL_CURLY_BRACKET) {
       printError(node->str_begin, node->str_end, "Bad body");
       return NULL;