Unify semantics of the built-in `swizzle` procedure with the selector expression semantics e.g. `.xyz`

6 files changed, 174 insertions, 27 deletions

diff --git a/src/check_builtin.cpp b/src/check_builtin.cpp
index ea89ec007..c0ba40503 100644
--- a/src/check_builtin.cpp
+++ b/src/check_builtin.cpp
@@ -778,12 +778,31 @@ bool check_builtin_procedure(CheckerContext *c, Operand *operand, Ast *call, i32
 		if (arg_count < max_count) {
 			operand->type = alloc_type_array(elem_type, arg_count);
 		}
-		operand->mode = Addressing_Value;
+		if (type->kind == Type_Array) {
+			if (operand->mode == Addressing_Variable) {
+				operand->mode = Addressing_SwizzleVariable;
+			} else {
+				operand->mode = Addressing_SwizzleValue;
+			}
+		} else {
+			operand->mode = Addressing_Value;
+		}
 
-		if (type_hint != nullptr && check_is_castable_to(c, operand, type_hint)) {
-			operand->type = type_hint;
+		Type *array_type = base_type(type_deref(operand->type));
+		GB_ASSERT(array_type->kind == Type_Array);
+
+		Type *swizzle_array_type = nullptr;
+		Type *bth = base_type(type_hint);
+		if (bth != nullptr && bth->kind == Type_Array &&
+		    bth->Array.count == arg_count &&
+		    are_types_identical(bth->Array.elem, array_type->Array.elem)) {
+			swizzle_array_type = type_hint;
+		} else {
+			swizzle_array_type = alloc_type_array(array_type->Array.elem, arg_count);
 		}
 
+		operand->type = swizzle_array_type;
+
 		break;
 	}
 
diff --git a/src/llvm_backend.hpp b/src/llvm_backend.hpp
index e0c2171fd..c5f5897c1 100644
--- a/src/llvm_backend.hpp
+++ b/src/llvm_backend.hpp
@@ -48,6 +48,7 @@ enum lbAddrKind {
 	lbAddr_RelativeSlice,
 
 	lbAddr_Swizzle,
+	lbAddr_SwizzleLarge,
 };
 
 struct lbAddr {
@@ -78,6 +79,10 @@ struct lbAddr {
 			u8 count;      // 2, 3, or 4 components
 			u8 indices[4];
 		} swizzle;
+		struct {
+			Type *type;
+			Slice<i32> indices;
+		} swizzle_large;
 	};
 };
 
diff --git a/src/llvm_backend_expr.cpp b/src/llvm_backend_expr.cpp
index e5a57d3d3..1c5f03722 100644
--- a/src/llvm_backend_expr.cpp
+++ b/src/llvm_backend_expr.cpp
@@ -2497,6 +2497,41 @@ lbAddr lb_build_addr_from_entity(lbProcedure *p, Entity *e, Ast *expr) {
 	return lb_addr(v);
 }
 
