From c7de22287e94df158b49811550c052dfac85b2ef Mon Sep 17 00:00:00 2001
From: David Mak <csdavidmak@ust.hk>
Date: Tue, 3 Oct 2023 18:02:45 +0800
Subject: [PATCH] core: Fix restoration of stack address

All allocas for temporary objects are now placed in the beginning of the
function. Allocas for on-temporary objects are not modified because
these variables may appear in a loop and thus must be uniquely
represented.
---
 nac3artiq/src/codegen.rs                       |  2 +-
 nac3core/src/codegen/expr.rs                   |  9 +++++----
 nac3standalone/demo/demo.c                     |  6 ++++++
 nac3standalone/demo/interpret_demo.py          |  5 +++++
 nac3standalone/demo/src/stack_addr_issue233.py | 15 +++++++++++++++
 5 files changed, 32 insertions(+), 5 deletions(-)
 create mode 100644 nac3standalone/demo/src/stack_addr_issue233.py

diff --git a/nac3artiq/src/codegen.rs b/nac3artiq/src/codegen.rs
index 4e38bdceb..112864358 100644
--- a/nac3artiq/src/codegen.rs
+++ b/nac3artiq/src/codegen.rs
@@ -385,7 +385,7 @@ fn rpc_codegen_callback_fn<'ctx, 'a>(
     }
 
     for (i, arg) in real_params.iter().enumerate() {
-        let arg_slot = ctx.builder.build_alloca(arg.get_type(), &format!("rpc.arg{}", i));
+        let arg_slot = generator.gen_var_alloc(ctx, arg.get_type(), Some(&format!("rpc.arg{i}"))).unwrap();
         ctx.builder.build_store(arg_slot, *arg);
         let arg_slot = ctx.builder.build_bitcast(arg_slot, ptr_type, "rpc.arg");
         let arg_ptr = unsafe {
diff --git a/nac3core/src/codegen/expr.rs b/nac3core/src/codegen/expr.rs
index d5602bc76..779a0e105 100644
--- a/nac3core/src/codegen/expr.rs
+++ b/nac3core/src/codegen/expr.rs
@@ -6,7 +6,7 @@ use crate::{
         get_llvm_type,
         get_llvm_abi_type,
         irrt::*,
-        stmt::gen_raise,
+        stmt::{gen_raise, gen_var},
         CodeGenContext, CodeGenTask,
     },
     symbol_resolver::{SymbolValue, ValueEnum},
@@ -357,13 +357,14 @@ impl<'ctx, 'a> CodeGenContext<'ctx, 'a> {
     }
 
     pub fn build_call_or_invoke(
-        &self,
+        &mut self,
         fun: FunctionValue<'ctx>,
         params: &[BasicValueEnum<'ctx>],
         call_name: &str,
     ) -> Option<BasicValueEnum<'ctx>> {
         let mut loc_params: Vec<BasicValueEnum<'ctx>> = Vec::new();
         let mut return_slot = None;
+
         if fun.count_params() > 0 {
             let sret_id = Attribute::get_named_enum_kind_id("sret");
             let byref_id = Attribute::get_named_enum_kind_id("byref");
@@ -384,7 +385,7 @@ impl<'ctx, 'a> CodeGenContext<'ctx, 'a> {
                     if loc_params.is_empty() {
                         loc_params.extend(params[0..i+offset].iter().copied());
                     }
-                    let slot = self.builder.build_alloca(param.get_type(), call_name);
+                    let slot = gen_var(self, param.get_type(), Some(call_name)).unwrap();
                     loc_params.push(slot.into());
                     self.builder.build_store(slot, *param);
                 } else if !loc_params.is_empty() {
@@ -451,7 +452,7 @@ impl<'ctx, 'a> CodeGenContext<'ctx, 'a> {
     ) {
         let ty = self.get_llvm_type(generator, self.primitives.exception).into_pointer_type();
         let zelf_ty: BasicTypeEnum = ty.get_element_type().into_struct_type().into();
-        let zelf = self.builder.build_alloca(zelf_ty, "alloca");
+        let zelf = generator.gen_var_alloc(self, zelf_ty, Some("exn")).unwrap();
         let int32 = self.ctx.i32_type();
         let zero = int32.const_zero();
         unsafe {
diff --git a/nac3standalone/demo/demo.c b/nac3standalone/demo/demo.c
index 3e402b263..fb112d880 100644
--- a/nac3standalone/demo/demo.c
+++ b/nac3standalone/demo/demo.c
@@ -74,6 +74,12 @@ void output_str(struct cslice *slice) {
   putchar('\n');
 }
 
+uint64_t dbg_stack_address(__attribute__((unused)) struct cslice *slice) {
+  int i;
+  void *ptr = (void *) &i;
+  return (uintptr_t) ptr;
+}
+
 uint32_t __nac3_personality(uint32_t state, uint32_t exception_object, uint32_t context) {
   printf("__nac3_personality(state: %u, exception_object: %u, context: %u\n", state, exception_object, context);
   exit(101);
diff --git a/nac3standalone/demo/interpret_demo.py b/nac3standalone/demo/interpret_demo.py
index 0a1007f05..6eac7bb34 100755
--- a/nac3standalone/demo/interpret_demo.py
+++ b/nac3standalone/demo/interpret_demo.py
@@ -51,6 +51,9 @@ def patch(module):
     def output_float(x):
         print("%f" % x)
 
+    def dbg_stack_address(_):
+        return 0
+
     def extern(fun):
         name = fun.__name__
         if name == "output_asciiart":
@@ -67,6 +70,8 @@ def patch(module):
             "output_str",
         }:
             return print
+        elif name == "dbg_stack_address":
+            return dbg_stack_address
         else:
             raise NotImplementedError
 
diff --git a/nac3standalone/demo/src/stack_addr_issue233.py b/nac3standalone/demo/src/stack_addr_issue233.py
new file mode 100644
index 000000000..e71a22dbd
--- /dev/null
+++ b/nac3standalone/demo/src/stack_addr_issue233.py
@@ -0,0 +1,15 @@
+@extern
+def output_bool(x: bool):
+	...
+
+@extern
+def dbg_stack_address(x: str) -> uint64:
+	...
+
+def run() -> int32:
+	a = dbg_stack_address("a")
+	b = dbg_stack_address("b")
+
+	output_bool(a == b)
+
+	return 0
\ No newline at end of file