Support by-reference semantics in `for value_ref, index in &some_array` and `for key, value_ref in &some_map`

2 files changed, 41 insertions, 9 deletions

diff --git a/src/check_stmt.cpp b/src/check_stmt.cpp
index e9b6869c8..be2f6b06a 100644
--- a/src/check_stmt.cpp
+++ b/src/check_stmt.cpp
@@ -1484,6 +1484,7 @@ void check_stmt_internal(CheckerContext *ctx, Ast *node, u32 flags) {
 		Entity *entities[2] = {};
 		isize entity_count = 0;
 		bool is_map = false;
+		bool use_by_reference_for_value = false;
 
 		Ast *expr = unparen_expr(rs->expr);
 
@@ -1529,26 +1530,31 @@ void check_stmt_internal(CheckerContext *ctx, Ast *node, u32 flags) {
 					break;
 
 				case Type_EnumeratedArray:
+					if (is_ptr) use_by_reference_for_value = true;
 					val0 = t->EnumeratedArray.elem;
 					val1 = t->EnumeratedArray.index;
 					break;
 
 				case Type_Array:
+					if (is_ptr) use_by_reference_for_value = true;
 					val0 = t->Array.elem;
 					val1 = t_int;
 					break;
 
 				case Type_DynamicArray:
+					if (is_ptr) use_by_reference_for_value = true;
 					val0 = t->DynamicArray.elem;
 					val1 = t_int;
 					break;
 
 				case Type_Slice:
+					if (is_ptr) use_by_reference_for_value = true;
 					val0 = t->Slice.elem;
 					val1 = t_int;
 					break;
 
 				case Type_Map:
+					if (is_ptr) use_by_reference_for_value = true;
 					is_map = true;
 					val0 = t->Map.key;
 					val1 = t->Map.value;
@@ -1628,6 +1634,14 @@ void check_stmt_internal(CheckerContext *ctx, Ast *node, u32 flags) {
 				if (found == nullptr) {
 					entity = alloc_entity_variable(ctx->scope, token, type, EntityState_Resolved);
 					entity->flags |= EntityFlag_Value;
+					if (use_by_reference_for_value) {
+						if (i == 0 && !is_map) {
+							entity->flags &= ~EntityFlag_Value;
+						} else if (i == 1 && is_map) {
+							entity->flags &= ~EntityFlag_Value;
+						}
+					}
+
 					add_entity_definition(&ctx->checker->info, name, entity);
 				} else {
 					TokenPos pos = found->token.pos;
diff --git a/src/ir.cpp b/src/ir.cpp
index f51538a2e..4736f2ba1 100644
--- a/src/ir.cpp
+++ b/src/ir.cpp
@@ -9542,6 +9542,29 @@ void ir_store_type_case_implicit(irProcedure *proc, Ast *clause, irValue *value)
 	}
 }
 
+irAddr ir_store_range_stmt_val(irProcedure *proc, Ast *stmt_val, irValue *value) {
+	Entity *e = entity_of_node(stmt_val);
+	if (e == nullptr) {
+		return {};
+	}
+
+	if ((e->flags & EntityFlag_Value) == 0) {
+		if (value->kind == irValue_Instr) {
+			if (value->Instr.kind == irInstr_Load) {
+				irValue *ptr = value->Instr.Load.address;
+				ir_module_add_value(proc->module, e, ptr);
+				return ir_addr(ptr);
+			}
+		}
+	}
+
+	// by value
+	irAddr addr = ir_addr(ir_add_local(proc, e, nullptr, false));
+	GB_ASSERT(are_types_identical(ir_type(value), e->type));
+	ir_addr_store(proc, addr, value);
+	return addr;
+}
+
 void ir_type_case_body(irProcedure *proc, Ast *label, Ast *clause, irBlock *body, irBlock *done) {
 	ast_node(cc, CaseClause, clause);
 
@@ -10076,17 +10099,12 @@ void ir_build_stmt_internal(irProcedure *proc, Ast *node) {
 		}
 
 
-		irAddr val0_addr = {};
-		irAddr val1_addr = {};
-		if (val0_type) val0_addr = ir_build_addr(proc, rs->val0);
-		if (val1_type) val1_addr = ir_build_addr(proc, rs->val1);
-
 		if (is_map) {
-			if (val0_type) ir_addr_store(proc, val0_addr, key);
-			if (val1_type) ir_addr_store(proc, val1_addr, val);
+			if (val0_type) ir_store_range_stmt_val(proc, rs->val0, key);
+			if (val1_type) ir_store_range_stmt_val(proc, rs->val1, val);
 		} else {
-			if (val0_type) ir_addr_store(proc, val0_addr, val);
-			if (val1_type) ir_addr_store(proc, val1_addr, key);
+			if (val0_type) ir_store_range_stmt_val(proc, rs->val0, val);
+			if (val1_type) ir_store_range_stmt_val(proc, rs->val1, key);
 		}
 
 		ir_push_target_list(proc, rs->label, done, loop, nullptr);