+lbAddr lb_build_array_swizzle_addr(lbProcedure *p, AstCallExpr *ce, TypeAndValue const &tv) {
+	isize index_count = ce->args.count-1;
+	lbAddr addr = lb_build_addr(p, ce->args[0]);
+	if (index_count == 0) {
+		return addr;
+	}
+	Type *type = base_type(lb_addr_type(addr));
+	GB_ASSERT(type->kind == Type_Array);
+	i64 count = type->Array.count;
+	if (count <= 4) {
+		u8 indices[4] = {};
+		u8 index_count = 0;
+		for (i32 i = 1; i < ce->args.count; i++) {
+			TypeAndValue tv = type_and_value_of_expr(ce->args[i]);
+			GB_ASSERT(is_type_integer(tv.type));
+			GB_ASSERT(tv.value.kind == ExactValue_Integer);
+
+			i64 src_index = big_int_to_i64(&tv.value.value_integer);
+			indices[index_count++] = cast(u8)src_index;
+		}
+		return lb_addr_swizzle(lb_addr_get_ptr(p, addr), tv.type, index_count, indices);
+	}
+	auto indices = slice_make<i32>(permanent_allocator(), ce->args.count-1);
+	isize index_index = 0;
+	for (i32 i = 1; i < ce->args.count; i++) {
+		TypeAndValue tv = type_and_value_of_expr(ce->args[i]);
+		GB_ASSERT(is_type_integer(tv.type));
+		GB_ASSERT(tv.value.kind == ExactValue_Integer);
+
+		i64 src_index = big_int_to_i64(&tv.value.value_integer);
+		indices[index_index++] = cast(i32)src_index;
+	}
+	return lb_addr_swizzle_large(lb_addr_get_ptr(p, addr), tv.type, indices);
+}
+
 
 lbAddr lb_build_addr(lbProcedure *p, Ast *expr) {
 	expr = unparen_expr(expr);
@@ -2617,6 +2652,10 @@ lbAddr lb_build_addr(lbProcedure *p, Ast *expr) {
 					GB_ASSERT(sel.index.count > 0);
 					// NOTE(bill): just patch the index in place
 					sel.index[0] = addr.swizzle.indices[sel.index[0]];
+				} else if (addr.kind == lbAddr_SwizzleLarge) {
+					GB_ASSERT(sel.index.count > 0);
+					// NOTE(bill): just patch the index in place
+					sel.index[0] = addr.swizzle.indices[sel.index[0]];
 				}
 				lbValue a = lb_addr_get_ptr(p, addr);
 				a = lb_emit_deep_field_gep(p, a, sel);
@@ -3028,6 +3067,23 @@ lbAddr lb_build_addr(lbProcedure *p, Ast *expr) {
 	case_end;
 
 	case_ast_node(ce, CallExpr, expr);
+		BuiltinProcId builtin_id = BuiltinProc_Invalid;
+		if (ce->proc->tav.mode == Addressing_Builtin) {
+			Entity *e = entity_of_node(ce->proc);
+			if (e != nullptr) {
+				builtin_id = cast(BuiltinProcId)e->Builtin.id;
+			} else {
+				builtin_id = BuiltinProc_DIRECTIVE;
+			}
+		}
+		auto const &tv = expr->tav;
+		if (builtin_id == BuiltinProc_swizzle &&
+		    is_type_array(tv.type)) {
+		    	// NOTE(bill, 2021-08-09): `swizzle` has some bizarre semantics so it needs to be
+		    	// specialized here for to be addressable
+			return lb_build_array_swizzle_addr(p, ce, tv);
+		}
+
 		// NOTE(bill): This is make sure you never need to have an 'array_ev'
 		lbValue e = lb_build_expr(p, expr);
 	#if 1
diff --git a/src/llvm_backend_general.cpp b/src/llvm_backend_general.cpp
index 61625a374..74aaafe94 100644
--- a/src/llvm_backend_general.cpp
+++ b/src/llvm_backend_general.cpp
@@ -301,6 +301,14 @@ lbAddr lb_addr_swizzle(lbValue addr, Type *array_type, u8 swizzle_count, u8 swiz
 	return v;
 }
 
+lbAddr lb_addr_swizzle_large(lbValue addr, Type *array_type, Slice<i32> const &swizzle_indices) {
+	GB_ASSERT_MSG(is_type_array(array_type), "%s", type_to_string(array_type));
+	lbAddr v = {lbAddr_SwizzleLarge, addr};
+	v.swizzle_large.type = array_type;
+	v.swizzle_large.indices = swizzle_indices;
+	return v;
+}
+
 Type *lb_addr_type(lbAddr const &addr) {
 	if (addr.addr.value == nullptr) {
 		return nullptr;
@@ -313,6 +321,9 @@ Type *lb_addr_type(lbAddr const &addr) {
 	if (addr.kind == lbAddr_Swizzle) {
 		return addr.swizzle.type;
 	}
+	if (addr.kind == lbAddr_SwizzleLarge) {
+		return addr.swizzle_large.type;
+	}
 	return type_deref(addr.addr.type);
 }
 LLVMTypeRef lb_addr_lb_type(lbAddr const &addr) {
@@ -372,6 +383,7 @@ lbValue lb_addr_get_ptr(lbProcedure *p, lbAddr const &addr) {
 		break;
 
 	case lbAddr_Swizzle:
+	case lbAddr_SwizzleLarge:
 		// TOOD(bill): is this good enough logic?
 		break;
 	}
@@ -692,6 +704,19 @@ void lb_addr_store(lbProcedure *p, lbAddr addr, lbValue value) {
 			}
 		}
 		return;
+	} else if (addr.kind == lbAddr_SwizzleLarge) {
+		GB_ASSERT(value.value != nullptr);
+		value = lb_emit_conv(p, value, lb_addr_type(addr));
+
+		lbValue dst = lb_addr_get_ptr(p, addr);
+		lbValue src = lb_address_from_load_or_generate_local(p, value);
+		for_array(i, addr.swizzle_large.indices) {
+			lbValue src_ptr = lb_emit_array_epi(p, src, i);
+			lbValue dst_ptr = lb_emit_array_epi(p, dst, addr.swizzle_large.indices[i]);
+			lbValue src_load = lb_emit_load(p, src_ptr);
+			lb_emit_store(p, dst_ptr, src_load);
+		}
+		return;
 	}
 
 	GB_ASSERT(value.value != nullptr);
@@ -1014,6 +1039,25 @@ lbValue lb_addr_load(lbProcedure *p, lbAddr const &addr) {
 			}
 		}
 		return lb_addr_load(p, res);
+	}  else if (addr.kind == lbAddr_SwizzleLarge) {
+		Type *array_type = base_type(addr.swizzle_large.type);
+		GB_ASSERT(array_type->kind == Type_Array);
+
+		unsigned res_align = cast(unsigned)type_align_of(addr.swizzle_large.type);
+
+		Type *elem_type = base_type(array_type->Array.elem);
+		lbAddr res = lb_add_local_generated(p, addr.swizzle_large.type, false);
+		lbValue ptr = lb_addr_get_ptr(p, res);
+		GB_ASSERT(is_type_pointer(ptr.type));
+
+		for_array(i, addr.swizzle_large.indices) {
+			i32 index = addr.swizzle_large.indices[i];
+			lbValue dst = lb_emit_array_epi(p, ptr, i);
+			lbValue src = lb_emit_array_epi(p, addr.addr, index);
+			lb_emit_store(p, dst, lb_emit_load(p, src));
+		}
+
+		return lb_addr_load(p, res);
 	}
 
 	if (is_type_proc(addr.addr.type)) {
diff --git a/src/llvm_backend_proc.cpp b/src/llvm_backend_proc.cpp
index 955cf03b1..99f824c35 100644
--- a/src/llvm_backend_proc.cpp
+++ b/src/llvm_backend_proc.cpp
@@ -1018,32 +1018,10 @@ lbValue lb_build_builtin_proc(lbProcedure *p, Ast *expr, TypeAndValue const &tv,
 			res.type = tv.type;
 			res.value = LLVMBuildShuffleVector(p->builder, v1, v2, mask, "");
 			return res;
-
-		}
-
-		lbAddr addr = lb_build_addr(p, ce->args[0]);
-		if (index_count == 0) {
-			return lb_addr_load(p, addr);
 		}
-		lbValue src = lb_addr_get_ptr(p, addr);
-		// TODO(bill): Should this be zeroed or not?
-		lbAddr dst = lb_add_local_generated(p, tv.type, true);
-		lbValue dst_ptr = lb_addr_get_ptr(p, dst);
-
-		for (i32 i = 1; i < ce->args.count; i++) {
-			TypeAndValue tv = type_and_value_of_expr(ce->args[i]);
-			GB_ASSERT(is_type_integer(tv.type));
-			GB_ASSERT(tv.value.kind == ExactValue_Integer);
 
-			i32 src_index = cast(i32)big_int_to_i64(&tv.value.value_integer);
-			i32 dst_index = i-1;
-
-			lbValue src_elem = lb_emit_array_epi(p, src, src_index);
-			lbValue dst_elem = lb_emit_array_epi(p, dst_ptr, dst_index);
-
-			lb_emit_store(p, dst_elem, lb_emit_load(p, src_elem));
-		}
-		return lb_addr_load(p, dst);
+		lbAddr addr = lb_build_array_swizzle_addr(p, ce, tv);
+		return lb_addr_load(p, addr);
 	}
 
 	case BuiltinProc_complex: {
diff --git a/src/llvm_backend_stmt.cpp b/src/llvm_backend_stmt.cpp
index 324b408c7..a220995a3 100644
--- a/src/llvm_backend_stmt.cpp
+++ b/src/llvm_backend_stmt.cpp
@@ -1818,6 +1818,51 @@ void lb_build_assign_stmt_array(lbProcedure *p, TokenKind op, lbAddr const &lhs,
 			lb_emit_store(p, lhs_ptrs[i], ops[i]);
 		}
 		return;
+	} else if (lhs.kind == lbAddr_SwizzleLarge) {
+		GB_ASSERT(is_type_array(lhs_type));
+
+		struct ValueAndIndex {
+			lbValue value;
+			u32     index;
+		};
+
+		Type *bt = base_type(lhs_type);
+		GB_ASSERT(bt->kind == Type_Array);
+
+		auto indices_handled = slice_make<bool>(temporary_allocator(), bt->Array.count);
+		auto indices = slice_make<i32>(temporary_allocator(), bt->Array.count);
+		i32 index_count = 0;
+		for_array(i, lhs.swizzle_large.indices) {
+			i32 index = lhs.swizzle_large.indices[i];
+			if (indices_handled[index]) {
+				continue;
+			}
+			indices[index_count++] = index;
+		}
+		gb_sort_array(indices.data, index_count, gb_i32_cmp(0));
+
+		lbValue lhs_ptrs[4] = {};
+		lbValue x_loads[4]  = {};
+		lbValue y_loads[4]  = {};
+		lbValue ops[4]      = {};
+
+		for (i32 i = 0; i < index_count; i++) {
+			lhs_ptrs[i] = lb_emit_array_epi(p, lhs.addr, indices[i]);
+		}
+		for (i32 i = 0; i < index_count; i++) {
+			x_loads[i] = lb_emit_load(p, lhs_ptrs[i]);
+		}
+		for (i32 i = 0; i < index_count; i++) {
+			y_loads[i].value = LLVMBuildExtractValue(p->builder, rhs.value, i, "");
+			y_loads[i].type = elem_type;
+		}
+		for (i32 i = 0; i < index_count; i++) {
+			ops[i] = lb_emit_arith(p, op, x_loads[i], y_loads[i], elem_type);
+		}
+		for (i32 i = 0; i < index_count; i++) {
+			lb_emit_store(p, lhs_ptrs[i], ops[i]);
+		}
+		return;
 	}