From d01e31b2034e53098f92dfb07b183a21857ded1e Mon Sep 17 00:00:00 2001
From: Jakub Doka <jakub.doka2@gmail.com>
Date: Wed, 13 Nov 2024 16:18:21 +0100
Subject: [PATCH] fixing stack return values

---
 lang/README.md                               | 13 ++++++++++
 lang/src/son.rs                              | 27 ++++++++++++++++----
 lang/tests/son_tests_inline_return_stack.txt |  6 +++++
 3 files changed, 41 insertions(+), 5 deletions(-)
 create mode 100644 lang/tests/son_tests_inline_return_stack.txt

diff --git a/lang/README.md b/lang/README.md
index 88b61841..1080c28c 100644
--- a/lang/README.md
+++ b/lang/README.md
@@ -283,6 +283,19 @@ main := fn(): void {
 }
 ```
 
+#### inline_return_stack
+```hb
+$fun := fn(): [uint; 3] {
+	res := [uint].(0, 1, 2)
+	return res
+}
+
+main := fn(): uint {
+	vl := fun()
+	return vl[0]
+}
+```
+
 #### struct_operators
 ```hb
 Point := struct {
diff --git a/lang/src/son.rs b/lang/src/son.rs
index 4e97cfbf..d6391d39 100644
--- a/lang/src/son.rs
+++ b/lang/src/son.rs
@@ -2158,7 +2158,7 @@ pub struct ItemCtx {
     inline_var_base: usize,
     inline_aclass_base: usize,
     inline_depth: u16,
-    inline_ret: Option<(Value, StrongRef, Scope)>,
+    inline_ret: Option<(Value, StrongRef, Scope, Option<AClass>)>,
     nodes: Nodes,
     ctrl: StrongRef,
     loops: Vec<Loop>,
@@ -2683,7 +2683,11 @@ impl<'a> Codegen<'a> {
                     self.ci.ctrl.set(NEVER, &mut self.ci.nodes);
                     self.ci.nodes[ret].pos = pos;
                     self.ci.nodes.bind(ret, NEVER);
-                } else if let Some((pv, ctrl, scope)) = &mut self.ci.inline_ret {
+                } else if let Some((pv, ctrl, scope, aclass)) = &mut self.ci.inline_ret {
+                    debug_assert!(
+                        aclass.is_none(),
+                        "TODO: oh no, we cant return structs from divergent branches"
+                    );
                     ctrl.set(
                         self.ci.nodes.new_node(
                             ty::Id::VOID,
@@ -2726,8 +2730,14 @@ impl<'a> Codegen<'a> {
                         .for_each(|v| v.remove(&mut self.ci.nodes));
 
                     let repl = StrongRef::new(NEVER, &mut self.ci.nodes);
+                    let (index, _) = self
+                        .ci
+                        .nodes
+                        .aclass_index(*self.ci.nodes[value.id].inputs.get(1).unwrap_or(&VOID));
+                    let aclass = (self.ci.inline_aclass_base <= index)
+                        .then(|| self.ci.scope.aclasses[index].dup(&mut self.ci.nodes));
                     self.ci.inline_ret =
-                        Some((value, mem::replace(&mut self.ci.ctrl, repl), scope));
+                        Some((value, mem::replace(&mut self.ci.ctrl, repl), scope, aclass));
                 }
 
                 None
@@ -3871,7 +3881,8 @@ impl<'a> Codegen<'a> {
                 var.remove(&mut self.ci.nodes);
             }
 
-            let (v, ctrl, mut scope) = mem::replace(&mut self.ci.inline_ret, prev_inline_ret)?;
+            let (v, ctrl, mut scope, aclass) =
+                mem::replace(&mut self.ci.inline_ret, prev_inline_ret)?;
             if is_inline
                 && ctrl.get() != prev_ctrl
                 && (!self.ci.nodes[ctrl.get()].kind.is_eca()
@@ -3882,11 +3893,16 @@ impl<'a> Codegen<'a> {
 
             scope.vars.drain(var_base..).for_each(|v| v.remove(&mut self.ci.nodes));
             scope.aclasses.drain(aclass_base..).for_each(|v| v.remove(&mut self.ci.nodes));
-            scope.aclasses.extend(self.ci.scope.aclasses.drain(aclass_base..));
             self.ci.nodes.unlock(v.id);
             self.ci.scope.clear(&mut self.ci.nodes);
             self.ci.scope = scope;
 
+            if let Some(aclass) = aclass {
+                let (_, reg) = self.ci.nodes.aclass_index(v.id);
+                self.ci.nodes[reg].aclass = self.ci.scope.aclasses.len() as _;
+                self.ci.scope.aclasses.push(aclass);
+            }
+
             mem::replace(&mut self.ci.ctrl, ctrl).remove(&mut self.ci.nodes);
 
             Some(v)
@@ -4868,6 +4884,7 @@ mod tests {
         request_page;
         tests_ptr_to_ptr_copy;
         global_variable_wiredness;
+        inline_return_stack;
 
         // Just Testing Optimizations;
         const_folding_with_arg;
diff --git a/lang/tests/son_tests_inline_return_stack.txt b/lang/tests/son_tests_inline_return_stack.txt
new file mode 100644
index 00000000..5aa16aa1
--- /dev/null
+++ b/lang/tests/son_tests_inline_return_stack.txt
@@ -0,0 +1,6 @@
+main:
+    CP r1, r0
+    JALA r0, r31, 0a
+code size: 22
+ret: 0
+status: Ok(())