Try to improve const `union` LLVM construction

4 files changed, 172 insertions, 17 deletions

diff --git a/src/llvm_backend.cpp b/src/llvm_backend.cpp
index 33b03e235..3af473c3a 100644
--- a/src/llvm_backend.cpp
+++ b/src/llvm_backend.cpp
@@ -1914,7 +1914,7 @@ gb_internal WORKER_TASK_PROC(lb_llvm_module_verification_worker_proc) {
 	lbModule *m = cast(lbModule *)data;
 
 	if (LLVMVerifyModule(m->mod, LLVMReturnStatusAction, &llvm_error)) {
-		gb_printf_err("LLVM Error:\n%s\n", llvm_error);
+		gb_printf_err("LLVM Error in module %s:\n%s\n", m->module_name, llvm_error);
 		if (build_context.keep_temp_files) {
 			TIME_SECTION("LLVM Print Module to File");
 			String filepath_ll = lb_filepath_ll_for_module(m);
diff --git a/src/llvm_backend_const.cpp b/src/llvm_backend_const.cpp
index 98e50dd78..781cc52b8 100644
--- a/src/llvm_backend_const.cpp
+++ b/src/llvm_backend_const.cpp
@@ -537,6 +537,100 @@ gb_internal bool lb_is_nested_possibly_constant(Type *ft, Selection const &sel,
 	return lb_is_elem_const(elem, ft);
 }
 
+gb_internal LLVMValueRef lb_construct_const_union(lbModule *m, LLVMValueRef variant_value, Type *variant_type, Type *union_type) {
+	Type *bt = base_type(union_type);
+	GB_ASSERT(bt->kind == Type_Union);
+	GB_ASSERT(lb_type(m, variant_type) == LLVMTypeOf(variant_value));
+
+	if (bt->Union.variants.count == 0) {
+		GB_ASSERT(LLVMIsNull(variant_value));
+		return variant_value;
+	}
+
+	LLVMTypeRef llvm_type = lb_type(m, union_type);
+	LLVMTypeRef llvm_variant_type = lb_type(m, variant_type);
+
+	if (is_type_union_maybe_pointer(bt)) {
+		GB_ASSERT(lb_sizeof(LLVMTypeOf(variant_value)) == lb_sizeof(llvm_type));
+		return LLVMConstBitCast(variant_value, llvm_type);
+	}
+
+	if (bt->Union.variants.count == 1) {
+		unsigned long long the_tag = cast(unsigned long long)union_variant_index(union_type, variant_type);
+		LLVMTypeRef tag_type = lb_type(m, union_tag_type(bt));
+
+		LLVMValueRef values[3] = {};
+		unsigned i = 0;
+		values[i++] = variant_value;
+		values[i++] = LLVMConstInt(tag_type, the_tag, false);
+
+		i64 used_size = bt->Union.variant_block_size + lb_sizeof(tag_type);
+		i64 padding = type_size_of(union_type) - used_size;
+		i64 align = type_align_of(union_type);
+		if (padding > 0) {
+			LLVMTypeRef padding_type = lb_type_padding_filler(m, padding, align);
+			values[i++] = LLVMConstNull(padding_type);
+		}
+
+		return LLVMConstNamedStruct(llvm_type, values, i);
+	}
+
+	LLVMTypeRef block_type = LLVMStructGetTypeAtIndex(llvm_type, 0);
+
+	LLVMTypeRef block_padding = nullptr;
+	i64 block_padding_size = bt->Union.variant_block_size - type_size_of(variant_type);
+	if (block_padding_size > 0) {
+		block_padding = lb_type_padding_filler(m, block_padding_size, type_align_of(variant_type));
+		return nullptr;
+	}
+
+	LLVMTypeRef tag_type = lb_type(m, union_tag_type(bt));
+
+	i64 used_size = bt->Union.variant_block_size + lb_sizeof(tag_type);
+	i64 padding = type_size_of(union_type) - used_size;
+	i64 align = type_align_of(union_type);
+	LLVMTypeRef padding_type = nullptr;
+	if (padding > 0) {
+		padding_type = lb_type_padding_filler(m, padding, align);
+	}
+
+
+	unsigned i = 0;
+	LLVMValueRef values[3] = {};
+
+	LLVMValueRef variant_value_wrapped = variant_value;
+
+	if (lb_sizeof(llvm_variant_type) == 0) {
+		variant_value_wrapped = LLVMConstNull(block_type);
+	} else if (block_type != llvm_variant_type) {
+		return nullptr;
+	}
+
+	values[i++] = variant_value_wrapped;
+
+	if (block_padding_size > 0) {
+		values[i++] = LLVMConstNull(block_padding);
+	}
+	unsigned long long the_tag = cast(unsigned long long)union_variant_index(union_type, variant_type);
+	values[i++] = LLVMConstInt(tag_type, the_tag, false);
+	if (padding > 0) {
+		values[i++] = LLVMConstNull(padding_type);
+	}
+	return LLVMConstNamedStruct(llvm_type, values, i);
+}
+
+gb_internal bool lb_try_construct_const_union(lbModule *m, lbValue *value, Type *variant_type, Type *union_type) {
+	if (lb_is_const(*value)) {
+		LLVMValueRef res = lb_construct_const_union(m, value->value, variant_type, union_type);
+		if (res != nullptr) {
+			*value = {res, union_type};
+			return true;
+		}
+	}
+	return false;
+}
+
+
 gb_internal lbValue lb_const_value(lbModule *m, Type *type, ExactValue value, lbConstContext cc) {
 	if (cc.allow_local) {
 		cc.is_rodata = false;
diff --git a/src/llvm_backend_expr.cpp b/src/llvm_backend_expr.cpp
index cff91813e..4a1f39c19 100644
--- a/src/llvm_backend_expr.cpp
+++ b/src/llvm_backend_expr.cpp
@@ -2495,6 +2495,11 @@ gb_internal lbValue lb_emit_conv(lbProcedure *p, lbValue value, Type *t) {
 			Type *vt = dst->Union.variants[0];
 			if (internal_check_is_assignable_to(src_type, vt)) {
 				value = lb_emit_conv(p, value, vt);
+				if (lb_is_const(value)) {
+					LLVMValueRef res = lb_construct_const_union(m, value.value, vt, t);
+					return {res, t};
+				}
+
 				lbAddr parent = lb_add_local_generated(p, t, true);
 				lb_emit_store_union_variant(p, parent.addr, value, vt);
 				return lb_addr_load(p, parent);
@@ -2503,11 +2508,18 @@ gb_internal lbValue lb_emit_conv(lbProcedure *p, lbValue value, Type *t) {
 		for (Type *vt : dst->Union.variants) {
 			if (src_type == t_llvm_bool && is_type_boolean(vt)) {
 				value = lb_emit_conv(p, value, vt);
+				if (lb_try_construct_const_union(m, &value, vt, t)) {
+					return value;
+				}
+
 				lbAddr parent = lb_add_local_generated(p, t, true);
 				lb_emit_store_union_variant(p, parent.addr, value, vt);
 				return lb_addr_load(p, parent);
 			}
 			if (are_types_identical(src_type, vt)) {
+				if (lb_try_construct_const_union(m, &value, vt, t)) {
+					return value;
+				}
 				lbAddr parent = lb_add_local_generated(p, t, true);
 				lb_emit_store_union_variant(p, parent.addr, value, vt);
 				return lb_addr_load(p, parent);
@@ -2545,6 +2557,9 @@ gb_internal lbValue lb_emit_conv(lbProcedure *p, lbValue value, Type *t) {
 		if (valid_count == 1) {
 			Type *vt = dst->Union.variants[first_success_index];
 			value = lb_emit_conv(p, value, vt);
+			if (lb_try_construct_const_union(m, &value, vt, t)) {
+					return value;
+				}
 			lbAddr parent = lb_add_local_generated(p, t, true);
 			lb_emit_store_union_variant(p, parent.addr, value, vt);
 			return lb_addr_load(p, parent);
diff --git a/src/llvm_backend_general.cpp b/src/llvm_backend_general.cpp
index 3fa7a076e..fc770102c 100644
--- a/src/llvm_backend_general.cpp
+++ b/src/llvm_backend_general.cpp
@@ -1719,8 +1719,69 @@ gb_internal LLVMTypeRef lb_type_internal_for_procedures_raw(lbModule *m, Type *t
 	map_set(&m->func_raw_types, type, new_abi_fn_type);
 
 	return new_abi_fn_type;
+}
+
+
+gb_internal LLVMTypeRef lb_type_internal_union_block_type(lbModule *m, Type *type) {
+	GB_ASSERT(type->kind == Type_Union);
+
+	if (type->Union.variants.count <= 0) {
+		return nullptr;
+	}
+	if (type->Union.variants.count == 1) {
+		return lb_type(m, type->Union.variants[0]);
+	}
+
+	i64 align = type_align_of(type);
+	i64 size = type_size_of(type);
+	gb_unused(size);
 
+	unsigned block_size = cast(unsigned)type->Union.variant_block_size;
+
+	bool all_pointers = align == build_context.ptr_size;
+	for (isize i = 0; all_pointers && i < type->Union.variants.count; i++) {
+		Type *t = type->Union.variants[i];
+		if (!is_type_internally_pointer_like(t)) {
+			all_pointers = false;
+		}
+	}
+	if (all_pointers) {
+		return lb_type(m, t_rawptr);
+	}
+
+	{
+		Type *pt = type->Union.variants[0];
+		for (isize i = 1; i < type->Union.variants.count; i++) {
+			Type *t = type->Union.variants[i];
+			if (!are_types_identical(pt, t)) {
+				goto end_check_for_all_the_same;
+			}
+		}
+		return lb_type(m, pt);
+	} end_check_for_all_the_same:;
+
+	{
+		Type *first_different = nullptr;
+		for (isize i = 0; i < type->Union.variants.count; i++) {
+			Type *t = type->Union.variants[i];
+			if (type_size_of(t) == 0) {
+				continue;
+			}
+			if (first_different == nullptr) {
+				first_different = t;
+			} else if (!are_types_identical(first_different, t)) {
+				goto end_rest_zero_except_one;
+			}
+		}
+		if (first_different != nullptr) {
+			return lb_type(m, first_different);
+		}
+	} end_rest_zero_except_one:;
+
+	return lb_type_padding_filler(m, block_size, align);
 }
+
+
 gb_internal LLVMTypeRef lb_type_internal(lbModule *m, Type *type) {
 	LLVMContextRef ctx = m->ctx;
 	i64 size = type_size_of(type); // Check size
@@ -2233,8 +2294,6 @@ gb_internal LLVMTypeRef lb_type_internal(lbModule *m, Type *type) {
 				return LLVMStructTypeInContext(ctx, fields, gb_count_of(fields), false);
 			}
 
-			unsigned block_size = cast(unsigned)type->Union.variant_block_size;
-
 			auto fields = array_make<LLVMTypeRef>(temporary_allocator(), 0, 3);
 			if (is_type_union_maybe_pointer(type)) {
 				LLVMTypeRef variant = lb_type(m, type->Union.variants[0]);
@@ -2252,20 +2311,7 @@ gb_internal LLVMTypeRef lb_type_internal(lbModule *m, Type *type) {
 					array_add(&fields, padding_type);
 				}
 			} else {
-				LLVMTypeRef block_type = nullptr;
-
-				bool all_pointers = align == build_context.ptr_size;
-				for (isize i = 0; all_pointers && i < type->Union.variants.count; i++) {
-					Type *t = type->Union.variants[i];
-					if (!is_type_internally_pointer_like(t)) {
-						all_pointers = false;
-					}
-				}
-				if (all_pointers) {
-					block_type = lb_type(m, t_rawptr);
-				} else {
-					block_type = lb_type_padding_filler(m, block_size, align);
-				}
+				LLVMTypeRef block_type = lb_type_internal_union_block_type(m, type);
 
 				LLVMTypeRef tag_type = lb_type(m, union_tag_type(type));
 				array_add(&fields, block_type);