parser: Implement variable assignment

This commit introduces variable assignment making it possible to
change a variable value. Example:

  myvar: i32 = 1;
  myvar = 2;

Signed-off-by: Johnny Richard <johnny@johnnyrichard.com>
Co-authored-by: Carlos Maniero <carlos@maniero.me>

6 files changed, 148 insertions, 23 deletions

diff --git a/examples/variables.pipa b/examples/variables.pipa
index 25cbd59..0723782 100644
--- a/examples/variables.pipa
+++ b/examples/variables.pipa
@@ -1,6 +1,8 @@
 main(): i32 {
-  a: i32 = 12;
-  b: i32 = 32;
+  a: i32 = 13;
+  a = a - 1;
+  b: i32 = 33;
+  b = 32;
   c: i32 = 2 * (b + a);
   d: i32 = (c - 33) + 1;
   e: i32 = d;
diff --git a/src/ast.c b/src/ast.c
index 55b65fa..b460f1a 100644
--- a/src/ast.c
+++ b/src/ast.c
@@ -57,6 +57,8 @@ ast_node_destroy(ast_node_t *node)
     case AST_RETURN_STMT:
       ast_node_destroy(node->data.return_stmt.argument);
       break;
+    case AST_VARIABLE_ASSIGNMENT:
+      ast_node_destroy(node->data.variable_assignment.expression);
     case AST_LITERAL:
     case AST_IDENTIFIER:
     case AST_UNKOWN_NODE:
diff --git a/src/ast.h b/src/ast.h
index 2bd700e..5095852 100644
--- a/src/ast.h
+++ b/src/ast.h
@@ -87,6 +87,12 @@ typedef struct ast_variable_declaration_t
   ast_node_t *value;
 } ast_variable_declaration_t;
 
+typedef struct ast_variable_assignment_t
+{
+  ast_identifier_t *identifier;
+  ast_node_t *expression;
+} ast_variable_assignment_t;
+
 typedef enum
 {
   AST_BINARY_OPERATION,
@@ -96,6 +102,7 @@ typedef enum
   AST_RETURN_STMT,
   AST_UNKOWN_NODE,
   AST_VARIABLE_DECLARATION,
+  AST_VARIABLE_ASSIGNMENT,
   AST_VARIABLE
 } ast_node_kind_t;
 
@@ -106,6 +113,7 @@ typedef union
   ast_literal_t literal;
   ast_return_stmt_t return_stmt;
   ast_variable_declaration_t variable_declaration;
+  ast_variable_assignment_t variable_assignment;
   ast_identifier_t identifier;
   ast_variable_t variable;
 } ast_node_data_t;
diff --git a/src/gas_assembly_generator.c b/src/gas_assembly_generator.c
index dd5d240..fca02ed 100644
--- a/src/gas_assembly_generator.c
+++ b/src/gas_assembly_generator.c
@@ -37,6 +37,9 @@ gas_assembly_generator_compile_return_stmt(gas_assembly_generator_t *gen, ast_re
 static void
 gas_assembly_generator_compile_variable_declaration(gas_assembly_generator_t *gen,
                                                     ast_variable_declaration_t *variable_declaration);
+static void
+gas_assembly_generator_compile_variable_assignment(gas_assembly_generator_t *gen,
+                                                   ast_variable_assignment_t *variable_assignment);
 
 static void
 gas_assembly_generator_compile_variable(gas_assembly_generator_t *gen, ast_variable_t *variable);
@@ -81,6 +84,18 @@ ref_entry_new(void)
   return entry;
 }
 
