From ce0b7b0cbc758da920473d1f05a274d3c91ffde9 Mon Sep 17 00:00:00 2001
From: mlokr <jdoka@crownsterling.io>
Date: Thu, 18 Jul 2024 17:55:55 +0200
Subject: [PATCH] fixing wide register returns

---
 Cargo.toml                                        |  7 +++++++
 hblang/README.md                                  | 15 ++++++++++++++-
 hblang/src/codegen.rs                             | 11 +++++++++++
 ...n_tests_struct_return_from_module_function.txt |  3 +++
 4 files changed, 35 insertions(+), 1 deletion(-)
 create mode 100644 hblang/tests/codegen_tests_struct_return_from_module_function.txt

diff --git a/Cargo.toml b/Cargo.toml
index 04c812b1..43735306 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,3 +1,10 @@
 [workspace]
 resolver = "2"
 members = ["hbbytecode", "hbvm", "hbxrt", "xtask", "hblang", "hbjit"]
+
+[profile.small]
+inherits = "release"
+opt-level = "z"
+strip = true
+lto = true
+codegen-units = 1
diff --git a/hblang/README.md b/hblang/README.md
index f3acba4e..7aca49d6 100644
--- a/hblang/README.md
+++ b/hblang/README.md
@@ -448,7 +448,6 @@ main := fn(): int {
 
 #### different_types
 ```hb
-
 Color := struct {
 	r: u8,
 	g: u8,
@@ -494,3 +493,17 @@ main := fn(): int {
 }
 ```
 
+#### struct_return_from_module_function
+```hb
+bar := @use("bar.hb");
+
+main := fn(): int {
+	return 7 - bar.foo().x - bar.foo().y - bar.foo().z;
+}
+
+// in module: bar.hb
+
+foo := fn(): struct { x: int, y: u32, z: u32 } {
+	return .{ x: 3, y: 2, z: 2 };
+}
+```
diff --git a/hblang/src/codegen.rs b/hblang/src/codegen.rs
index 81cda8a7..383a3a24 100644
--- a/hblang/src/codegen.rs
+++ b/hblang/src/codegen.rs
@@ -2592,6 +2592,16 @@ impl Codegen {
             (&Loc::Ct { value }, lpat!(false, reg, 0, None)) => {
                 self.output.emit(li64(reg.get(), u64::from_ne_bytes(value)))
             }
+            (&Loc::Ct { value }, lpat!(false, reg, 8, None)) if reg.get() == 1 && size == 8 => {
+                self.output.emit(li64(reg.get() + 1, u64::from_ne_bytes(value)));
+            }
+            (&Loc::Ct { value }, lpat!(false, reg, off, None)) if reg.get() == 1 => {
+                let freg = reg.get() + (off / 8) as u8;
+                let mask = !(((1u64 << (8 * size)) - 1) << (8 * (off % 8)));
+                self.output.emit(andi(freg, freg, mask));
+                let value = u64::from_ne_bytes(value) << (8 * (off % 8));
+                self.output.emit(ori(freg, freg, value));
+            }
             (lpat!(true, src, soff, ref ssta), lpat!(true, dst, doff, ref dsta)) => {
                 // TODO: some oportuinies to ellit more optimal code
                 let src_off = self.ci.regs.allocate();
@@ -3131,5 +3141,6 @@ mod tests {
         c_strings => README;
         struct_patterns => README;
         arrays => README;
+        struct_return_from_module_function => README;
     }
 }
diff --git a/hblang/tests/codegen_tests_struct_return_from_module_function.txt b/hblang/tests/codegen_tests_struct_return_from_module_function.txt
new file mode 100644
index 00000000..de423263
--- /dev/null
+++ b/hblang/tests/codegen_tests_struct_return_from_module_function.txt
@@ -0,0 +1,3 @@
+code size: 330
+ret: 0
+status: Ok(())