Add support for `#soa[N]T` compound literals

2 files changed, 249 insertions, 77 deletions

diff --git a/src/check_expr.cpp b/src/check_expr.cpp
index 99d464da5..9308eab8e 100644
--- a/src/check_expr.cpp
+++ b/src/check_expr.cpp
@@ -9464,6 +9464,7 @@ gb_internal ExprKind check_compound_literal(CheckerContext *c, Operand *o, Ast *
 	}
 	bool is_to_be_determined_array_count = false;
 	bool is_constant = true;
+	bool is_soa = false;
 
 	Ast *type_expr = cl->type;
 
@@ -9496,8 +9497,14 @@ gb_internal ExprKind check_compound_literal(CheckerContext *c, Operand *o, Ast *
 					GB_ASSERT(tag->kind == Ast_BasicDirective);
 					String name = tag->BasicDirective.name.string;
 					if (name == "soa") {
-						error(node, "#soa arrays are not supported for compound literals");
-						return kind;
+						is_soa = true;
+						if (count == nullptr) {
+							error(node, "#soa slices are not supported for compound literals");
+							return kind;
+						} else if (count->kind == Ast_UnaryExpr &&
+						           count->UnaryExpr.op.kind == Token_Question) {
+							error(node, "#soa fixed length arrays must specify their length and cannot use ?");
+						}
 					}
 				}
 			}
@@ -9507,7 +9514,8 @@ gb_internal ExprKind check_compound_literal(CheckerContext *c, Operand *o, Ast *
 				GB_ASSERT(tag->kind == Ast_BasicDirective);
 				String name = tag->BasicDirective.name.string;
 				if (name == "soa") {
-					error(node, "#soa arrays are not supported for compound literals");
+					is_soa = true;
+					error(node, "#soa dynamic arrays are not supported for compound literals");
 					return kind;
 				}
 			}
