From 5d73e9231c6251cfe91b179976625a1ecfd2c4e4 Mon Sep 17 00:00:00 2001
From: Brennen Shaughnessy <pinsaasin@gmail.com>
Date: Sat, 14 Dec 2024 14:44:32 -0500
Subject: [PATCH] Port homogenous type aggregation from arm64 to arm32

---
 src/llvm_abi.cpp | 91 ++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 91 insertions(+)

diff --git a/src/llvm_abi.cpp b/src/llvm_abi.cpp
index 0b2bb7956a4..6f611bddb54 100644
--- a/src/llvm_abi.cpp
+++ b/src/llvm_abi.cpp
@@ -1420,6 +1420,7 @@ namespace lbAbiWasm {
 
 namespace lbAbiArm32 {
 	gb_internal Array<lbArgType> compute_arg_types(LLVMContextRef c, LLVMTypeRef *arg_types, unsigned arg_count, ProcCallingConvention calling_convention);
+	gb_internal bool is_homogenous_aggregate(LLVMContextRef c, LLVMTypeRef type, LLVMTypeRef *base_type_, unsigned *member_count_);
 	gb_internal lbArgType compute_return_type(LLVMContextRef c, LLVMTypeRef return_type, bool return_is_defined);
 
 	gb_internal LB_ABI_INFO(abi_info) {
@@ -1459,14 +1460,104 @@ namespace lbAbiArm32 {
 		}
 		return lb_arg_type_direct(type, nullptr, nullptr, attr);
 	}
+	
+	gb_internal bool is_homogenous_array(LLVMContextRef c, LLVMTypeRef type, LLVMTypeRef *base_type_, unsigned *member_count_) {
+		GB_ASSERT(lb_is_type_kind(type, LLVMArrayTypeKind));
+		unsigned len = LLVMGetArrayLength(type);
+		if (len == 0) {
+			return false;
+		}
+		LLVMTypeRef elem = OdinLLVMGetArrayElementType(type);
+		LLVMTypeRef base_type = nullptr;
+		unsigned member_count = 0;
+		if (is_homogenous_aggregate(c, elem, &base_type, &member_count)) {
+			if (base_type_) *base_type_ = base_type;
+			if (member_count_) *member_count_ = member_count * len;
+			return true;
+
+		}
+		return false;
+	}
+	gb_internal bool is_homogenous_struct(LLVMContextRef c, LLVMTypeRef type, LLVMTypeRef *base_type_, unsigned *member_count_) {
+		GB_ASSERT(lb_is_type_kind(type, LLVMStructTypeKind));
+		unsigned elem_count = LLVMCountStructElementTypes(type);
+		if (elem_count == 0) {
+			return false;
+		}
+		LLVMTypeRef base_type = nullptr;
+		unsigned member_count = 0;
+
+		for (unsigned i = 0; i < elem_count; i++) {
+			LLVMTypeRef field_type = nullptr;
+			unsigned field_member_count = 0;
+
+			LLVMTypeRef elem = LLVMStructGetTypeAtIndex(type, i);
+			if (!is_homogenous_aggregate(c, elem, &field_type, &field_member_count)) {
+				return false;
+			}
+
+			if (base_type == nullptr) {
+				base_type = field_type;
+				member_count = field_member_count;
+			} else {
+				if (base_type != field_type) {
+					return false;
+				}
+				member_count += field_member_count;
+			}
+		}
+
+		if (base_type == nullptr) {
+			return false;
+		}
+
+		if (lb_sizeof(type) == lb_sizeof(base_type) * member_count) {
+			if (base_type_) *base_type_ = base_type;
+			if (member_count_) *member_count_ = member_count;
+			return true;
+		}
+
+		return false;
+	}
+
+
+	gb_internal bool is_homogenous_aggregate(LLVMContextRef c, LLVMTypeRef type, LLVMTypeRef *base_type_, unsigned *member_count_) {
+		LLVMTypeKind kind = LLVMGetTypeKind(type);
+		switch (kind) {
+		case LLVMFloatTypeKind:
+		case LLVMDoubleTypeKind:
+			if (base_type_) *base_type_ = type;
+			if (member_count_) *member_count_ = 1;
+			return true;
+		case LLVMArrayTypeKind:
+			return is_homogenous_array(c, type, base_type_, member_count_);
+		case LLVMStructTypeKind:
+			return is_homogenous_struct(c, type, base_type_, member_count_);
+		}
+		return false;
+	}
+
+	gb_internal unsigned is_homogenous_aggregate_small_enough(LLVMTypeRef base_type, unsigned member_count) {
+		return (member_count <= 4);
+	}
 
 	gb_internal Array<lbArgType> compute_arg_types(LLVMContextRef c, LLVMTypeRef *arg_types, unsigned arg_count, ProcCallingConvention calling_convention) {
 		auto args = array_make<lbArgType>(lb_function_type_args_allocator(), arg_count);
 
 		for (unsigned i = 0; i < arg_count; i++) {
 			LLVMTypeRef t = arg_types[i];
+			
+			LLVMTypeRef homo_base_type = {};
+			unsigned homo_member_count = 0;
+			
 			if (is_register(t, false)) {
 				args[i] = non_struct(c, t, false);
+			} else if (is_homogenous_aggregate(c, t, &homo_base_type, &homo_member_count)) {
+				if (is_homogenous_aggregate_small_enough(homo_base_type, homo_member_count)) {
+					args[i] = lb_arg_type_direct(t, llvm_array_type(homo_base_type, homo_member_count), nullptr, nullptr);
+				} else {
+					args[i] = lb_arg_type_indirect(t, nullptr);;
+				}
 			} else {
 				i64 sz = lb_sizeof(t);
 				i64 a = lb_alignof(t);