parser: Refactor return statement to return an ast_node

During the refactoring process, I identified a memory leak where the
return argument was allocated but not freed in case of an error.

It also introduces the concept of keyword tokens. Where return is now a
keyword simplifying the parser.

Signed-off-by: Carlos Maniero <carlos@maniero.me>

4 files changed, 33 insertions, 16 deletions

diff --git a/src/lexer.c b/src/lexer.c
index b60fbb5..9f2a57b 100644
--- a/src/lexer.c
+++ b/src/lexer.c
@@ -85,6 +85,7 @@ lexer_tokenize_name(lexer_t *lexer, token_t *token)
   while (lexer_is_not_eof(lexer) && isalnum(lexer_current_char(lexer))) {
     lexer_drop_char(lexer);
   }
+
   token->kind = TOKEN_NAME;
   token->value = string_view_new(lexer->src + begin, lexer->cur - begin);
   token->filepath = lexer->filepath;
@@ -93,6 +94,15 @@ lexer_tokenize_name(lexer_t *lexer, token_t *token)
   token->bol = lexer->bol;
 }
 
+static void
+lexer_token_process_keyword(token_t *token)
+{
+  if (string_view_eq(string_view_from_str("return"), token->value)) {
+    token->kind = TOKEN_KEYWORD_RETURN;
+    return;
+  }
+}
+
 void
 lexer_next_token(lexer_t *lexer, token_t *token)
 {
@@ -111,6 +121,7 @@ lexer_next_token(lexer_t *lexer, token_t *token)
 
   if (isalpha(lexer_current_char(lexer))) {
     lexer_tokenize_name(lexer, token);
+    lexer_token_process_keyword(token);
     return;
   }
 
@@ -292,11 +303,12 @@ token_kind_to_str(token_kind_t kind)
       return "/";
     case TOKEN_EQUAL:
       return "=";
+    case TOKEN_KEYWORD_RETURN:
+      return "return";
     case TOKEN_EOF:
       return "TOKEN_EOF";
     case TOKEN_UNKNOWN:
       return "TOKEN_UNKNOWN";
-    default:
-      assert(false && "unreachable");
   }
+  assert(false && "unreachable");
 }
diff --git a/src/lexer.h b/src/lexer.h
index abee53a..70d4b0d 100644
--- a/src/lexer.h
+++ b/src/lexer.h
@@ -37,6 +37,7 @@ typedef enum
   TOKEN_STAR,
   TOKEN_SLASH,
   TOKEN_EQUAL,
+  TOKEN_KEYWORD_RETURN,
   TOKEN_EOF,
   TOKEN_UNKNOWN
 } token_kind_t;
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:
diff --git a/test/parser_test.c b/test/parser_test.c
index f923c46..52437ce 100644
--- a/test/parser_test.c
+++ b/test/parser_test.c
@@ -208,7 +208,7 @@ test_parse_basic_syntax_errors(const MunitParameter params[], void *user_data_or
   assert_parser_error("main(: i32 { return 42; }", "expected ')' but got ':'");
   assert_parser_error("main() i32 { return 42; }", "expected ':' but got 'TOKEN_NAME'");
   assert_parser_error("main(): { return 42; }", "expected 'TOKEN_NAME' but got '{'");
-  assert_parser_error("main(): i32 return 42; }", "expected '{' but got 'TOKEN_NAME'");
+  assert_parser_error("main(): i32 return 42; }", "expected '{' but got 'return'");
   assert_parser_error("main(): i32 { 42; }", "expected 'TOKEN_NAME' but got 'TOKEN_NUMBER'");
   assert_parser_error("main(): i32 { return; }", "unexpected '; (;)' token");
   assert_parser_error("main(): i32 { return 42;", "expected '}' but got end of file");