Fix calculation of vector type sizes

4 files changed, 22 insertions, 20 deletions

diff --git a/src/check_expr.c b/src/check_expr.c
index a2267f96c..56899bff4 100644
--- a/src/check_expr.c
+++ b/src/check_expr.c
@@ -5380,7 +5380,7 @@ ExprKind check_expr_base_internal(Checker *c, Operand *o, AstNode *node, Type *t
 						fields_visited[sel.index.e[0]] = true;
 						check_expr(c, o, fv->value);
 
-						if (base_type(field->type) == t_any) {
+						if (is_type_any(field->type) || is_type_union(field->type) || is_type_raw_union(field->type)) {
 							is_constant = false;
 						}
 						if (is_constant) {
@@ -5418,7 +5418,7 @@ ExprKind check_expr_base_internal(Checker *c, Operand *o, AstNode *node, Type *t
 							continue;
 						}
 
-						if (base_type(field->type) == t_any) {
+						if (is_type_any(field->type) || is_type_union(field->type) || is_type_raw_union(field->type)) {
 							is_constant = false;
 						}
 						if (is_constant) {
diff --git a/src/ir.c b/src/ir.c
index 956877ff8..bb9d8b7e8 100644
--- a/src/ir.c
+++ b/src/ir.c
@@ -5203,7 +5203,9 @@ irAddr ir_build_addr(irProcedure *proc, AstNode *expr) {
 					}
 
 					field = st->fields[index];
-					if (!is_union && ir_is_elem_const(proc->module, elem, field->type)) {
+					Type *ft = field->type;
+					if (!is_union && !is_type_union(ft) &&
+					    ir_is_elem_const(proc->module, elem, ft)) {
 						continue;
 					}
 
@@ -5211,7 +5213,9 @@ irAddr ir_build_addr(irProcedure *proc, AstNode *expr) {
 
 					GB_ASSERT(ir_type(field_expr)->kind != Type_Tuple);
 
-					Type *ft = field->type;
+					if (is_type_union(ft)) {
+						// gb_printf_err("HERE! %s\n", type_to_string(ft));
+					}
 					irValue *fv = ir_emit_conv(proc, field_expr, ft);
 					irValue *gep = ir_emit_struct_ep(proc, v, index);
 					ir_emit_store(proc, gep, fv);
diff --git a/src/ir_print.c b/src/ir_print.c
index 8a069c856..92833eaf4 100644
--- a/src/ir_print.c
+++ b/src/ir_print.c
@@ -259,15 +259,7 @@ void ir_print_type(irFileBuffer *f, irModule *m, Type *t) {
 			i64 align = type_align_of(heap_allocator(), t);
 			i64 total_size = type_size_of(heap_allocator(), t);
 		#if 1
-			i64 fields_size = 0;
-			if (t->Record.field_count > 0) {
-				type_set_offsets(m->allocator, t);
-				isize end_index = t->Record.field_count-1;
-				i64 end_offset = t->Record.offsets[end_index];
-				isize end_size = type_size_of(m->allocator, t->Record.fields[end_index]->type);
-				fields_size = align_formula(end_offset + end_size, build_context.word_size);
-			}
-			i64 block_size = total_size - fields_size -  build_context.word_size;
+			i64 block_size =  t->Record.variant_block_size;
 
 			ir_fprintf(f, "{[0 x <%lld x i8>], ", align);
 			for (isize i = 0; i < t->Record.field_count; i++) {
diff --git a/src/types.c b/src/types.c
index 4001253f8..e363c6b2c 100644
--- a/src/types.c
+++ b/src/types.c
@@ -93,6 +93,7 @@ typedef struct TypeRecord {
 	Entity **variants;
 	i32      variant_count;
 	Entity * union__tag;
+	i64      variant_block_size; // NOTE(bill): Internal use only
 
 
 	i64 *    offsets;
@@ -1817,18 +1818,18 @@ i64 type_size_of_internal(gbAllocator allocator, Type *t, TypePath *path) {
 		total_size = (total_size_in_bits+7)/8;
 		return total_size;
 #else
-		i64 count, align, size, alignment;
-		count = t->Vector.count;
+		i64 count = t->Vector.count;
 		if (count == 0) {
 			return 0;
 		}
-		align = type_align_of_internal(allocator, t->Vector.elem, path);
+		i64 elem_align = type_align_of_internal(allocator, t->Vector.elem, path);
 		if (path->failure) {
 			return FAILURE_SIZE;
 		}
-		size  = type_size_of_internal( allocator, t->Vector.elem, path);
-		alignment = align_formula(size, align);
-		return alignment*(count-1) + size;
+		i64 vector_align = type_align_of_internal(allocator, t, path);
+		i64 elem_size = type_size_of_internal(allocator, t->Vector.elem, path);
+		i64 alignment = align_formula(elem_size, elem_align);
+		return align_formula(alignment*(count-1) + elem_size, vector_align);
 #endif
 	} break;
 
@@ -1904,6 +1905,7 @@ i64 type_size_of_internal(gbAllocator allocator, Type *t, TypePath *path) {
 				i64 end_size = type_size_of_internal(allocator, end_type, path);
 				max = end_offset + end_size ;
 			}
+			i64 field_size = max;
 
 			for (isize i = 1; i < variant_count; i++) {
 				Type *variant_type = t->Record.variants[i]->type;
@@ -1912,9 +1914,13 @@ i64 type_size_of_internal(gbAllocator allocator, Type *t, TypePath *path) {
 					max = size;
 				}
 			}
+
 			// NOTE(bill): Align to int
 			i64 size = align_formula(max, build_context.word_size);
-			size += type_size_of_internal(allocator, t_int, path);
+			// NOTE(bill): Calculate the padding between the common fields and the tag
+			t->Record.variant_block_size = size - field_size;
+
+			size += type_size_of(allocator, t_int);
 			size = align_formula(size, align);
 			return size;
 		} break;