From 6f9b466d7cea5eea13e8e261b64b3798bf794e94 Mon Sep 17 00:00:00 2001
From: Christoffer Lerno <christoffer@aegik.com>
Date: Sat, 4 Jan 2025 21:57:58 +0100
Subject: [PATCH] Optimize recast.

---
 src/compiler/compiler_internal.h            | 35 ++++++++++++++++-----
 src/compiler/llvm_codegen_expr.c            | 12 +++++++
 test/test_suite/cast/cast_bitstruct_etc.c3t | 13 +++-----
 3 files changed, 44 insertions(+), 16 deletions(-)

diff --git a/src/compiler/compiler_internal.h b/src/compiler/compiler_internal.h
index 3a36cf51e..7556c3d00 100644
--- a/src/compiler/compiler_internal.h
+++ b/src/compiler/compiler_internal.h
@@ -3593,15 +3593,10 @@ INLINE Ast *ast_next(AstId *current_ptr)
 
 INLINE void expr_rewrite_recast(Expr *expr, Type *type)
 {
-	if (expr->expr_kind == EXPR_RECAST)
+	if (expr->expr_kind == EXPR_RECAST || expr->expr_kind == EXPR_ADDR_CONVERSION || expr->expr_kind == EXPR_RVALUE)
 	{
-		Expr *inner = expr->inner_expr;
-		if (type_flatten(inner->type) == type_flatten(type))
-		{
-			expr_replace(expr, inner);
-			expr->type = type;
-			return;
-		}
+		expr->type = type;
+		return;
 	}
 
 	Expr *inner = expr_copy(expr);
@@ -3612,6 +3607,18 @@ INLINE void expr_rewrite_recast(Expr *expr, Type *type)
 
 INLINE void expr_rewrite_rvalue(Expr *expr, Type *type)
 {
+	switch (expr->expr_kind)
+	{
+		case EXPR_RECAST:
+			expr->expr_kind = EXPR_RVALUE;
+			expr->type = type;
+			return;
+		case EXPR_RVALUE:
+			expr->type = type;
+			return;
+		default:
+			break;
+	}
 	Expr *inner = expr_copy(expr);
 	expr->expr_kind = EXPR_RVALUE;
 	expr->inner_expr = inner;
@@ -3620,6 +3627,18 @@ INLINE void expr_rewrite_rvalue(Expr *expr, Type *type)
 
 INLINE void expr_rewrite_addr_conversion(Expr *expr, Type *type)
 {
+	switch (expr->expr_kind)
+	{
+		case EXPR_RECAST:
+			expr->expr_kind = EXPR_ADDR_CONVERSION;
+			expr->type = type;
+			return;
+		case EXPR_ADDR_CONVERSION:
+			expr->type = type;
+			return;
+		default:
+			break;
+	}
 	Expr *inner = expr_copy(expr);
 	expr->expr_kind = EXPR_ADDR_CONVERSION;
 	expr->inner_expr = inner;
diff --git a/src/compiler/llvm_codegen_expr.c b/src/compiler/llvm_codegen_expr.c
index 00b218d35..74158b46a 100644
--- a/src/compiler/llvm_codegen_expr.c
+++ b/src/compiler/llvm_codegen_expr.c
@@ -94,12 +94,24 @@ LLVMValueRef llvm_emit_expect_raw(GenContext *c, LLVMValueRef expect_true)
 	return llvm_emit_call_intrinsic(c, intrinsic_id.expect, &c->bool_type, 1, values, 2);
 }
 
+Expr *expr_remove_recast(Expr *expr)
+{
+	Type *main_type = type_lowering(expr->type);
+	while (expr->expr_kind == EXPR_RECAST && main_type == type_lowering(expr->inner_expr->type))
+	{
+		expr = expr->inner_expr;
+	}
+	return expr;
+}
+
 BEValue llvm_emit_assign_expr(GenContext *c, BEValue *ref, Expr *expr, LLVMValueRef optional, bool is_init)
 {
 	ASSERT0(llvm_value_is_addr(ref));
 
 	assert((optional || !IS_OPTIONAL(expr)) && "Assumed an optional address if it's an optional expression.");
 
+	expr = expr_remove_recast(expr);
+
 	// Special optimization of handling of optional
 	if (expr->expr_kind == EXPR_OPTIONAL)
 	{
diff --git a/test/test_suite/cast/cast_bitstruct_etc.c3t b/test/test_suite/cast/cast_bitstruct_etc.c3t
index bb47fea83..e34922498 100644
--- a/test/test_suite/cast/cast_bitstruct_etc.c3t
+++ b/test/test_suite/cast/cast_bitstruct_etc.c3t
@@ -22,18 +22,15 @@ bitstruct Foo : ushort
 
   %e = alloca i8, align 1
   %x = alloca i16, align 2
-  %literal = alloca i16, align 2
   %z = alloca i16, align 2
   store i8 0, ptr %e, align 1
   %0 = load i8, ptr %e, align 1
   %zext = zext i8 %0 to i16
   %1 = and i16 %zext, 255
-  store i16 %1, ptr %literal, align 2
-  %2 = load i16, ptr %literal, align 2
-  store i16 %2, ptr %x, align 2
-  %3 = load i8, ptr %e, align 1
-  %zext1 = zext i8 %3 to i16
-  %4 = and i16 %zext1, 255
-  store i16 %4, ptr %z, align 2
+  store i16 %1, ptr %x, align 2
+  %2 = load i8, ptr %e, align 1
+  %zext1 = zext i8 %2 to i16
+  %3 = and i16 %zext1, 255
+  store i16 %3, ptr %z, align 2
   ret i64 0
 }
\ No newline at end of file