[PATCH olang 5/5] codegen: preserve function's variable stack location

[Carlos Maniero <carlos@maniero.me>]

Since the x86 expressions uses stack instructions (push/pop), a prologue
is required to ensure the space reserved for variables are not modified
by these instructions.

Signed-off-by: Carlos Maniero <carlos@maniero.me>
Co-authored-by: Johnny Richard <johnny@johnnyrichard.com>
---
 src/codegen_linux_x86_64.c                    | 43 +++++++++++++++++++
 .../tests/0025_var_definition_nested.ol       | 35 +++++++++++++++
 2 files changed, 78 insertions(+)
 create mode 100644 tests/integration/tests/0025_var_definition_nested.ol

diff --git a/src/codegen_linux_x86_64.c b/src/codegen_linux_x86_64.c
index 640c9fb..3ce11a7 100644
--- a/src/codegen_linux_x86_64.c
+++ b/src/codegen_linux_x86_64.c
@@ -346,6 +346,7 @@ codegen_linux_x86_64_emit_block(codegen_x86_64_t *codegen, ast_block_t *block)
 
                 codegen_linux_x86_64_emit_expression(codegen, expr);
 
+                fprintf(codegen->out, "    mov %%rbp, %%rsp\n");
                 fprintf(codegen->out, "    ret\n");
 
                 break;
@@ -418,6 +419,40 @@ codegen_linux_x86_64_emit_block(codegen_x86_64_t *codegen, ast_block_t *block)
     codegen->base_offset = block_offset;
 }
 
+static size_t
+calculate_fn_local_size(scope_t *scope)
+{
+    assert(scope);
+
+    size_t local_size = 0;
+
+    map_kv_t kvs[scope->symbols->size];
+
+    map_get_kvs(scope->symbols, (map_kv_t **)kvs);
+
+    for (size_t i = 0; i < scope->symbols->size; ++i) {
+        // FIXME: symbols must have their types. Since we just have 8bytes
+        //        variables it is hard coded.
+        local_size += 8;
+    }
+
+    size_t max_child_local_size = 0;
+
+    list_item_t *item = list_head(scope->children);
+
+    while (item != NULL) {
+        size_t child_local_size = calculate_fn_local_size((scope_t *)item->value);
+
+        if (child_local_size > max_child_local_size) {
+            max_child_local_size = child_local_size;
+        }
+
+        item = list_next(item);
+    }
+
+    return local_size + max_child_local_size;
+}
+
 static void
 codegen_linux_x86_64_emit_function(codegen_x86_64_t *codegen, ast_fn_definition_t *fn)
 {
@@ -427,6 +462,14 @@ codegen_linux_x86_64_emit_function(codegen_x86_64_t *codegen, ast_fn_definition_
 
     fprintf(codegen->out, "    mov %%rsp, %%rbp\n");
 
+    size_t local_size = calculate_fn_local_size(fn->scope);
+
+    // TODO: get the local_size from function scope
+
+    if (local_size != 0) {
+        fprintf(codegen->out, "    sub $%ld, %%rsp\n", local_size);
+    }
+
     assert(block_node->kind == AST_NODE_BLOCK);
     ast_block_t block = block_node->as_block;
 
diff --git a/tests/integration/tests/0025_var_definition_nested.ol b/tests/integration/tests/0025_var_definition_nested.ol
new file mode 100644
index 0000000..fdbe903
--- /dev/null
+++ b/tests/integration/tests/0025_var_definition_nested.ol
@@ -0,0 +1,35 @@
+# Copyright (C) 2024 olang mantainers
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+fn main(): u32 {
+  var a: u32 = 1
+
+  if a == 1 {
+    var b: u32 = 43
+    var c: u32 = 2
+
+    return a + b - c
+  } else {
+    var b: u32 = 41
+
+    return a + b
+  }
+
+  return a
+}
+
+# TEST test_compile(exit_code=0)
+
+# TEST test_run_binary(exit_code=42)
-- 
2.34.1