From 9c43dafcf5cb44e30d88bff257bb2fb1327dc9c2 Mon Sep 17 00:00:00 2001
From: Jakub Doka <jakub.doka2@gmail.com>
Date: Thu, 24 Oct 2024 14:08:17 +0200
Subject: [PATCH] fixing @as misbehaving

---
 lang/README.md                                |  8 ++---
 lang/src/son.rs                               |  9 +++++-
 .../son_tests_returning_global_struct.txt     | 21 +++++++------
 lang/tests/son_tests_small_struct_bitcast.txt | 30 ++++++++++++-------
 4 files changed, 41 insertions(+), 27 deletions(-)

diff --git a/lang/README.md b/lang/README.md
index 2fe1e99..d0f667b 100644
--- a/lang/README.md
+++ b/lang/README.md
@@ -542,14 +542,12 @@ main := fn(): int {
 
 #### small_struct_bitcast
 ```hb
-
 Color := struct {r: u8, g: u8, b: u8, a: u8}
 white := Color.(255, 255, 255, 255)
-u32_to_color := fn(v: u32): Color {
-	return @bitcast(v)
-}
+u32_to_color := fn(v: u32): Color return @bitcast(v)
+u32_to_u32 := fn(v: u32): u32 return v
 main := fn(): int {
-	return u32_to_color(@bitcast(white)).r
+	return u32_to_color(@bitcast(white)).r + @as(Color, @bitcast(u32_to_u32(@bitcast(white)))).g
 }
 ```
 
diff --git a/lang/src/son.rs b/lang/src/son.rs
index 0aa3c44..8fd1179 100644
--- a/lang/src/son.rs
+++ b/lang/src/son.rs
@@ -2539,7 +2539,8 @@ impl<'a> Codegen<'a> {
             Expr::Directive { name: "as", args: [ty, expr], .. } => {
                 let ty = self.ty(ty);
                 let ctx = Ctx::default().with_ty(ty);
-                let mut val = self.expr_ctx(expr, ctx)?;
+                let mut val = self.raw_expr_ctx(expr, ctx)?;
+                self.strip_var(&mut val);
                 self.assert_ty(expr.pos(), &mut val, ty, "hinted expr");
                 Some(val)
             }
@@ -3538,6 +3539,9 @@ impl<'a> Codegen<'a> {
 
             if let Some(oper) = to_correct {
                 oper.ty = upcasted;
+                if core::mem::take(&mut oper.ptr) {
+                    oper.id = self.load_mem(oper.id, oper.ty);
+                }
                 oper.id = self.ci.nodes.new_node(upcasted, Kind::Extend, [VOID, oper.id]);
                 if matches!(op, TokenKind::Add | TokenKind::Sub)
                     && let Some(elem) = self.tys.base_of(upcasted)
@@ -3586,6 +3590,9 @@ impl<'a> Codegen<'a> {
                     self.ty_display(upcasted)
                 );
                 src.ty = upcasted;
+                if core::mem::take(&mut src.ptr) {
+                    src.id = self.load_mem(src.id, src.ty);
+                }
                 src.id = self.ci.nodes.new_node(upcasted, Kind::Extend, [VOID, src.id]);
             }
             true
diff --git a/lang/tests/son_tests_returning_global_struct.txt b/lang/tests/son_tests_returning_global_struct.txt
index 21d52e8..b6b214c 100644
--- a/lang/tests/son_tests_returning_global_struct.txt
+++ b/lang/tests/son_tests_returning_global_struct.txt
@@ -6,17 +6,16 @@ main:
     ST r1, r254, 4a, 4h
     ADDI64 r5, r254, 0d
     BMC r32, r5, 4h
-    LD r4, r254, 3a, 1h
-    ANDI r6, r4, 255d
-    LD r3, r254, 2a, 1h
+    LD r3, r254, 3a, 1h
     ANDI r5, r3, 255d
-    LD r2, r254, 1a, 1h
+    LD r2, r254, 2a, 1h
     ANDI r4, r2, 255d
-    LD r2, r254, 0a, 1h
-    ANDI r7, r2, 255d
-    ADD64 r8, r4, r7
-    ADD64 r12, r8, r5
-    ADD64 r1, r12, r6
+    LD r1, r254, 1a, 1h
+    ANDI r3, r1, 255d
+    LD r2, r254, 0a, 8h
+    ADD64 r6, r2, r3
+    ADD64 r10, r6, r4
+    ADD64 r1, r10, r5
     LD r31, r254, 8a, 16h
     ADDI64 r254, r254, 24d
     JALA r0, r31, 0a
@@ -24,6 +23,6 @@ random_color:
     LRA r1, r0, :white
     LD r1, r1, 0a, 4h
     JALA r0, r31, 0a
-code size: 257
-ret: 1020
+code size: 246
+ret: 764
 status: Ok(())
diff --git a/lang/tests/son_tests_small_struct_bitcast.txt b/lang/tests/son_tests_small_struct_bitcast.txt
index 9d7251d..942c78a 100644
--- a/lang/tests/son_tests_small_struct_bitcast.txt
+++ b/lang/tests/son_tests_small_struct_bitcast.txt
@@ -1,15 +1,22 @@
 main:
-    ADDI64 r254, r254, -12d
-    ST r31, r254, 4a, 8h
+    ADDI64 r254, r254, -24d
+    ST r31, r254, 8a, 16h
     LRA r1, r0, :white
-    LD r2, r1, 0a, 4h
-    ADDI64 r5, r254, 0d
+    LD r32, r1, 0a, 4h
+    ADDI64 r5, r254, 4d
+    CP r2, r32
     JAL r31, r0, :u32_to_color
+    ST r1, r254, 4a, 4h
+    CP r2, r32
+    JAL r31, r0, :u32_to_u32
+    ADDI64 r12, r254, 0d
+    LD r11, r254, 4a, 1h
     ST r1, r254, 0a, 4h
-    LD r9, r254, 0a, 1h
-    ANDI r1, r9, 255d
-    LD r31, r254, 4a, 8h
-    ADDI64 r254, r254, 12d
+    LD r4, r254, 1a, 1h
+    ADD8 r7, r4, r11
+    ANDI r1, r7, 255d
+    LD r31, r254, 8a, 16h
+    ADDI64 r254, r254, 24d
     JALA r0, r31, 0a
 u32_to_color:
     ADDI64 r254, r254, -4d
@@ -18,6 +25,9 @@ u32_to_color:
     LD r1, r3, 0a, 4h
     ADDI64 r254, r254, 4d
     JALA r0, r31, 0a
-code size: 216
-ret: 255
+u32_to_u32:
+    CP r1, r2
+    JALA r0, r31, 0a
+code size: 284
+ret: 254
 status: Ok(())