Merge pull request #7270 from corleypc/soa-array-elem-indexing-fix

Fixes chain indexing for array elements in soa containers

src/check_expr.cpp

@@ -3004,8 +3004,13 @@ gb_internal void check_unary_expr(CheckerContext *c, Operand *o, Token op, Ast *
 			if (ast_node_expect(index_expr, Ast_IndexExpr)) {
 				ast_node(ie, IndexExpr, index_expr);
 				Type *soa_type = type_deref(type_of_expr(ie->expr));
-				GB_ASSERT(is_type_soa_struct(soa_type));
-				o->type = alloc_type_soa_pointer(soa_type);
+				if (is_type_soa_struct(soa_type)) {
+					o->type = alloc_type_soa_pointer(soa_type);
+				} else {
+					// &soa[i][j]
+					GB_ASSERT_MSG(is_type_array(soa_type), "%s", type_to_string(soa_type));
+					o->type = alloc_type_pointer(o->type);
+				}
 			} else {
 				o->type = alloc_type_pointer(o->type);
 			}
@@ -9167,7 +9172,11 @@ gb_internal bool check_set_index_data(Operand *o, Type *t, bool indirection, i64
 		if (indirection) {
 			o->mode = Addressing_Variable;
 		} else if (o->mode != Addressing_Variable &&
+		           o->mode != Addressing_SoaVariable &&
 		           o->mode != Addressing_Constant) {
+			// NOTE: an #soa element of array type keeps SoaVariable, so soa[i][j] stays an
+			// lvalue. Its components are one per lane rather than contiguous, but a single
+			// component still has a real address, the same one soa[i].y denotes.
 			o->mode = Addressing_Value;
 		}
 		o->type = t->Array.elem;

src/llvm_backend_expr.cpp

@@ -5068,6 +5068,47 @@ gb_internal void lb_build_addr_compound_lit_assign_array(lbProcedure *p, Array<l
 	}
 }
 
+// address of one component of an #soa container's array element, for soa[i][j] and for the
+// same element reached through a for-in looping variable, v[j]
+gb_internal lbAddr lb_build_addr_soa_elem_index(lbProcedure *p, Ast *expr, lbAddr const &soa_addr, Type *elem_type) {
+	ast_node(ie, IndexExpr, expr);
+	GB_ASSERT(soa_addr.kind == lbAddr_SoaVariable);
+	GB_ASSERT_MSG(elem_type->kind == Type_Array, "%s", type_to_string(elem_type));
+
+	// this is a no-op when the index is in range by construction, like in a for-in loop
+	lb_emit_soa_index_bounds_check(p, soa_addr.addr, soa_addr.soa.index, soa_addr.soa.index_expr);
+
+	i64 const component_count = elem_type->Array.count;
+
+	auto const index_tv = type_and_value_of_expr(ie->index);
+	if (index_tv.mode == Addressing_Constant && index_tv.value.kind == ExactValue_Integer) {
+		i64 component = big_int_to_i64(&index_tv.value.value_integer);
+		GB_ASSERT_MSG(0 <= component && component < component_count, "%s", expr_to_string(expr));
+		return lb_addr_soa_field_elem(lb_soa_field_elem_ptr(p, soa_addr.addr, cast(i32)component, soa_addr.soa.index));
+	}
+
+	// for non-constant index do chain select between the component pointers;
+	// an element array holds at most 4 components, so this is at most 3 compares and 3 selects (branchless), 
+	// and should be folded if j turns out constant after inlining;
+	// TODO: more efficient codegen can be done, most definitely for fixed kind, possibly for slice/dynamic,
+	// (but this works for both kinds)
+	//
+	// Note: a temp holding the whole element would be simpler but would not be an lvalue,
+	// and writes go through here
+	lbValue index = lb_emit_conv(p, lb_build_expr(p, ie->index), t_int);
+	lb_emit_bounds_check(p, ast_token(ie->index), index, lb_const_int(p->module, t_int, component_count));
+
+	lbValue ptr = lb_soa_field_elem_ptr(p, soa_addr.addr, 0, soa_addr.soa.index);
+	// lb_emit_select evaluates both arms, so every candidate address is formed;
+	// only the selected one is loaded or stored through
+	for (i64 component = 1; component < component_count; component++) {
+		lbValue candidate = lb_soa_field_elem_ptr(p, soa_addr.addr, cast(i32)component, soa_addr.soa.index);
+		lbValue is_component = lb_emit_comp(p, Token_CmpEq, index, lb_const_int(p->module, t_int, component));
+		ptr = lb_emit_select(p, is_component, candidate, ptr);
+	}
+	return lb_addr_soa_field_elem(ptr);
+}
+
 gb_internal lbAddr lb_build_addr_index_expr(lbProcedure *p, Ast *expr) {
 	ast_node(ie, IndexExpr, expr);
 
@@ -5086,10 +5127,10 @@ gb_internal lbAddr lb_build_addr_index_expr(lbProcedure *p, Ast *expr) {
 		return lb_addr_soa_variable(val, index, ie->index);
 	}
 
-	if (ie->expr->tav.mode == Addressing_SoaVariable) {
-		// SOA Structures for slices/dynamic arrays
-		GB_ASSERT_MSG(is_type_multi_pointer(type_of_expr(ie->expr)), "%s", type_to_string(type_of_expr(ie->expr)));
-
+	if (ie->expr->tav.mode == Addressing_SoaVariable && is_type_multi_pointer(type_of_expr(ie->expr))) {
+		// soa.x[i], indexing one field's multipointer;
+		// the soa element of array type, soa[i][j] carries Addressing_SoaVariable too but has
+		// no multipointer to index; it is handled in the Type_Array case below
 		lbValue field = lb_build_expr(p, ie->expr);
 		lbValue index = lb_build_expr(p, ie->index);
 
@@ -5145,8 +5186,15 @@ gb_internal lbAddr lb_build_addr_index_expr(lbProcedure *p, Ast *expr) {
 
 	switch (t->kind) {
 	case Type_Array: {
-		lbValue array = {};
-		array = lb_build_addr_ptr(p, ie->expr);
+		lbAddr array_addr = lb_build_addr(p, ie->expr);
+		// soa[i][j], or v[j] in a for-in loop, with v being the looping var over soa container
+		//
+		// keyed on the lowered lbAddr kind rather than the addressing mode,
+		// because, unlike soa[i], v is seen by the checker as Addressing_Variable
+		if (array_addr.kind == lbAddr_SoaVariable) {
+			return lb_build_addr_soa_elem_index(p, expr, array_addr, t);
+		}
+		lbValue array = lb_addr_get_ptr(p, array_addr);
 		if (deref) {
 			array = lb_emit_load(p, array);
 		}

src/llvm_backend_general.cpp

@@ -617,6 +617,12 @@ gb_internal lbAddr lb_addr_soa_variable(lbValue addr, lbValue index, Ast *index_
 }
 
 // pointer to the index element of the field_index component
+//
+// the returned pointer type depends on the soa kind (because the field types do):
+// ^T for StructSoa_Fixed (field is [N]T array), but [^]T for the slice and dynamic
+// kinds (field is the [^]T this offsets);
+// loads and stores work with either, but an lbAddr must not hold the multipointer,
+// so use lb_addr_soa_field_elem to build an lbAddr from this
 gb_internal lbValue lb_soa_field_elem_ptr(lbProcedure *p, lbValue soa_ptr, i32 field_index, lbValue index) {
 	Type *t = base_type(type_deref(soa_ptr.type));
 	GB_ASSERT_MSG(t->kind == Type_Struct && t->Struct.soa_kind != StructSoa_None, "%s", type_to_string(t));
@@ -628,6 +634,14 @@ gb_internal lbValue lb_soa_field_elem_ptr(lbProcedure *p, lbValue soa_ptr, i32 f
 	return lb_emit_ptr_offset(p, lb_emit_load(p, field), index);
 }
 
+// lbAddr over an lb_soa_field_elem_ptr pointer, retyped ^T when it came as [^]T
+gb_internal lbAddr lb_addr_soa_field_elem(lbValue ptr) {
+	if (is_type_multi_pointer(ptr.type)) {
+		ptr.type = alloc_type_multi_pointer_to_pointer(ptr.type);
+	}
+	return lb_addr(ptr);
+}
+
 // bounds check for an #soa element index
 gb_internal void lb_emit_soa_index_bounds_check(lbProcedure *p, lbValue soa_ptr, lbValue index, Ast *index_expr) {
 	if (index_expr == nullptr) {

tests/core/runtime/test_core_runtime.odin

@@ -384,6 +384,214 @@ test_soa_array_elem_swizzle :: proc(t: ^testing.T) {
 	testing.expect_value(t, fixed[1].xy, [2]u16{10, 9})
 }
 
+// chained indexing of #soa container when element type is an array -> soa[i][j]
+@(test)
+test_soa_array_elem_chained_indexing :: proc(t: ^testing.T) {
+
+	fixed: #soa[3][4]u16
+	for i in 0 ..< 3 {
+		fixed.x[i] = u16(i*10 + 0)
+		fixed.y[i] = u16(i*10 + 1)
+		fixed.z[i] = u16(i*10 + 2)
+		fixed.w[i] = u16(i*10 + 3)
+	}
+
+	// const inner index
+	testing.expect_value(t, fixed[1][0], 10)
+	testing.expect_value(t, fixed[2][3], 23)
+
+	// var inner index
+	for i in 0 ..< 3 {
+		tmp := fixed[i]
+		for j in 0 ..< 4 {
+			testing.expect_value(t, fixed[i][j], tmp[j])
+			testing.expect_value(t, fixed[i][j], u16(i*10 + j))
+		}
+	}
+
+	// index and the .x/y/z/w name must agree
+	for i in 0 ..< 3 {
+		testing.expect_value(t, fixed[i][0], fixed[i].x)
+		testing.expect_value(t, fixed[i][3], fixed[i].w)
+	}
+
+	// dynamic + slice
+	dyn := make(#soa[dynamic][4]u16, 2)
+	defer delete(dyn)
+	dyn[0] = [4]u16{10, 11, 12, 13}
+	dyn[1] = [4]u16{20, 21, 22, 23}
+
+	testing.expect_value(t, dyn[0][3], 13)
+	testing.expect_value(t, dyn[1][0], 20)
+
+	slice := dyn[:]
+	testing.expect_value(t, slice[0][3], 13)
+	testing.expect_value(t, slice[1][2], 22)
+
+	for i in 0 ..< 2 {
+		dyn_tmp := dyn[i]
+		slice_tmp := slice[i]
+		for j in 0 ..< 4 {
+			testing.expect_value(t, dyn[i][j], dyn_tmp[j])
+			testing.expect_value(t, slice[i][j], dyn_tmp[j])
+			testing.expect_value(t, dyn[i][j], slice_tmp[j])
+			testing.expect_value(t, slice[i][j], slice_tmp[j])
+		}
+	}
+
+	// soa[i][j] must equal v[j] where v is the for-in looping variable
+	// test fixed, dynamic and sliec
+	for v, i in fixed {
+		testing.expect_value(t, v[3], u16(i*10 + 3))
+		for j in 0 ..< 4 {
+			testing.expect_value(t, v[j], fixed[i][j])
+		}
+	}
+	for &v, i in fixed {
+		testing.expect_value(t, v[1], u16(i*10 + 1))
+		j := 2
+		testing.expect_value(t, v[j], u16(i*10 + j))
+	}
+	for v, i in dyn {
+		for j in 0 ..< 4 {
+			testing.expect_value(t, v[j], u16((i + 1)*10 + j))
+		}
+	}
+	for &v, i in dyn {
+		testing.expect_value(t, v[1], u16((i + 1)*10 + 1))
+		j := 3
+		testing.expect_value(t, v[j], u16((i + 1)*10 + j))
+	}	
+	for v, i in slice {
+		testing.expect_value(t, v[2], slice[i][2])
+	}
+	for &v, i in slice {
+		testing.expect_value(t, v[2], slice[i][2])
+		j := 0
+		testing.expect_value(t, v[j], u16((i + 1)*10 + j))
+	}		
+
+	// write access must work through the looping var
+	fixed_scatter: #soa[2][4]u16
+	for &v, i in fixed_scatter {
+		v[0] = u16(i)
+		for j in 1 ..< 4 {
+			v[j] = u16(i*10 + j)
+		}
+		v[3] += 5
+	}
+	for i in 0 ..< 2 {
+		testing.expect_value(t, fixed_scatter.x[i], u16(i))
+		testing.expect_value(t, fixed_scatter.y[i], u16(i*10 + 1))
+		testing.expect_value(t, fixed_scatter.z[i], u16(i*10 + 2))
+		testing.expect_value(t, fixed_scatter.w[i], u16(i*10 + 3 + 5))
+	}
+	// for fixed soa you can get fancy and select a whole .x/y/z/w lane
+	testing.expect_value(t, fixed_scatter.x, [2]u16{0, 1})
+	testing.expect_value(t, fixed_scatter.y, [2]u16{1, 11})
+	testing.expect_value(t, fixed_scatter.z, [2]u16{2, 12})
+	testing.expect_value(t, fixed_scatter.w, [2]u16{3 + 5, 13 + 5})
+
+	dyn_scatter := make(#soa[dynamic][4]u16, 2)
+	defer delete(dyn_scatter)
+	for &v, i in dyn_scatter {
+		v[0] = u16(i + 1)
+		for j in 1 ..< 4 {
+			v[j] = u16((i + 1)*10 + j)
+		}
+		v[3] += 5
+	}
+	for i in 0 ..< 2 {
+		testing.expect_value(t, dyn_scatter.x[i], u16(i + 1))
+		testing.expect_value(t, dyn_scatter.y[i], u16((i + 1)*10 + 1))
+		testing.expect_value(t, dyn_scatter.z[i], u16((i + 1)*10 + 2))
+		testing.expect_value(t, dyn_scatter.w[i], u16((i + 1)*10 + 3 + 5))
+	}
+
+	slice_scatter := dyn_scatter[:]
+	for &v in slice_scatter {
+		v[1] += 100
+	}
+	testing.expect_value(t, dyn_scatter.y[0], 111)
+	testing.expect_value(t, dyn_scatter.y[1], 121)
+	testing.expect_value(t, dyn_scatter.x[0], 1)
+
+	// soa[i][j] writes
+	fixed_write: #soa[3][4]u16
+	fixed_write[1][0] = 5
+	k := 2
+	fixed_write[1][k] = 6
+	fixed_write[1][1] += 7
+	fixed_write[1][0], fixed_write[1][3] = fixed_write[1][3], fixed_write[1][0]
+	testing.expect_value(t, fixed_write.x, [3]u16{0, 0, 0})
+	testing.expect_value(t, fixed_write.y, [3]u16{0, 7, 0})
+	testing.expect_value(t, fixed_write.z, [3]u16{0, 6, 0})
+	testing.expect_value(t, fixed_write.w, [3]u16{0, 5, 0})
+
+	dyn_write := make(#soa[dynamic][4]u16, 2)
+	defer delete(dyn_write)
+	dyn_write[0][3] = 41
+	dyn_write[1][k] = 42
+	slice_write := dyn_write[:]
+	slice_write[0][1] = 43
+	slice_write[0][k] = 44
+	testing.expect_value(t, dyn_write.w[0], 41)
+	testing.expect_value(t, dyn_write.z[1], 42)
+	testing.expect_value(t, dyn_write.y[0], 43)
+	testing.expect_value(t, dyn_write.z[0], 44)
+	testing.expect_value(t, dyn_write.x[0], 0)
+
+	// a single component has a real address
+	testing.expect_value(t, &fixed_write[1][2], &fixed_write.z[1])
+	pw := &fixed_write[1][2]
+	pw^ = 60
+	testing.expect_value(t, fixed_write.z[1], 60)
+	for j in 0 ..< 4 {
+		p := &fixed_write[1][j]
+		p^ = u16(70 + j)
+	}
+	testing.expect_value(t, fixed_write.x, [3]u16{0, 70, 0})
+	testing.expect_value(t, fixed_write.y, [3]u16{0, 71, 0})
+	testing.expect_value(t, fixed_write.z, [3]u16{0, 72, 0})
+	testing.expect_value(t, fixed_write.w, [3]u16{0, 73, 0})
+
+	for j in 0 ..< 4 {
+		p := &dyn_write[0][j]
+		p^ = u16(80 + j)
+	}
+	testing.expect_value(t, dyn_write.x[0], 80)
+	testing.expect_value(t, dyn_write.w[0], 83)
+	for j in 0 ..< 4 {
+		p := &slice_write[1][j]
+		p^ = u16(90 + j)
+	}
+	testing.expect_value(t, dyn_write.x[1], 90)
+	testing.expect_value(t, dyn_write.w[1], 93)
+
+	// multi dim array element type, only the outer index is scattered
+	nested: #soa[2][3][2]u16
+	nested[1].x = {1, 3}
+	nested[1].z = {7, 11}
+	testing.expect_value(t, nested[1][0][1], 3)
+	testing.expect_value(t, nested[1][2][0], 7)	
+
+	nested[1][0][1] = 8
+	nested[1][2][0] = 9
+	testing.expect_value(t, nested[1][0], [2]u16{1, 8})
+	testing.expect_value(t, nested[1][2], [2]u16{9, 11})
+
+	for j in 0 ..< 3 {
+		nested[1][j][0] = u16(50 + j)
+	}
+	for j in 0 ..< 3 {
+		testing.expect_value(t, nested[1][j][0], u16(50 + j))
+	}
+	testing.expect_value(t, nested.x[1][0], 50)
+	testing.expect_value(t, nested.y[1][0], 51)
+	testing.expect_value(t, nested.z[1][0], 52)
+	testing.expect_value(t, nested[0][0], [2]u16{0, 0})
+}
+
 // "using" on an #soa for-in looping variable
 @(test)
 test_soa_for_in_using :: proc(t: ^testing.T) {