1 files changed, 17 insertions, 13 deletions

diff --git a/src/parser.c b/src/parser.c
index ca5600d..8218c87 100644
--- a/src/parser.c
+++ b/src/parser.c
@@ -221,18 +221,23 @@ parser_parse_expression(parser_t *parser, ast_node_t *node)
   return true;
 }
 
-static bool
-parser_parse_return_stmt(parser_t *parser, ast_node_t *node)
+static ast_node_t *
+parser_parse_return_stmt(parser_t *parser)
 {
   ast_node_t *argument_token = ast_node_new();
-  if (!parser_parse_expression(parser, argument_token))
-    return false;
+  if (!parser_parse_expression(parser, argument_token)) {
+    ast_node_destroy(argument_token);
+    return NULL;
+  }
 
-  if (!drop_expected_token(parser, TOKEN_SEMICOLON))
-    return false;
+  if (!drop_expected_token(parser, TOKEN_SEMICOLON)) {
+    ast_node_destroy(argument_token);
+    return NULL;
+  }
 
+  ast_node_t *node = ast_node_new();
   ast_node_init_return_stmt(node, argument_token);
-  return true;
+  return node;
 }
 
 static bool
@@ -311,19 +316,17 @@ parser_parse_block_declarations(parser_t *parser)
   vector_t *body = vector_new();
 
   while (current_token.kind != TOKEN_CCURLY && current_token.kind != TOKEN_EOF) {
-    if (current_token.kind != TOKEN_NAME) {
+    if (current_token.kind != TOKEN_NAME && current_token.kind != TOKEN_KEYWORD_RETURN) {
       parser_error_push_unexpected_kind(parser, &current_token, TOKEN_NAME);
       scope_leave(parser->scope);
       vector_destroy(body);
       return NULL;
     }
 
-    if (string_view_eq(current_token.value, string_view_from_str("return"))) {
-      ast_node_t *return_node = ast_node_new();
-      bool parsed_return = parser_parse_return_stmt(parser, return_node);
+    if (current_token.kind == TOKEN_KEYWORD_RETURN) {
+      ast_node_t *return_node = parser_parse_return_stmt(parser);
 
-      if (!parsed_return) {
-        ast_node_destroy(return_node);
+      if (return_node == NULL) {
         scope_leave(parser->scope);
         vector_destroy(body);
         return NULL;
@@ -367,6 +370,7 @@ parser_parse_block_declarations(parser_t *parser)
         case TOKEN_CCURLY:
         case TOKEN_NUMBER:
         case TOKEN_PLUS:
+        case TOKEN_KEYWORD_RETURN:
         case TOKEN_MINUS:
         case TOKEN_STAR:
         case TOKEN_SLASH: