Correct and improve type inference for swizzling expressions

3 files changed, 35 insertions, 31 deletions

diff --git a/examples/demo/demo.odin b/examples/demo/demo.odin
index 6431f1d7a..2590fdc01 100644
--- a/examples/demo/demo.odin
+++ b/examples/demo/demo.odin
@@ -8,6 +8,9 @@ import "core:time"
 import "core:reflect"
 import "core:runtime"
 import "intrinsics"
+import "core:encoding/json"
+
+_ :: json;
 
 /*
 	The Odin programming language is fast, concise, readable, pragmatic and open sourced.
@@ -1203,6 +1206,12 @@ array_programming :: proc() {
 			return i - j;
 		}
 
+		cross_shorter :: proc(a, b: Vector3) -> Vector3 {
+			i := a.yzx * b.zxy;
+			j := a.zxy * b.yzx;
+			return i - j;
+		}
+
 		blah :: proc(a: Vector3) -> f32 {
 			return a.x + a.y + a.z;
 		}
diff --git a/src/check_builtin.cpp b/src/check_builtin.cpp
index c0ba40503..88b643d10 100644
--- a/src/check_builtin.cpp
+++ b/src/check_builtin.cpp
@@ -718,7 +718,8 @@ bool check_builtin_procedure(CheckerContext *c, Operand *operand, Ast *call, i32
 
 	case BuiltinProc_swizzle: {
 		// swizzle :: proc(v: [N]T, ..int) -> [M]T
-		Type *type = base_type(operand->type);
+		Type *original_type = operand->type;
+		Type *type = base_type(original_type);
 		i64 max_count = 0;
 		Type *elem_type = nullptr;
 
@@ -775,9 +776,6 @@ bool check_builtin_procedure(CheckerContext *c, Operand *operand, Ast *call, i32
 			return false;
 		}
 
-		if (arg_count < max_count) {
-			operand->type = alloc_type_array(elem_type, arg_count);
-		}
 		if (type->kind == Type_Array) {
 			if (operand->mode == Addressing_Variable) {
 				operand->mode = Addressing_SwizzleVariable;
@@ -788,21 +786,7 @@ bool check_builtin_procedure(CheckerContext *c, Operand *operand, Ast *call, i32
 			operand->mode = Addressing_Value;
 		}
 
-		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;
-
+		operand->type = determine_swizzle_array_type(original_type, type_hint, arg_count);
 		break;
 	}
 
diff --git a/src/check_expr.cpp b/src/check_expr.cpp
index b2cd18430..a440c3540 100644
--- a/src/check_expr.cpp
+++ b/src/check_expr.cpp
@@ -3632,6 +3632,27 @@ void check_did_you_mean_scope(String const &name, Scope *scope) {
 	check_did_you_mean_print(&d);
 }
 
+Type *determine_swizzle_array_type(Type *original_type, Type *type_hint, isize new_count) {
+	Type *array_type = base_type(type_deref(original_type));
+	GB_ASSERT(array_type->kind == Type_Array);
+	Type *elem_type = array_type->Array.elem;
+
+	Type *swizzle_array_type = nullptr;
+	Type *bth = base_type(type_deref(type_hint));
+	if (bth != nullptr && bth->kind == Type_Array &&
+	    bth->Array.count == new_count &&
+	    are_types_identical(bth->Array.elem, elem_type)) {
+		swizzle_array_type = type_hint;
+	} else {
+		i64 max_count = array_type->Array.count;
+		if (new_count == max_count) {
+			swizzle_array_type = original_type;
+		} else {
+			swizzle_array_type = alloc_type_array(elem_type, new_count);
+		}
+	}
+	return swizzle_array_type;
+}
 
 
 Entity *check_selector(CheckerContext *c, Operand *operand, Ast *node, Type *type_hint) {
@@ -3842,21 +3863,11 @@ Entity *check_selector(CheckerContext *c, Operand *operand, Ast *node, Type *typ
 			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 *original_type = operand->type;
 
-			Type *swizzle_array_type = nullptr;
-			Type *bth = base_type(type_hint);
-			if (bth != nullptr && bth->kind == Type_Array &&
-			    bth->Array.count == index_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, index_count);
-			}
 			AddressingMode prev_mode = operand->mode;
 			operand->mode = Addressing_SwizzleValue;
-			operand->type = swizzle_array_type;
+			operand->type = determine_swizzle_array_type(original_type, type_hint, index_count);
 			operand->expr = node;
 
 			switch (prev_mode) {