@@ -9536,101 +9544,101 @@ gb_internal ExprKind check_compound_literal(CheckerContext *c, Operand *o, Ast *
 
 
 	switch (t->kind) {
-	case Type_Struct: {
+	case Type_Struct:
 		if (cl->elems.count == 0) {
 			break; // NOTE(bill): No need to init
 		}
 
-		if (t->Struct.soa_kind != StructSoa_None) {
-			error(node, "#soa arrays are not supported for compound literals");
-			break;
-		}
-
-		if (t->Struct.is_raw_union) {
-			if (cl->elems.count > 0) {
-				// NOTE: unions cannot be constant
-				is_constant = false;
+		if (t->Struct.soa_kind == StructSoa_None) {
+			if (t->Struct.is_raw_union) {
+				if (cl->elems.count > 0) {
+					// NOTE: unions cannot be constant
+					is_constant = false;
 
-				if (cl->elems[0]->kind != Ast_FieldValue) {
-					gbString type_str = type_to_string(type);
-					error(node, "%s ('struct #raw_union') compound literals are only allowed to contain 'field = value' elements", type_str);
-					gb_string_free(type_str);
-				} else {
-					if (cl->elems.count != 1) {
+					if (cl->elems[0]->kind != Ast_FieldValue) {
 						gbString type_str = type_to_string(type);
-						error(node, "%s ('struct #raw_union') compound literals are only allowed to contain up to 1 'field = value' element, got %td", type_str, cl->elems.count);
+						error(node, "%s ('struct #raw_union') compound literals are only allowed to contain 'field = value' elements", type_str);
 						gb_string_free(type_str);
 					} else {
-						check_compound_literal_field_values(c, cl->elems, o, type, is_constant);
+						if (cl->elems.count != 1) {
+							gbString type_str = type_to_string(type);
+							error(node, "%s ('struct #raw_union') compound literals are only allowed to contain up to 1 'field = value' element, got %td", type_str, cl->elems.count);
+							gb_string_free(type_str);
+						} else {
+							check_compound_literal_field_values(c, cl->elems, o, type, is_constant);
+						}
 					}
 				}
-			}
-			break;
-		}
-
-		wait_signal_until_available(&t->Struct.fields_wait_signal);
-		isize field_count = t->Struct.fields.count;
-		isize min_field_count = t->Struct.fields.count;
-		for (isize i = min_field_count-1; i >= 0; i--) {
-			Entity *e = t->Struct.fields[i];
-			GB_ASSERT(e->kind == Entity_Variable);
-			if (e->Variable.param_value.kind != ParameterValue_Invalid) {
-				min_field_count--;
-			} else {
 				break;
 			}
-		}
 
-		if (cl->elems[0]->kind == Ast_FieldValue) {
-			check_compound_literal_field_values(c, cl->elems, o, type, is_constant);
-		} else {
-			bool seen_field_value = false;
-
-			for_array(index, cl->elems) {
-				Entity *field = nullptr;
-				Ast *elem = cl->elems[index];
-				if (elem->kind == Ast_FieldValue) {
-					seen_field_value = true;
-					error(elem, "Mixture of 'field = value' and value elements in a literal is not allowed");
-					continue;
-				} else if (seen_field_value) {
-					error(elem, "Value elements cannot be used after a 'field = value'");
-					continue;
-				}
-				if (index >= field_count) {
-					error(elem, "Too many values in structure literal, expected %td, got %td", field_count, cl->elems.count);
+			wait_signal_until_available(&t->Struct.fields_wait_signal);
+			isize field_count = t->Struct.fields.count;
+			isize min_field_count = t->Struct.fields.count;
+			for (isize i = min_field_count-1; i >= 0; i--) {
+				Entity *e = t->Struct.fields[i];
+				GB_ASSERT(e->kind == Entity_Variable);
+				if (e->Variable.param_value.kind != ParameterValue_Invalid) {
+					min_field_count--;
+				} else {
 					break;
 				}
+			}
 
-				if (field == nullptr) {
-					field = t->Struct.fields[index];
-				}
+			if (cl->elems[0]->kind == Ast_FieldValue) {
+				check_compound_literal_field_values(c, cl->elems, o, type, is_constant);
+			} else {
+				bool seen_field_value = false;
 
-				Operand o = {};
-				check_expr_or_type(c, &o, elem, field->type);
+				for_array(index, cl->elems) {
+					Entity *field = nullptr;
+					Ast *elem = cl->elems[index];
+					if (elem->kind == Ast_FieldValue) {
+						seen_field_value = true;
+						error(elem, "Mixture of 'field = value' and value elements in a literal is not allowed");
+						continue;
+					} else if (seen_field_value) {
+						error(elem, "Value elements cannot be used after a 'field = value'");
+						continue;
+					}
+					if (index >= field_count) {
+						error(elem, "Too many values in structure literal, expected %td, got %td", field_count, cl->elems.count);
+						break;
+					}
 
-				if (is_type_any(field->type) || is_type_union(field->type) || is_type_raw_union(field->type) || is_type_typeid(field->type)) {
-					is_constant = false;
-				}
-				if (is_constant) {
-					is_constant = check_is_operand_compound_lit_constant(c, &o);
-				}
+					if (field == nullptr) {
+						field = t->Struct.fields[index];
+					}
 
-				check_assignment(c, &o, field->type, str_lit("structure literal"));
-			}
-			if (cl->elems.count < field_count) {
-				if (min_field_count < field_count) {
-				    if (cl->elems.count < min_field_count) {
-						error(cl->close, "Too few values in structure literal, expected at least %td, got %td", min_field_count, cl->elems.count);
-				    }
-				} else {
-					error(cl->close, "Too few values in structure literal, expected %td, got %td", field_count, cl->elems.count);
+					Operand o = {};
+					check_expr_or_type(c, &o, elem, field->type);
+
+					if (is_type_any(field->type) || is_type_union(field->type) || is_type_raw_union(field->type) || is_type_typeid(field->type)) {
+						is_constant = false;
+					}
+					if (is_constant) {
+						is_constant = check_is_operand_compound_lit_constant(c, &o);
+					}
+
+					check_assignment(c, &o, field->type, str_lit("structure literal"));
+				}
+				if (cl->elems.count < field_count) {
+					if (min_field_count < field_count) {
+					    if (cl->elems.count < min_field_count) {
+							error(cl->close, "Too few values in structure literal, expected at least %td, got %td", min_field_count, cl->elems.count);
+					    }
+					} else {
+						error(cl->close, "Too few values in structure literal, expected %td, got %td", field_count, cl->elems.count);
+					}
 				}
 			}
-		}
 
-		break;
-	}
+			break;
+		} else if (t->Struct.soa_kind != StructSoa_Fixed) {
+			error(node, "#soa slices and dynamic arrays are not supported for compound literals");
+			break;
+		}
+		/*fallthrough*/
 
 	case Type_Slice:
 	case Type_Array:
@@ -9641,7 +9649,14 @@ gb_internal ExprKind check_compound_literal(CheckerContext *c, Operand *o, Ast *
 		Type *elem_type = nullptr;
 		String context_name = {};
 		i64 max_type_count = -1;
-		if (t->kind == Type_Slice) {
+		if (t->kind == Type_Struct) {
+			GB_ASSERT(t->Struct.soa_kind == StructSoa_Fixed);
+			elem_type = t->Struct.soa_elem;
+			context_name = str_lit("#soa array literal");
+			if (!is_to_be_determined_array_count) {
+				max_type_count = t->Struct.soa_count;
+			}
+		} else if (t->kind == Type_Slice) {
 			elem_type = t->Slice.elem;
 			context_name = str_lit("slice literal");
 		} else if (t->kind == Type_Array) {
@@ -9650,6 +9665,12 @@ gb_internal ExprKind check_compound_literal(CheckerContext *c, Operand *o, Ast *
 			if (!is_to_be_determined_array_count) {
 				max_type_count = t->Array.count;
 			}
+		} else if (t->kind == Type_Array) {
+			elem_type = t->Array.elem;
+			context_name = str_lit("array literal");
+			if (!is_to_be_determined_array_count) {
+				max_type_count = t->Array.count;
+			}
 		} else if (t->kind == Type_DynamicArray) {
 			elem_type = t->DynamicArray.elem;
 			context_name = str_lit("dynamic array literal");
@@ -9817,6 +9838,15 @@ gb_internal ExprKind check_compound_literal(CheckerContext *c, Operand *o, Ast *
 					error(node, "Expected %lld values for this array literal, got %lld", cast(long long)t->Array.count, cast(long long)max);
 				}
 			}
+		} else if (t->kind == Type_Struct) {
+			GB_ASSERT(t->Struct.soa_kind == StructSoa_Fixed);
+			if (is_to_be_determined_array_count) {
+				t->Struct.soa_count = cast(i32)max;
+			} else if (cl->elems.count > 0 && cl->elems[0]->kind != Ast_FieldValue) {
+				if (0 < max && max < t->Struct.soa_count) {
+					error(node, "Expected %lld values for this #soa array literal, got %lld", cast(long long)t->Struct.soa_count, cast(long long)max);
+				}
+			}
 		}
 
 
diff --git a/src/llvm_backend_const.cpp b/src/llvm_backend_const.cpp
index 51c8a4449..02bb7473c 100644
--- a/src/llvm_backend_const.cpp
+++ b/src/llvm_backend_const.cpp
@@ -851,6 +851,148 @@ gb_internal lbValue lb_const_value(lbModule *m, Type *type, ExactValue value, lb
 	case ExactValue_Compound:
 		if (is_type_slice(type)) {
 			return lb_const_value(m, type, value, cc);
+		} else if (is_type_soa_struct(type)) {
+			GB_ASSERT(type->kind == Type_Struct);
+			GB_ASSERT(type->Struct.soa_kind == StructSoa_Fixed);
+			ast_node(cl, CompoundLit, value.value_compound);
+			Type *elem_type = type->Struct.soa_elem;
+			isize elem_count = cl->elems.count;
+			if (elem_count == 0 || !elem_type_can_be_constant(elem_type)) {
+				return lb_const_nil(m, original_type);
+			}
+			if (cl->elems[0]->kind == Ast_FieldValue) {
+				TEMPORARY_ALLOCATOR_GUARD();
+
+				// TODO(bill): This is O(N*M) and will be quite slow; it should probably be sorted before hand
+
+				isize elem_count = cast(isize)type->Struct.soa_count;
+
+				LLVMValueRef *aos_values = gb_alloc_array(temporary_allocator(), LLVMValueRef, elem_count);
+
+				isize value_index = 0;
+				for (i64 i = 0; i < elem_count; i++) {
+					bool found = false;
+
+					for (isize j = 0; j < elem_count; j++) {
+						Ast *elem = cl->elems[j];
+						ast_node(fv, FieldValue, elem);
+						if (is_ast_range(fv->field)) {
+							ast_node(ie, BinaryExpr, fv->field);
+							TypeAndValue lo_tav = ie->left->tav;
+							TypeAndValue hi_tav = ie->right->tav;
+							GB_ASSERT(lo_tav.mode == Addressing_Constant);
+							GB_ASSERT(hi_tav.mode == Addressing_Constant);
+
+							TokenKind op = ie->op.kind;
+							i64 lo = exact_value_to_i64(lo_tav.value);
+							i64 hi = exact_value_to_i64(hi_tav.value);
+							if (op != Token_RangeHalf) {
+								hi += 1;
+							}
+							if (lo == i) {
+								TypeAndValue tav = fv->value->tav;
+								LLVMValueRef val = lb_const_value(m, elem_type, tav.value, cc).value;
+								for (i64 k = lo; k < hi; k++) {
+									aos_values[value_index++] = val;
+								}
+
+								found = true;
+								i += (hi-lo-1);
+								break;
+							}
+						} else {
+							TypeAndValue index_tav = fv->field->tav;
+							GB_ASSERT(index_tav.mode == Addressing_Constant);
+							i64 index = exact_value_to_i64(index_tav.value);
+							if (index == i) {
+								TypeAndValue tav = fv->value->tav;
+								LLVMValueRef val = lb_const_value(m, elem_type, tav.value, cc).value;
+								aos_values[value_index++] = val;
+								found = true;
+								break;
+							}
+						}
+					}
+
+					if (!found) {
+						aos_values[value_index++] = nullptr;
+					}
+				}
+
+
+				isize field_count = type->Struct.fields.count;
+				LLVMValueRef *soa_values = gb_alloc_array(temporary_allocator(), LLVMValueRef, field_count);
+
+				for (isize i = 0; i < field_count; i++) {
+					TEMPORARY_ALLOCATOR_GUARD();
+
+					LLVMValueRef *values = gb_alloc_array(temporary_allocator(), LLVMValueRef, elem_count);
+
+					Entity *f = type->Struct.fields[i];
+					Type *array_type = f->type;
+					GB_ASSERT(array_type->kind == Type_Array);
+					Type *field_type = array_type->Array.elem;
+
+					for (isize j = 0; j < elem_count; j++) {
+						LLVMValueRef v = aos_values[j];
+						if (v != nullptr) {
+							values[j] = llvm_const_extract_value(m, v, cast(unsigned)i);
+						} else {
+							values[j] = LLVMConstNull(lb_type(m, field_type));
+						}
+					}
+
+					soa_values[i] = lb_build_constant_array_values(m, array_type, field_type, elem_count, values, cc);
+				}
+
+				res.value = llvm_const_named_struct(m, type, soa_values, field_count);
+				return res;
+			} else {
+				GB_ASSERT_MSG(elem_count == type->Struct.soa_count, "%td != %td", elem_count, type->Struct.soa_count);
+
+				TEMPORARY_ALLOCATOR_GUARD();
+
+				isize elem_count = cast(isize)type->Struct.soa_count;
+
+				LLVMValueRef *aos_values = gb_alloc_array(temporary_allocator(), LLVMValueRef, elem_count);
+
+				for (isize i = 0; i < elem_count; i++) {
+					TypeAndValue tav = cl->elems[i]->tav;
+					GB_ASSERT(tav.mode != Addressing_Invalid);
+					aos_values[i] = lb_const_value(m, elem_type, tav.value, cc).value;
+				}
+				for (isize i = elem_count; i < type->Struct.soa_count; i++) {
+					aos_values[i] = nullptr;
+				}
+
+				isize field_count = type->Struct.fields.count;
+				LLVMValueRef *soa_values = gb_alloc_array(temporary_allocator(), LLVMValueRef, field_count);
+
+				for (isize i = 0; i < field_count; i++) {
+					TEMPORARY_ALLOCATOR_GUARD();
+
+					LLVMValueRef *values = gb_alloc_array(temporary_allocator(), LLVMValueRef, elem_count);
+
+					Entity *f = type->Struct.fields[i];
+					Type *array_type = f->type;
+					GB_ASSERT(array_type->kind == Type_Array);
+					Type *field_type = array_type->Array.elem;
+
+					for (isize j = 0; j < elem_count; j++) {
+						LLVMValueRef v = aos_values[j];
+						if (v != nullptr) {
+							values[j] = llvm_const_extract_value(m, v, cast(unsigned)i);
+						} else {
+							values[j] = LLVMConstNull(lb_type(m, field_type));
+						}
+					}
+
+					soa_values[i] = lb_build_constant_array_values(m, array_type, field_type, elem_count, values, cc);
+				}
+
+				res.value = llvm_const_named_struct(m, type, soa_values, field_count);
+				return res;
+			}
 		} else if (is_type_array(type)) {
 			ast_node(cl, CompoundLit, value.value_compound);
 			Type *elem_type = type->Array.elem;