Support swizzle selector syntax `.xyzw` for `#simd` vectors

4 files changed, 82 insertions, 32 deletions

diff --git a/src/check_expr.cpp b/src/check_expr.cpp
index 3a8cdf0b1..0911e48cf 100644
--- a/src/check_expr.cpp
+++ b/src/check_expr.cpp
@@ -4920,7 +4920,7 @@ gb_internal Entity *check_selector(CheckerContext *c, Operand *operand, Ast *nod
 		}
 	}
 
-	if (entity == nullptr && selector->kind == Ast_Ident && is_type_array(type_deref(operand->type))) {
+	if (entity == nullptr && selector->kind == Ast_Ident && (is_type_array(type_deref(operand->type)) || is_type_simd_vector(type_deref(operand->type)))) {
 		String field_name = selector->Ident.token.string;
 		if (1 < field_name.len && field_name.len <= 4) {
 			u8 swizzles_xyzw[4] = {'x', 'y', 'z', 'w'};
@@ -4975,8 +4975,10 @@ gb_internal Entity *check_selector(CheckerContext *c, Operand *operand, Ast *nod
 
 			Type *original_type = operand->type;
 			Type *array_type = base_type(type_deref(original_type));
-			GB_ASSERT(array_type->kind == Type_Array);
-			i64 array_count = array_type->Array.count;
+			GB_ASSERT(array_type->kind == Type_Array || array_type->kind == Type_SimdVector);
+
+			i64 array_count = get_array_type_count(array_type);
+
 			for (u8 i = 0; i < index_count; i++) {
 				u8 idx = indices>>(i*2) & 3;
 				if (idx >= array_count) {
@@ -4996,7 +4998,6 @@ gb_internal Entity *check_selector(CheckerContext *c, Operand *operand, Ast *nod
 			se->swizzle_count = index_count;
 			se->swizzle_indices = indices;
 
-
 			AddressingMode prev_mode = operand->mode;
 			operand->mode = Addressing_SwizzleValue;
 			operand->type = determine_swizzle_array_type(original_type, type_hint, index_count);
@@ -5010,6 +5011,10 @@ gb_internal Entity *check_selector(CheckerContext *c, Operand *operand, Ast *nod
 				break;
 			}
 
+			if (array_type->kind == Type_SimdVector) {
+				operand->mode = Addressing_Value;
+			}
+
 			Entity *swizzle_entity = alloc_entity_variable(nullptr, make_token_ident(field_name), operand->type, EntityState_Resolved);
 			add_type_and_value(c, operand->expr, operand->mode, operand->type, operand->value);
 			return swizzle_entity;
diff --git a/src/llvm_backend_expr.cpp b/src/llvm_backend_expr.cpp
index 5bc961af2..98618798b 100644
--- a/src/llvm_backend_expr.cpp
+++ b/src/llvm_backend_expr.cpp
@@ -4655,7 +4655,7 @@ gb_internal lbAddr lb_build_addr_internal(lbProcedure *p, Ast *expr) {
 
 			if (se->swizzle_count > 0) {
 				Type *array_type = base_type(type_deref(tav.type));
-				GB_ASSERT(array_type->kind == Type_Array);
+				GB_ASSERT(array_type->kind == Type_Array || array_type->kind == Type_SimdVector);
 				u8 swizzle_count = se->swizzle_count;
 				u8 swizzle_indices_raw = se->swizzle_indices;
 				u8 swizzle_indices[4] = {};
@@ -4671,7 +4671,7 @@ gb_internal lbAddr lb_build_addr_internal(lbProcedure *p, Ast *expr) {
 					a = lb_addr_get_ptr(p, addr);
 				}
 
-				GB_ASSERT(is_type_array(expr->tav.type));
+				GB_ASSERT(is_type_array(expr->tav.type) || is_type_simd_vector(expr->tav.type));
 				return lb_addr_swizzle(a, expr->tav.type, swizzle_count, swizzle_indices);
 			}
 
diff --git a/src/llvm_backend_general.cpp b/src/llvm_backend_general.cpp
index 4ff8482a7..09de90dc9 100644
--- a/src/llvm_backend_general.cpp
+++ b/src/llvm_backend_general.cpp
@@ -434,7 +434,7 @@ gb_internal lbAddr lb_addr_soa_variable(lbValue addr, lbValue index, Ast *index_
 }
 
 gb_internal lbAddr lb_addr_swizzle(lbValue addr, Type *array_type, u8 swizzle_count, u8 swizzle_indices[4]) {
-	GB_ASSERT(is_type_array(array_type));
+	GB_ASSERT(is_type_array(array_type) || is_type_simd_vector(array_type));
 	GB_ASSERT(1 < swizzle_count && swizzle_count <= 4);
 	lbAddr v = {lbAddr_Swizzle, addr};
 	v.swizzle.type = array_type;
@@ -1264,6 +1264,30 @@ gb_internal lbValue lb_addr_load(lbProcedure *p, lbAddr const &addr) {
 		return lb_addr_load(p, res);
 	} else if (addr.kind == lbAddr_Swizzle) {
 		Type *array_type = base_type(addr.swizzle.type);
+		if (array_type->kind == Type_SimdVector) {
+			lbValue vec = lb_emit_load(p, addr.addr);
+			u8 index_count = addr.swizzle.count;
+			if (index_count == 0) {
+				return vec;
+			}
+
+			unsigned mask_len = cast(unsigned)index_count;
+			LLVMValueRef *mask_elems = gb_alloc_array(permanent_allocator(), LLVMValueRef, index_count);
+			for (isize i = 0; i < index_count; i++) {
+				mask_elems[i] = LLVMConstInt(lb_type(p->module, t_u32), addr.swizzle.indices[i], false);
+			}
+
+			LLVMValueRef mask = LLVMConstVector(mask_elems, mask_len);
+
+			LLVMValueRef v1 = vec.value;
+			LLVMValueRef v2 = vec.value;
+
+			lbValue res = {};
+			res.type = addr.swizzle.type;
+			res.value = LLVMBuildShuffleVector(p->builder, v1, v2, mask, "");
+			return res;
+		}
+
 		GB_ASSERT(array_type->kind == Type_Array);
 
 		unsigned res_align = cast(unsigned)type_align_of(addr.swizzle.type);
diff --git a/src/types.cpp b/src/types.cpp
index e9e91dcd4..5a3ad5d6b 100644
--- a/src/types.cpp
+++ b/src/types.cpp
@@ -3430,31 +3430,6 @@ gb_internal Selection lookup_field_with_selection(Type *type_, String field_name
 		}
 
 		return sel;
-	} else if (type->kind == Type_Array) {
-		if (type->Array.count <= 4) {
-			// HACK(bill): Memory leak
-			switch (type->Array.count) {
-			#define _ARRAY_FIELD_CASE_IF(_length, _name) \
-				if (field_name == (_name)) { \
-					selection_add_index(&sel, (_length)-1); \
-					sel.entity = alloc_entity_array_elem(nullptr, make_token_ident(str_lit(_name)), type->Array.elem, (_length)-1); \
-					return sel; \
-				}
-			#define _ARRAY_FIELD_CASE(_length, _name0, _name1) \
-			case (_length): \
-				_ARRAY_FIELD_CASE_IF(_length, _name0); \
-				_ARRAY_FIELD_CASE_IF(_length, _name1); \
-				/*fallthrough*/
-
-			_ARRAY_FIELD_CASE(4, "w", "a");
-			_ARRAY_FIELD_CASE(3, "z", "b");
-			_ARRAY_FIELD_CASE(2, "y", "g");
-			_ARRAY_FIELD_CASE(1, "x", "r");
-			default: break;
-
-			#undef _ARRAY_FIELD_CASE
-			}
-		}
 	} else if (type->kind == Type_DynamicArray) {
 		GB_ASSERT(t_allocator != nullptr);
 		String allocator_str = str_lit("allocator");
@@ -3475,7 +3450,53 @@ gb_internal Selection lookup_field_with_selection(Type *type_, String field_name
 			sel.entity = entity__allocator;
 			return sel;
 		}
+
+
+#define _ARRAY_FIELD_CASE_IF(_length, _name) \
+	if (field_name == (_name)) { \
+		selection_add_index(&sel, (_length)-1); \
+		sel.entity = alloc_entity_array_elem(nullptr, make_token_ident(str_lit(_name)), elem, (_length)-1); \
+		return sel; \
 	}
+#define _ARRAY_FIELD_CASE(_length, _name0, _name1) \
+case (_length): \
+	_ARRAY_FIELD_CASE_IF(_length, _name0); \
+	_ARRAY_FIELD_CASE_IF(_length, _name1); \
+	/*fallthrough*/
+
+
+	} else if (type->kind == Type_Array) {
+
+		Type *elem = type->Array.elem;
+
+		if (type->Array.count <= 4) {
+			// HACK(bill): Memory leak
+			switch (type->Array.count) {
+
+			_ARRAY_FIELD_CASE(4, "w", "a");
+			_ARRAY_FIELD_CASE(3, "z", "b");
+			_ARRAY_FIELD_CASE(2, "y", "g");
+			_ARRAY_FIELD_CASE(1, "x", "r");
+			default: break;
+			}
+		}
+	} else if (type->kind == Type_SimdVector) {
+
+		Type *elem = type->SimdVector.elem;
+		if (type->SimdVector.count <= 4) {
+			// HACK(bill): Memory leak
+			switch (type->SimdVector.count) {
+			_ARRAY_FIELD_CASE(4, "w", "a");
+			_ARRAY_FIELD_CASE(3, "z", "b");
+			_ARRAY_FIELD_CASE(2, "y", "g");
+			_ARRAY_FIELD_CASE(1, "x", "r");
+			default: break;
+			}
+		}
+	}
+
+#undef _ARRAY_FIELD_CASE
+#undef _ARRAY_FIELD_CASE
 
 	return sel;
 }