+static ref_entry_t *
+find_ref_entry(vector_t *refs, ast_identifier_t *identifier)
+{
+  for (int i = refs->size - 1; i >= 0; --i) {
+    ref_entry_t *entry = (ref_entry_t *)vector_at(refs, i);
+    if (entry->id == identifier) {
+      return entry;
+    }
+  }
+  return NULL;
+}
+
 void
 gas_assembly_generator_init(gas_assembly_generator_t *gen, FILE *stream)
 {
@@ -114,6 +129,9 @@ gas_assembly_generator_compile(gas_assembly_generator_t *gen, ast_node_t *ast)
     case AST_VARIABLE_DECLARATION:
       gas_assembly_generator_compile_variable_declaration(gen, &ast->data.variable_declaration);
       break;
+    case AST_VARIABLE_ASSIGNMENT:
+      gas_assembly_generator_compile_variable_assignment(gen, &ast->data.variable_assignment);
+      break;
     case AST_VARIABLE:
       gas_assembly_generator_compile_variable(gen, &ast->data.variable);
       break;
@@ -190,16 +208,42 @@ gas_assembly_generator_compile_variable_declaration(gas_assembly_generator_t *ge
 }
 
 static void
-gas_assembly_generator_compile_variable(gas_assembly_generator_t *gen, ast_variable_t *variable)
+gas_assembly_generator_compile_variable_assignment(gas_assembly_generator_t *gen,
+                                                   ast_variable_assignment_t *variable_assignment)
 {
-  ref_entry_t *entry = NULL;
-  for (int i = gen->refs->size - 1; i >= 0; --i) {
-    ref_entry_t *entry_tmp = (ref_entry_t *)vector_at(gen->refs, i);
-    if (entry_tmp->id == variable->identifier) {
-      entry = entry_tmp;
+  ref_entry_t *entry = find_ref_entry(gen->refs, variable_assignment->identifier);
+  assert(entry && "reference not found");
+
+  gas_assembly_generator_compile(gen, variable_assignment->expression);
+
+  switch (gen->latest_evaluation.kind) {
+    case EVALUATION_RESULT_AS_LITERAL_INTEGER:
+      fprintf(gen->stream, "  movq $%ld, %d(%%rbp)\n", gen->latest_evaluation.data.literal_int, entry->stack_offset);
       break;
-    }
+    case EVALUATION_RESULT_ON_RAX:
+      fprintf(gen->stream, "  mov %%rax, %d(%%rbp)\n", entry->stack_offset);
+      break;
+    case EVALUATION_RESULT_ON_STACK:
+      fprintf(gen->stream, "  movq %d(%%rbp), %%rax\n", gen->latest_evaluation.data.stack_offset);
+      fprintf(gen->stream, "  movq %%rax, %d(%%rbp)\n", entry->stack_offset);
+      break;
+    case EVALUATION_RESULT_VOID:
+      assert(false && "Unexpected void result for variable declaration");
+      // FIXME: store the latest node at latest evaluation and print_ast
   }
+  // Today we dont't support chain-assignment like this:
+  //
+  //   a = b = 1;
+  //
+  // If we start supporting it, we need to change the latest evaluation kind to
+  // stack.
+  gen->latest_evaluation.kind = EVALUATION_RESULT_VOID;
+}
+
+static void
+gas_assembly_generator_compile_variable(gas_assembly_generator_t *gen, ast_variable_t *variable)
+{
+  ref_entry_t *entry = find_ref_entry(gen->refs, variable->identifier);
   assert(entry && "reference not found");
   gas_assembly_generator_set_latest_evaluation_on_stack(gen, entry->stack_offset);
 }
diff --git a/src/parser.c b/src/parser.c
index af372fb..ad0fdc9 100644
--- a/src/parser.c
+++ b/src/parser.c
@@ -236,6 +236,40 @@ parser_parse_return_stmt(parser_t *parser, ast_node_t *node)
 }
 
 static bool
+parser_parse_variable_assignment(parser_t *parser, token_t variable_token, ast_node_t *node)
+{
+  if (!drop_expected_token(parser, TOKEN_EQUAL))
+    return false;
+
+  ast_node_t *expression = ast_node_new();
+
+  if (!parser_parse_expression(parser, expression) || !drop_expected_token(parser, TOKEN_SEMICOLON)) {
+    ast_node_destroy(expression);
+    return false;
+  }
+
+  ast_node_t *variable_declaration_node = scope_get(parser->scope, variable_token.value);
+
+  if (variable_declaration_node == NULL) {
+    parser_error_t error;
+    error.token = variable_token;
+
+    sprintf(error.message, "trying to assign '" SVFMT "' before defining it.", SVARG(&variable_token.value));
+
+    parser->errors[parser->errors_len++] = error;
+    return false;
+  }
+
+  assert(variable_declaration_node->kind == AST_VARIABLE_DECLARATION);
+
+  node->kind = AST_VARIABLE_ASSIGNMENT;
+  node->data = (ast_node_data_t){ .variable_assignment = {
+                                    .identifier = &variable_declaration_node->data.variable_declaration.identifier,
+                                    .expression = expression } };
+  return true;
+}
+
+static bool
 parser_parse_variable_definition(parser_t *parser, string_view_t variable_name, ast_node_t *node)
 {
   if (!drop_expected_token(parser, TOKEN_COLON))
@@ -268,9 +302,12 @@ parser_parse_block_declarations(parser_t *parser, vector_t *body)
   token_t current_token;
   lexer_next_token(parser->lexer, &current_token);
 
+  scope_enter(parser->scope);
+
   while (current_token.kind != TOKEN_CCURLY && current_token.kind != TOKEN_EOF) {
     if (current_token.kind != TOKEN_NAME) {
       parser_error_push_unexpected_kind(parser, &current_token, TOKEN_NAME);
+      scope_leave(parser->scope);
       return false;
     }
 
@@ -280,20 +317,53 @@ parser_parse_block_declarations(parser_t *parser, vector_t *body)
 
       if (!parsed_return) {
         ast_node_destroy(return_node);
+        scope_leave(parser->scope);
         return false;
       }
 
       vector_push_back(body, return_node);
     } else {
-      ast_node_t *variable_node = ast_node_new();
-      bool parsed_variable = parser_parse_variable_definition(parser, current_token.value, variable_node);
-
-      if (!parsed_variable) {
-        ast_node_destroy(variable_node);
-        return false;
+      token_t token;
+      lexer_peek_next_token(parser->lexer, &token);
+
+      switch (token.kind) {
+        case TOKEN_COLON: {
+          ast_node_t *variable_node = ast_node_new();
+
+          if (!parser_parse_variable_definition(parser, current_token.value, variable_node)) {
+            ast_node_destroy(variable_node);
+            return false;
+          }
+
+          vector_push_back(body, variable_node);
+          break;
+        }
+        case TOKEN_EQUAL: {
+          ast_node_t *variable_assignment = ast_node_new();
+
+          if (!parser_parse_variable_assignment(parser, current_token, variable_assignment)) {
+            ast_node_destroy(variable_assignment);
+            return false;
+          }
+
+          vector_push_back(body, variable_assignment);
+          break;
+        }
+        case TOKEN_NAME:
+        case TOKEN_OPAREN:
+        case TOKEN_CPAREN:
+        case TOKEN_SEMICOLON:
+        case TOKEN_OCURLY:
+        case TOKEN_CCURLY:
+        case TOKEN_NUMBER:
+        case TOKEN_PLUS:
+        case TOKEN_MINUS:
+        case TOKEN_STAR:
+        case TOKEN_SLASH:
+        case TOKEN_EOF:
+        case TOKEN_UNKNOWN:
+          break;
       }
-
-      vector_push_back(body, variable_node);
     }
 
     lexer_next_token(parser->lexer, &current_token);
@@ -301,6 +371,7 @@ parser_parse_block_declarations(parser_t *parser, vector_t *body)
 
   if (current_token.kind != TOKEN_CCURLY) {
     parser_error_push_unexpected_kind(parser, &current_token, TOKEN_CCURLY);
+    scope_leave(parser->scope);
     return false;
   }
 
@@ -313,8 +384,10 @@ parser_parse_block_declarations(parser_t *parser, vector_t *body)
     sprintf(error.message, "expected 'return' keyword.");
 
     parser->errors[parser->errors_len++] = error;
+    scope_leave(parser->scope);
     return false;
   }
+  scope_leave(parser->scope);
   return true;
 }
 
@@ -347,14 +420,9 @@ parser_parse_function_declaration(parser_t *parser, ast_node_t *node)
   ast_node_init_function_declaration(node, func_name_token.value, return_type, body);
   scope_push(parser->scope, &node->data.function.identifier, node);
 
-  scope_enter(parser->scope);
-
   if (!parser_parse_block_declarations(parser, body)) {
-    scope_leave(parser->scope);
     return false;
   }
 
-  scope_leave(parser->scope);
-
   return true;
 }
diff --git a/test/parser_test.c b/test/parser_test.c
index 5d615ec..8f75e07 100644
--- a/test/parser_test.c
+++ b/test/parser_test.c
@@ -220,9 +220,10 @@ test_parse_basic_syntax_errors(const MunitParameter params[], void *user_data_or
   assert_parser_error("main(): i32 { return 42;", "expected '}' but got end of file");
   assert_parser_error("main(): beff { return 42; }", "type 'beff' is not defined");
   assert_parser_error("main(): i32 { return b; }", "identifier 'b' not defined");
+  assert_parser_error("main(): i32 { b = 1; return b; }", "trying to assign 'b' before defining it.");
   // FIXME: once function calls are implemented, this error should inform that
   // neither a variable or function call was found.
-  assert_parser_error("main(): i32 { oxi 42; }", "expected ':' but got 'TOKEN_NUMBER'");
+  assert_parser_error("main(): i32 { oxi 42; }", "expected 'TOKEN_NAME' but got 'TOKEN_NUMBER'");
 
   return MUNIT_OK;
 }