Experimental support for inline swizzling for array types of len <= 4 e.g. `v.xyz`, `v.argb`, `v.xxx`

7 files changed, 185 insertions, 74 deletions

diff --git a/src/check_expr.cpp b/src/check_expr.cpp
index e31343e29..eddd7f369 100644
--- a/src/check_expr.cpp
+++ b/src/check_expr.cpp
@@ -3698,6 +3698,94 @@ Entity *check_selector(CheckerContext *c, Operand *operand, Ast *node, Type *typ
 		}
 	}
 
+	if (entity == nullptr  && selector->kind == Ast_Ident && is_type_array(type_deref(operand->type))) {
+		// TODO(bill): Simd_Vector swizzling
+
+		String field_name = selector->Ident.token.string;
+		if (1 < field_name.len && field_name.len <= 4) {
+			u8 swizzles_xyzw[4] = {'x', 'y', 'z', 'w'};
+			u8 swizzles_rgba[4] = {'r', 'g', 'b', 'a'};
+			bool found_xyzw = false;
+			bool found_rgba = false;
+			for (isize i = 0; i < field_name.len; i++) {
+				bool valid = false;
+				for (isize j = 0; j < 4; j++) {
+					if (field_name.text[i] == swizzles_xyzw[j]) {
+						found_xyzw = true;
+						valid = true;
+						break;
+					}
+					if (field_name.text[i] == swizzles_rgba[j]) {
+						found_rgba = true;
+						valid = true;
+						break;
+					}
+				}
+				if (!valid) {
+					goto end_of_array_selector_swizzle;
+				}
+			}
+
+			u8 *swizzles = nullptr;
+
+			u8 index_count = cast(u8)field_name.len;
+			if (found_xyzw && found_rgba) {
+				gbString op_str = expr_to_string(op_expr);
+				error(op_expr, "Mixture of swizzle kinds for field index, got %s", op_str);
+				gb_string_free(op_str);
+				operand->mode = Addressing_Invalid;
+				operand->expr = node;
+				return nullptr;
+			}
+			u8 indices = 0;
+
+			if (found_xyzw) {
+				swizzles = swizzles_xyzw;
+			} else if (found_rgba) {
+				swizzles = swizzles_rgba;
+			}
+			for (isize i = 0; i < field_name.len; i++) {
+				for (isize j = 0; j < 4; j++) {
+					if (field_name.text[i] == swizzles[j]) {
+						indices |= cast(u8)(j)<<(i*2);
+						break;
+					}
+				}
+			}
+
+			se->swizzle_count = index_count;
+			se->swizzle_indices = indices;
+
+			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 == index_count) {
+				swizzle_array_type = type_hint;
+			} else {
+				swizzle_array_type = alloc_type_array(array_type->Array.elem, index_count);
+			}
+			AddressingMode prev_mode = operand->mode;
+			operand->mode = Addressing_SwizzleValue;
+			operand->type = swizzle_array_type;
+			operand->expr = node;
+
+			switch (prev_mode) {
+			case Addressing_Variable:
+			case Addressing_SoaVariable:
+			case Addressing_SwizzleVariable:
+				operand->mode = Addressing_SwizzleVariable;
+				break;
+			}
+
+			Entity *swizzle_entity = alloc_entity_variable(nullptr, make_token_ident(field_name), operand->type, EntityState_Resolved);
+			add_type_and_value(c->info, operand->expr, operand->mode, operand->type, operand->value);
+			return swizzle_entity;
+		}
+	end_of_array_selector_swizzle:;
+	}
+
 	if (entity == nullptr) {
 		gbString op_str   = expr_to_string(op_expr);
 		gbString type_str = type_to_string(operand->type);
diff --git a/src/check_stmt.cpp b/src/check_stmt.cpp
index 7772b5c97..a4b7f9b24 100644
--- a/src/check_stmt.cpp
+++ b/src/check_stmt.cpp
@@ -370,6 +370,9 @@ Type *check_assignment_variable(CheckerContext *ctx, Operand *lhs, Operand *rhs)
 	case Addressing_SoaVariable:
 		break;
 
+	case Addressing_SwizzleVariable:
+		break;
+
 	default: {
 		if (lhs->expr->kind == Ast_SelectorExpr) {
 			// NOTE(bill): Extra error checks
diff --git a/src/entity.cpp b/src/entity.cpp
index 173a3fcd0..30bbef9dc 100644
--- a/src/entity.cpp
+++ b/src/entity.cpp
@@ -40,27 +40,28 @@ enum EntityFlag : u64 {
 	EntityFlag_Param         = 1ull<<4,
 	EntityFlag_Result        = 1ull<<5,
 	EntityFlag_ArrayElem     = 1ull<<6,
-	EntityFlag_Ellipsis      = 1ull<<7,
-	EntityFlag_NoAlias       = 1ull<<8,
-	EntityFlag_TypeField     = 1ull<<9,
-	EntityFlag_Value         = 1ull<<10,
-	EntityFlag_Sret          = 1ull<<11,
-	EntityFlag_ByVal         = 1ull<<12,
-	EntityFlag_BitFieldValue = 1ull<<13,
-	EntityFlag_PolyConst     = 1ull<<14,
-	EntityFlag_NotExported   = 1ull<<15,
-	EntityFlag_ConstInput    = 1ull<<16,
+	EntityFlag_ArraySwizzle  = 1ull<<7,
+	EntityFlag_Ellipsis      = 1ull<<8,
+	EntityFlag_NoAlias       = 1ull<<9,
+	EntityFlag_TypeField     = 1ull<<10,
+	EntityFlag_Value         = 1ull<<11,
+	EntityFlag_Sret          = 1ull<<12,
+	EntityFlag_ByVal         = 1ull<<13,
+	EntityFlag_BitFieldValue = 1ull<<14,
+	EntityFlag_PolyConst     = 1ull<<15,
+	EntityFlag_NotExported   = 1ull<<16,
+	EntityFlag_ConstInput    = 1ull<<17,
 
-	EntityFlag_Static        = 1ull<<17,
+	EntityFlag_Static        = 1ull<<18,
 
-	EntityFlag_ImplicitReference = 1ull<<18, // NOTE(bill): equivalent to `const &` in C++
+	EntityFlag_ImplicitReference = 1ull<<19, // NOTE(bill): equivalent to `const &` in C++
 
-	EntityFlag_SoaPtrField   = 1ull<<19, // to allow s.x[0] where `s.x` is a pointer rather than a slice
+	EntityFlag_SoaPtrField   = 1ull<<20, // to allow s.x[0] where `s.x` is a pointer rather than a slice
 
-	EntityFlag_ProcBodyChecked = 1ull<<20,
+	EntityFlag_ProcBodyChecked = 1ull<<21,
 
-	EntityFlag_CVarArg       = 1ull<<21,
-	EntityFlag_AutoCast      = 1ull<<22,
+	EntityFlag_CVarArg       = 1ull<<22,
+	EntityFlag_AutoCast      = 1ull<<23,
 
 	EntityFlag_Disabled      = 1ull<<24,
 	EntityFlag_Cold          = 1ull<<25, // procedure is rarely called
diff --git a/src/llvm_backend.cpp b/src/llvm_backend.cpp
index eaa621291..c77590294 100644
--- a/src/llvm_backend.cpp
+++ b/src/llvm_backend.cpp
@@ -134,6 +134,15 @@ lbAddr lb_addr_soa_variable(lbValue addr, lbValue index, Ast *index_expr) {
 	return v;
 }
 
+lbAddr lb_addr_swizzle(lbValue addr, Type *array_type, u8 swizzle_count, u8 swizzle_indices[4]) {
+	GB_ASSERT(1 < swizzle_count && swizzle_count <= 4);
+	lbAddr v = {lbAddr_Swizzle, addr};
+	v.swizzle.type = array_type;
+	v.swizzle.count = swizzle_count;
+	gb_memmove(v.swizzle.indices, swizzle_indices, swizzle_count);
+	return v;
+}
+
 Type *lb_addr_type(lbAddr const &addr) {
 	if (addr.addr.value == nullptr) {
 		return nullptr;
@@ -143,6 +152,9 @@ Type *lb_addr_type(lbAddr const &addr) {
 		GB_ASSERT(is_type_map(t));
 		return t->Map.value;
 	}
+	if (addr.kind == lbAddr_Swizzle) {
+		return addr.swizzle.type;
+	}
 	return type_deref(addr.addr.type);
 }
 LLVMTypeRef lb_addr_lb_type(lbAddr const &addr) {
@@ -193,13 +205,17 @@ lbValue lb_addr_get_ptr(lbProcedure *p, lbAddr const &addr) {
 		return final_ptr;
 	}
 
-	case lbAddr_SoaVariable: {
+	case lbAddr_SoaVariable:
 		// TODO(bill): FIX THIS HACK
 		return lb_address_from_load(p, lb_addr_load(p, addr));
-	}
 
 	case lbAddr_Context:
 		GB_PANIC("lbAddr_Context should be handled elsewhere");
+		break;
+
+	case lbAddr_Swizzle:
+		// TOOD(bill): is this good enough logic?
+		break;
 	}
 
 	return addr.addr;
@@ -354,59 +370,6 @@ void lb_addr_store(lbProcedure *p, lbAddr addr, lbValue value) {
 		LLVMBuildStore(p->builder, len.value, len_ptr.value);
 
 		return;
-
-	} else if (addr.kind == lbAddr_AtomOp_index_set) {
-		lbValue ptr = addr.addr;
-		lbValue index = addr.index_set.index;
-		Ast *node = addr.index_set.node;
-
-		ast_node(ce, CallExpr, node);
-		Type *proc_type = type_and_value_of_expr(ce->proc).type;
-		proc_type = base_type(proc_type);
-		GB_ASSERT(is_type_proc(proc_type));
-		TypeProc *pt = &proc_type->Proc;
-
-		isize arg_count = 3;
-		isize param_count = 0;
-		if (pt->params) {
-			GB_ASSERT(pt->params->kind == Type_Tuple);
-			param_count = pt->params->Tuple.variables.count;
-		}
-
-
-		auto args = array_make<lbValue>(permanent_allocator(), gb_max(arg_count, param_count));
-		args[0] = ptr;
-		args[1] = index;
-		args[2] = value;
-
-		isize arg_index = arg_count;
-		if (arg_count < param_count) {
-			lbModule *m = p->module;
-			String proc_name = {};
-			if (p->entity != nullptr) {
-				proc_name = p->entity->token.string;
-			}
-			TokenPos pos = ast_token(ce->proc).pos;
-
-			TypeTuple *param_tuple = &pt->params->Tuple;
-
-			isize end = cast(isize)param_count;
-			while (arg_index < end) {
-				Entity *e = param_tuple->variables[arg_index];
-				GB_ASSERT(e->kind == Entity_Variable);
-				args[arg_index++] = lb_handle_param_value(p, e->type, e->Variable.param_value, pos);
-			}
-		}
-
-		Entity *e = entity_from_expr(ce->proc);
-		GB_ASSERT(e != nullptr);
-		GB_ASSERT(is_type_polymorphic(e->type));
-
-		{
-			lb_emit_call(p, lb_find_procedure_value_from_entity(p->module, e), args);
-		}
-
-		return;
 	} else if (addr.kind == lbAddr_Map) {
 		lb_insert_dynamic_map_key_and_value(p, addr, addr.map.type, addr.map.key, value, p->curr_stmt);
 		return;
@@ -495,6 +458,17 @@ void lb_addr_store(lbProcedure *p, lbAddr addr, lbValue value) {
 			}
 		}
 		return;
+	} else if (addr.kind == lbAddr_Swizzle) {
+		lbValue ptr = lb_addr_get_ptr(p, addr);
+		lbValue src_ptr = lb_address_from_load_or_generate_local(p, value);
+
+		for (u8 i = 0; i < addr.swizzle.count; i++) {
+			u8 index = addr.swizzle.indices[i];
+			lbValue dst = lb_emit_array_epi(p, ptr, index);
+			lbValue src = lb_emit_array_epi(p, src_ptr, i);
+			lb_emit_store(p, dst, lb_emit_load(p, src));
+		}
+		return;
 	}
 
 	GB_ASSERT(value.value != nullptr);
@@ -771,6 +745,17 @@ lbValue lb_addr_load(lbProcedure *p, lbAddr const &addr) {
 		}
 
 		return lb_addr_load(p, res);
+	} else if (addr.kind == lbAddr_Swizzle) {
+		lbAddr res = lb_add_local_generated(p, addr.swizzle.type, false);
+		lbValue ptr = lb_addr_get_ptr(p, res);
+
+		for (u8 i = 0; i < addr.swizzle.count; i++) {
+			u8 index = addr.swizzle.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)) {
@@ -12389,6 +12374,23 @@ lbAddr lb_build_addr(lbProcedure *p, Ast *expr) {
 				GB_PANIC("Unreachable");
 			}
 
+			if (se->swizzle_count > 0) {
+				Type *array_type = base_type(type_deref(tav.type));
+				GB_ASSERT(array_type->kind == Type_Array);
+				u8 swizzle_count = se->swizzle_count;
+				u8 swizzle_indices_raw = se->swizzle_indices;
+				u8 swizzle_indices[4] = {};
+				for (u8 i = 0; i < swizzle_count; i++) {
+					u8 index = swizzle_indices_raw>>(i*2) & 3;
+					swizzle_indices[i] = index;
+				}
+				lbAddr addr = lb_build_addr(p, se->expr);
+				lbValue a = lb_addr_get_ptr(p, addr);
+
+				GB_ASSERT(is_type_array(expr->tav.type));
+				return lb_addr_swizzle(a, expr->tav.type, swizzle_count, swizzle_indices);
+			}
+
 			Selection sel = lookup_field(type, selector, false);
 			GB_ASSERT(sel.entity != nullptr);
 
@@ -12419,7 +12421,6 @@ lbAddr lb_build_addr(lbProcedure *p, Ast *expr) {
 					Type *t = base_type(type_deref(addr.addr.type));
 					GB_ASSERT(is_type_soa_struct(t));
 
-					// TODO(bill): Bounds check
 					if (addr.soa.index_expr != nullptr && (!lb_is_const(addr.soa.index) || t->Struct.soa_kind != StructSoa_Fixed)) {
 						lbValue len = lb_soa_struct_len(p, addr.addr);
 						lb_emit_bounds_check(p, ast_token(addr.soa.index_expr), addr.soa.index, len);
@@ -12436,6 +12437,10 @@ lbAddr lb_build_addr(lbProcedure *p, Ast *expr) {
 						item = lb_emit_deep_field_gep(p, item, sub_sel);
 					}
 					return lb_addr(item);
+				} else if (addr.kind == lbAddr_Swizzle) {
+					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);
diff --git a/src/llvm_backend.hpp b/src/llvm_backend.hpp
index 8f50650a8..e8d0f9daa 100644
--- a/src/llvm_backend.hpp
+++ b/src/llvm_backend.hpp
@@ -47,7 +47,7 @@ enum lbAddrKind {
 	lbAddr_RelativePointer,
 	lbAddr_RelativeSlice,
 
-	lbAddr_AtomOp_index_set,
+	lbAddr_Swizzle,
 };
 
 struct lbAddr {
@@ -73,6 +73,11 @@ struct lbAddr {
 		struct {
 			bool deref;
 		} relative;
+		struct {
+			Type *type;
+			u8 count;      // 2, 3, or 4 components
+			u8 indices[4];
+		} swizzle;
 	};
 };
 
diff --git a/src/parser.hpp b/src/parser.hpp
index 89f714aaa..1d353b04c 100644
--- a/src/parser.hpp
+++ b/src/parser.hpp
@@ -22,6 +22,8 @@ enum AddressingMode {
 	Addressing_OptionalOk  = 10, // rhs: acts like a value with an optional boolean part (for existence check)
 	Addressing_SoaVariable = 11, // Struct-Of-Arrays indexed variable
 
+	Addressing_SwizzleValue    = 12, // Swizzle indexed value
+	Addressing_SwizzleVariable = 13, // Swizzle indexed variable
 };
 
 struct TypeAndValue {
@@ -320,7 +322,12 @@ AST_KIND(_ExprBegin,  "",  bool) \
 	AST_KIND(UnaryExpr,    "unary expression",       struct { Token op; Ast *expr; }) \
 	AST_KIND(BinaryExpr,   "binary expression",      struct { Token op; Ast *left, *right; } ) \
 	AST_KIND(ParenExpr,    "parentheses expression", struct { Ast *expr; Token open, close; }) \
-	AST_KIND(SelectorExpr, "selector expression",    struct { Token token; Ast *expr, *selector; }) \
+	AST_KIND(SelectorExpr, "selector expression",    struct { \
+		Token token; \
+		Ast *expr, *selector; \
+		u8 swizzle_count; /*maximum of 4 components, if set, count >= 2*/ \
+		u8 swizzle_indices; /*2 bits per component*/ \
+	}) \
 	AST_KIND(ImplicitSelectorExpr, "implicit selector expression",    struct { Token token; Ast *selector; }) \
 	AST_KIND(SelectorCallExpr, "selector call expression",    struct { Token token; Ast *expr, *call; bool modified_call; }) \
 	AST_KIND(IndexExpr,    "index expression",       struct { Ast *expr, *index; Token open, close; }) \
diff --git a/src/types.cpp b/src/types.cpp
index 8a78e08d1..d6d4c943c 100644
--- a/src/types.cpp
+++ b/src/types.cpp
@@ -371,6 +371,8 @@ struct Selection {
 	Entity *   entity;
 	Array<i32> index;
 	bool       indirect; // Set if there was a pointer deref anywhere down the line
+	u8 swizzle_count;    // maximum components = 4
+	u8 swizzle_indices;  // 2 bits per component, representing which swizzle index
 };
 Selection empty_selection = {0};