From f9c47f86ad3667a1f1a21419b7217fdc19eb2161 Mon Sep 17 00:00:00 2001
From: Jakub Doka <jakub.doka2@gmail.com>
Date: Sun, 15 Dec 2024 17:42:20 +0100
Subject: [PATCH] fixing a glaring bug, where the elidded offset is also offset

Signed-off-by: Jakub Doka <jakub.doka2@gmail.com>
---
 lang/README.md                                | 20 +++++++++++++
 lang/src/backend/hbvm/regalloc.rs             | 22 +++++++++++----
 lang/src/son.rs                               |  1 +
 ...stack_offsets_for_parameters_correctly.txt | 28 +++++++++++++++++++
 4 files changed, 66 insertions(+), 5 deletions(-)
 create mode 100644 lang/tests/son_tests_elide_stack_offsets_for_parameters_correctly.txt

diff --git a/lang/README.md b/lang/README.md
index 3c3918ba..97e592e6 100644
--- a/lang/README.md
+++ b/lang/README.md
@@ -1673,6 +1673,26 @@ main := fn(): void {
 
 ### Just Testing Optimizations
 
+#### elide_stack_offsets_for_parameters_correctly
+```hb
+A := struct {
+	f: uint,
+	s: B,
+}
+
+B := struct {
+	a: uint,
+	b: uint,
+}
+
+main := fn(): uint {
+	a := A.(1, .(0, 0))
+	return pass(a.s)
+}
+
+pass := fn(s: B): uint return s.a + s.b
+```
+
 #### null_check_test
 ```hb
 get_ptr := fn(): ?^uint return null
diff --git a/lang/src/backend/hbvm/regalloc.rs b/lang/src/backend/hbvm/regalloc.rs
index 009e2b4c..290d9193 100644
--- a/lang/src/backend/hbvm/regalloc.rs
+++ b/lang/src/backend/hbvm/regalloc.rs
@@ -4,7 +4,7 @@ use {
             reg::{self, Reg},
             HbvmBackend, Nid, Nodes, PLoc, Reloc, TypedReloc,
         },
-        lexer::TokenKind,
+        lexer::{Token, TokenKind},
         parser, quad_sort,
         son::{Kind, ARG_START, MEM, VOID},
         ty::{self, Arg, Loc, Module, Offset, Sig, Types},
@@ -133,10 +133,22 @@ impl HbvmBackend {
                     res.node_to_reg[allc as usize]
                 };
 
-                let offset_atr = |allc: Nid, offsets: &[Offset]| {
-                    let allc = strip_load(allc);
-                    if nodes.is_locked(allc) && nodes[allc].kind == Kind::Stck {
-                        return (reg::STACK_PTR, offsets[allc as usize] as u64);
+                let offset_atr = |pallc: Nid, offsets: &[Offset]| {
+                    let allc = strip_load(pallc);
+                    if nodes.is_locked(allc) {
+                        let (region, offset) = nodes.strip_offset(allc, ty::Id::VOID, tys);
+                        match nodes[region].kind {
+                            Kind::Stck => {
+                                return (
+                                    reg::STACK_PTR,
+                                    offsets[region as usize] as u64 + offset as u64,
+                                )
+                            }
+                            _ => {
+                                assert_alloc_use(region);
+                                return (res.node_to_reg[region as usize], offset as u64);
+                            }
+                        }
                     }
 
                     assert_alloc_use(allc);
diff --git a/lang/src/son.rs b/lang/src/son.rs
index 0a8727e4..b0cd8a2b 100644
--- a/lang/src/son.rs
+++ b/lang/src/son.rs
@@ -6037,6 +6037,7 @@ mod tests {
         inline_return_stack;
 
         // Just Testing Optimizations;
+        elide_stack_offsets_for_parameters_correctly;
         const_folding_with_arg;
         branch_assignments;
         exhaustive_loop_testing;
diff --git a/lang/tests/son_tests_elide_stack_offsets_for_parameters_correctly.txt b/lang/tests/son_tests_elide_stack_offsets_for_parameters_correctly.txt
new file mode 100644
index 00000000..bf3337ff
--- /dev/null
+++ b/lang/tests/son_tests_elide_stack_offsets_for_parameters_correctly.txt
@@ -0,0 +1,28 @@
+main:
+    ADDI64 r254, r254, -40d
+    ST r31, r254, 24a, 16h
+    LI64 r32, 1d
+    ST r32, r254, 0a, 8h
+    ST r0, r254, 8a, 8h
+    ST r0, r254, 16a, 8h
+    LD r2, r254, 8a, 16h
+    JAL r31, r0, :pass
+    CP r32, r1
+    CP r1, r32
+    LD r31, r254, 24a, 16h
+    ADDI64 r254, r254, 40d
+    JALA r0, r31, 0a
+pass:
+    ADDI64 r254, r254, -16d
+    ST r2, r254, 0a, 16h
+    ADDI64 r2, r254, 0d
+    CP r13, r2
+    LD r14, r13, 0a, 8h
+    LD r13, r13, 8a, 8h
+    ADD64 r13, r13, r14
+    CP r1, r13
+    ADDI64 r254, r254, 16d
+    JALA r0, r31, 0a
+code size: 235
+ret: 0
+status: Ok(())