1 files changed, 36 insertions, 7 deletions

diff --git a/src/parser.c b/src/parser.c
index 4646de0..176a4a8 100644
--- a/src/parser.c
+++ b/src/parser.c
@@ -99,6 +99,39 @@ parser_parse_type(type_t *type, parser_t *parser)
 }
 
 bool
+parser_literal_integer_node(ast_node_t *node, token_t *token)
+{
+  char number_as_str[token->value.size];
+  string_view_to_str(&token->value, number_as_str);
+
+  ast_literal_integer_create(node, atoi(number_as_str));
+  return true;
+}
+
+bool
+parser_parse_next_expression(parser_t *parser, ast_node_t *node)
+{
+  token_t token;
+  lexer_next_token(parser->lexer, &token);
+
+  if (token.kind == TOKEN_NUMBER)
+  {
+    return parser_literal_integer_node(node, &token);
+  }
+
+  parser_error_t error;
+  error.token = token;
+  sprintf(
+      error.message,
+      "unexpected '%s ("SVFMT")' token",
+      token_kind_to_str(token.kind),
+      SVARG(&token.value)
+  );
+  parser->errors[parser->errors_len++] = error;
+  return false;
+}
+
+bool
 parser_parse_return_stmt(parser_t *parser, ast_node_t *node)
 {
   token_t return_keyword_token;
@@ -121,17 +154,13 @@ parser_parse_return_stmt(parser_t *parser, ast_node_t *node)
     return false;
   }
 
-  token_t number_token;
-
-  if (!expected_token(&number_token, parser, TOKEN_NUMBER)) return false;
+  ast_node_t *argument_token = ast_node_new();
+  if (!parser_parse_next_expression(parser, argument_token)) return false;
 
   if (!drop_expected_token(parser, TOKEN_SEMICOLON)) return false;
   if (!drop_expected_token(parser, TOKEN_CCURLY)) return false;
 
-  char number_as_str[number_token.value.size];
-  string_view_to_str(&number_token.value, number_as_str);
-
-  ast_node_init_return_stmt(node, atoi(number_as_str));
+  ast_node_init_return_stmt(node, argument_token);
   return true;
 }