Support `conj` on array and matrix types

2 files changed, 62 insertions, 28 deletions

diff --git a/src/check_builtin.cpp b/src/check_builtin.cpp
index b60509c03..7dc4784f8 100644
--- a/src/check_builtin.cpp
+++ b/src/check_builtin.cpp
@@ -1266,7 +1266,10 @@ bool check_builtin_procedure(CheckerContext *c, Operand *operand, Ast *call, i32
 	case BuiltinProc_conj: {
 		// conj :: proc(x: type) -> type
 		Operand *x = operand;
-		if (is_type_complex(x->type)) {
+		Type *t = x->type;
+		Type *elem = core_array_type(t);
+		
+		if (is_type_complex(t)) {
 			if (x->mode == Addressing_Constant) {
 				ExactValue v = exact_value_to_complex(x->value);
 				f64 r = v.value_complex->real;
@@ -1276,7 +1279,7 @@ bool check_builtin_procedure(CheckerContext *c, Operand *operand, Ast *call, i32
 			} else {
 				x->mode = Addressing_Value;
 			}
-		} else if (is_type_quaternion(x->type)) {
+		} else if (is_type_quaternion(t)) {
 			if (x->mode == Addressing_Constant) {
 				ExactValue v = exact_value_to_quaternion(x->value);
 				f64 r = +v.value_quaternion->real;
@@ -1288,7 +1291,11 @@ bool check_builtin_procedure(CheckerContext *c, Operand *operand, Ast *call, i32
 			} else {
 				x->mode = Addressing_Value;
 			}
-		} else {
+		} else if (is_type_array_like(t) && (is_type_complex(elem) || is_type_quaternion(elem))) {
+			x->mode = Addressing_Value;
+		} else if (is_type_matrix(t) && (is_type_complex(elem) || is_type_quaternion(elem))) {
+			x->mode = Addressing_Value;
+		}else {
 			gbString s = type_to_string(x->type);
 			error(call, "Expected a complex or quaternion, got '%s'", s);
 			gb_string_free(s);
diff --git a/src/llvm_backend_proc.cpp b/src/llvm_backend_proc.cpp
index 8686b3262..72ba3982c 100644
--- a/src/llvm_backend_proc.cpp
+++ b/src/llvm_backend_proc.cpp
@@ -725,6 +725,57 @@ lbValue lb_emit_runtime_call(lbProcedure *p, char const *c_name, Array<lbValue>
 	return lb_emit_call(p, proc, args);
 }
 
+lbValue lb_emit_conjugate(lbProcedure *p, lbValue val, Type *type) {
+	lbValue res = {};
+	Type *t = val.type;
+	if (is_type_complex(t)) {
+		res = lb_addr_get_ptr(p, lb_add_local_generated(p, type, false));
+		lbValue real = lb_emit_struct_ev(p, val, 0);
+		lbValue imag = lb_emit_struct_ev(p, val, 1);
+		imag = lb_emit_unary_arith(p, Token_Sub, imag, imag.type);
+		lb_emit_store(p, lb_emit_struct_ep(p, res, 0), real);
+		lb_emit_store(p, lb_emit_struct_ep(p, res, 1), imag);
+	} else if (is_type_quaternion(t)) {
+		// @QuaternionLayout
+		res = lb_addr_get_ptr(p, lb_add_local_generated(p, type, false));
+		lbValue real = lb_emit_struct_ev(p, val, 3);
+		lbValue imag = lb_emit_struct_ev(p, val, 0);
+		lbValue jmag = lb_emit_struct_ev(p, val, 1);
+		lbValue kmag = lb_emit_struct_ev(p, val, 2);
+		imag = lb_emit_unary_arith(p, Token_Sub, imag, imag.type);
+		jmag = lb_emit_unary_arith(p, Token_Sub, jmag, jmag.type);
+		kmag = lb_emit_unary_arith(p, Token_Sub, kmag, kmag.type);
+		lb_emit_store(p, lb_emit_struct_ep(p, res, 3), real);
+		lb_emit_store(p, lb_emit_struct_ep(p, res, 0), imag);
+		lb_emit_store(p, lb_emit_struct_ep(p, res, 1), jmag);
+		lb_emit_store(p, lb_emit_struct_ep(p, res, 2), kmag);
+	} else if (is_type_array_like(t)) {
+		res = lb_addr_get_ptr(p, lb_add_local_generated(p, type, true));
+		Type *elem_type = base_array_type(t);
+		i64 count = get_array_type_count(t);
+		for (i64 i = 0; i < count; i++) {
+			lbValue dst = lb_emit_array_epi(p, res, i);
+			lbValue elem = lb_emit_struct_ev(p, val, cast(i32)i);
+			elem = lb_emit_conjugate(p, elem, elem_type);
+			lb_emit_store(p, dst, elem);
+		}
+	} else if (is_type_matrix(t)) {
+		Type *mt = base_type(t);
+		GB_ASSERT(mt->kind == Type_Matrix);
+		Type *elem_type = mt->Matrix.elem;
+		res = lb_addr_get_ptr(p, lb_add_local_generated(p, type, true));
+		for (i64 j = 0; j < mt->Matrix.column_count; j++) {
+			for (i64 i = 0; i < mt->Matrix.row_count; i++) {
+				lbValue dst = lb_emit_matrix_epi(p, res, i, j);
+				lbValue elem = lb_emit_matrix_ev(p, val, i, j);
+				elem = lb_emit_conjugate(p, elem, elem_type);
+				lb_emit_store(p, dst, elem);
+			}
+		}
+	}
+	return lb_emit_load(p, res);
+}
+
 lbValue lb_emit_call(lbProcedure *p, lbValue value, Array<lbValue> const &args, ProcInlining inlining, bool use_copy_elision_hint) {
 	lbModule *m = p->module;
 
@@ -1117,31 +1168,7 @@ lbValue lb_build_builtin_proc(lbProcedure *p, Ast *expr, TypeAndValue const &tv,
 
 	case BuiltinProc_conj: {
 		lbValue val = lb_build_expr(p, ce->args[0]);
-		lbValue res = {};
-		Type *t = val.type;
-		if (is_type_complex(t)) {
-			res = lb_addr_get_ptr(p, lb_add_local_generated(p, tv.type, false));
-			lbValue real = lb_emit_struct_ev(p, val, 0);
-			lbValue imag = lb_emit_struct_ev(p, val, 1);
-			imag = lb_emit_unary_arith(p, Token_Sub, imag, imag.type);
-			lb_emit_store(p, lb_emit_struct_ep(p, res, 0), real);
-			lb_emit_store(p, lb_emit_struct_ep(p, res, 1), imag);
-		} else if (is_type_quaternion(t)) {
-			// @QuaternionLayout
-			res = lb_addr_get_ptr(p, lb_add_local_generated(p, tv.type, false));
-			lbValue real = lb_emit_struct_ev(p, val, 3);
-			lbValue imag = lb_emit_struct_ev(p, val, 0);
-			lbValue jmag = lb_emit_struct_ev(p, val, 1);
-			lbValue kmag = lb_emit_struct_ev(p, val, 2);
-			imag = lb_emit_unary_arith(p, Token_Sub, imag, imag.type);
-			jmag = lb_emit_unary_arith(p, Token_Sub, jmag, jmag.type);
-			kmag = lb_emit_unary_arith(p, Token_Sub, kmag, kmag.type);
-			lb_emit_store(p, lb_emit_struct_ep(p, res, 3), real);
-			lb_emit_store(p, lb_emit_struct_ep(p, res, 0), imag);
-			lb_emit_store(p, lb_emit_struct_ep(p, res, 1), jmag);
-			lb_emit_store(p, lb_emit_struct_ep(p, res, 2), kmag);
-		}
-		return lb_emit_load(p, res);
+		return lb_emit_conjugate(p, val, tv.type);
 	}
 
 	case BuiltinProc_expand_to_tuple: {