`for &e, i in array` and `for k, &v in map` (alternative to passing the iterable by pointer)

1 files changed, 29 insertions, 4 deletions

diff --git a/src/check_stmt.cpp b/src/check_stmt.cpp
index cf6f998e5..30eef4f25 100644
--- a/src/check_stmt.cpp
+++ b/src/check_stmt.cpp
@@ -1469,12 +1469,15 @@ gb_internal void check_range_stmt(CheckerContext *ctx, Ast *node, u32 mod_flags)
 
 	Ast *expr = unparen_expr(rs->expr);
 
+	bool is_possibly_addressable = true;
 	isize max_val_count = 2;
 	if (is_ast_range(expr)) {
 		ast_node(ie, BinaryExpr, expr);
 		Operand x = {};
 		Operand y = {};
 
+		is_possibly_addressable = false;
+
 		bool ok = check_range(ctx, expr, true, &x, &y, nullptr);
 		if (!ok) {
 			goto skip_expr_range_stmt;
@@ -1497,6 +1500,8 @@ gb_internal void check_range_stmt(CheckerContext *ctx, Ast *node, u32 mod_flags)
 				gb_string_free(t);
 				goto skip_expr_range_stmt;
 			} else {
+				is_possibly_addressable = false;
+
 				if (is_reverse) {
 					error(node, "#reverse for is not supported for enum types");
 				}
@@ -1510,7 +1515,8 @@ gb_internal void check_range_stmt(CheckerContext *ctx, Ast *node, u32 mod_flags)
 			Type *t = base_type(type_deref(operand.type));
 			switch (t->kind) {
 			case Type_Basic:
-				if (is_type_string(t) && t->Basic.kind != Basic_cstring) {
+				if (t->Basic.kind == Basic_string) {
+					is_possibly_addressable = false;
 					array_add(&vals, t_rune);
 					array_add(&vals, t_int);
 					if (is_reverse) {
@@ -1529,6 +1535,7 @@ gb_internal void check_range_stmt(CheckerContext *ctx, Ast *node, u32 mod_flags)
 
 			case Type_Array:
 				if (is_ptr) use_by_reference_for_value = true;
+				if (!is_ptr) is_possibly_addressable = operand.mode == Addressing_Variable;
 				array_add(&vals, t->Array.elem);
 				array_add(&vals, t_int);
 				break;
@@ -1575,6 +1582,8 @@ gb_internal void check_range_stmt(CheckerContext *ctx, Ast *node, u32 mod_flags)
 						array_add(&vals, e->type);
 					}
 
+					is_possibly_addressable = false;
+
 					if (rs->vals.count > 1 && rs->vals[1] != nullptr && count < 3) {
 						gbString s = type_to_string(t);
 						error(operand.expr, "Expected a 3-valued expression on the rhs, got (%s)", s);
@@ -1644,8 +1653,13 @@ gb_internal void check_range_stmt(CheckerContext *ctx, Ast *node, u32 mod_flags)
 		}
 		Ast * name = lhs[i];
 		Type *type = rhs[i];
-
 		Entity *entity = nullptr;
+
+		bool is_addressed = false;
+		if (name->kind == Ast_UnaryExpr && name->UnaryExpr.op.kind == Token_And) {
+			is_addressed = true;
+			name = name->UnaryExpr.expr;
+		}
 		if (name->kind == Ast_Ident) {
 			Token token = name->Ident.token;
 			String str = token.string;
@@ -1659,7 +1673,16 @@ gb_internal void check_range_stmt(CheckerContext *ctx, Ast *node, u32 mod_flags)
 				entity->flags |= EntityFlag_ForValue;
 				entity->flags |= EntityFlag_Value;
 				entity->identifier = name;
-				if (i == addressable_index && use_by_reference_for_value) {
+				entity->Variable.for_loop_parent_type = type_of_expr(expr);
+
+				if (is_addressed) {
+					if (is_possibly_addressable && i == addressable_index) {
+						entity->flags &= ~EntityFlag_Value;
+					} else {
+						char const *idx_name = is_map ? "key" : "index";
+						error(token, "The %s variable '%.*s' cannot be made addressable", idx_name, LIT(str));
+					}
+				} else if (i == addressable_index && use_by_reference_for_value) {
 					entity->flags &= ~EntityFlag_Value;
 				}
 				if (is_soa) {
@@ -1678,7 +1701,9 @@ gb_internal void check_range_stmt(CheckerContext *ctx, Ast *node, u32 mod_flags)
 				entity = found;
 			}
 		} else {
-			error(name, "A variable declaration must be an identifier");
+			gbString s = expr_to_string(lhs[i]);
+			error(name, "A variable declaration must be an identifier, got %s", s);
+			gb_string_free(s);
 		}
 
 		if (entity == nullptr) {