Allow `switch &v in ...` work to be consistent with `for &e in ...`

3 files changed, 53 insertions, 5 deletions

diff --git a/src/check_expr.cpp b/src/check_expr.cpp
index 532555b0a..e0e78b441 100644
--- a/src/check_expr.cpp
+++ b/src/check_expr.cpp
@@ -2246,6 +2246,33 @@ gb_internal bool check_is_not_addressable(CheckerContext *c, Operand *o) {
 	return o->mode != Addressing_Variable && o->mode != Addressing_SoaVariable;
 }
 
+gb_internal void check_old_for_or_switch_value_usage(Ast *expr) {
+	Entity *e = entity_of_node(expr);
+	if (e != nullptr && (e->flags & EntityFlag_OldForOrSwitchValue) != 0) {
+		GB_ASSERT(e->kind == Entity_Variable);
+
+		begin_error_block();
+		defer (end_error_block());
+
+		if ((e->flags & EntityFlag_ForValue) != 0) {
+			Type *parent_type = type_deref(e->Variable.for_loop_parent_type);
+
+			error(expr, "Assuming a for-in defined value is addressable as the iterable is passed by value has been disallowed with '-strict-style'.");
+
+			if (is_type_map(parent_type)) {
+				error_line("\tSuggestion: Prefer doing 'for key, &%.*s in ...'\n", LIT(e->token.string));
+			} else {
+				error_line("\tSuggestion: Prefer doing 'for &%.*s in ...'\n", LIT(e->token.string));
+			}
+		} else {
+			GB_ASSERT((e->flags & EntityFlag_SwitchValue) != 0);
+
+			error(expr, "Assuming a switch-in defined value is addressable as the iterable is passed by value has been disallowed with '-strict-style'.");
+			error_line("\tSuggestion: Prefer doing 'switch &%.*s in ...'\n", LIT(e->token.string));
+		}
+	}
+}
+
 gb_internal void check_unary_expr(CheckerContext *c, Operand *o, Token op, Ast *node) {
 	switch (op.kind) {
 	case Token_And: { // Pointer address
@@ -2255,7 +2282,7 @@ gb_internal void check_unary_expr(CheckerContext *c, Operand *o, Token op, Ast *
 				gbString str = expr_to_string(ue->expr);
 				defer (gb_string_free(str));
 
-				Entity *e = entity_of_node(o->expr);
+				Entity *e = entity_of_node(ue->expr);
 				if (e != nullptr && (e->flags & EntityFlag_Param) != 0) {
 					error(op, "Cannot take the pointer address of '%s' which is a procedure parameter", str);
 				} else {
@@ -2306,6 +2333,11 @@ gb_internal void check_unary_expr(CheckerContext *c, Operand *o, Token op, Ast *
 				o->type = alloc_type_pointer(o->type);
 			}
 		} else {
+			if (build_context.strict_style && ast_node_expect(node, Ast_UnaryExpr)) {
+				ast_node(ue, UnaryExpr, node);
+				check_old_for_or_switch_value_usage(ue->expr);
+			}
+
 			o->type = alloc_type_pointer(o->type);
 		}
 
diff --git a/src/check_stmt.cpp b/src/check_stmt.cpp
index 30eef4f25..964ff1842 100644
--- a/src/check_stmt.cpp
+++ b/src/check_stmt.cpp
@@ -417,6 +417,7 @@ gb_internal Type *check_assignment_variable(CheckerContext *ctx, Operand *lhs, O
 		return nullptr;
 
 	case Addressing_Variable:
+		check_old_for_or_switch_value_usage(lhs->expr);
 		break;
 
 	case Addressing_MapIndex: {
@@ -1141,8 +1142,14 @@ gb_internal void check_type_switch_stmt(CheckerContext *ctx, Ast *node, u32 mod_
 		syntax_error(as_token, "Expected 1 expression after 'in'");
 		return;
 	}
+	bool is_addressed = false;
+
 	Ast *lhs = as->lhs[0];
 	Ast *rhs = as->rhs[0];
+	if (lhs->kind == Ast_UnaryExpr && lhs->UnaryExpr.op.kind == Token_And) {
+		is_addressed = true;
+		lhs = lhs->UnaryExpr.expr;
+	}
 
 	check_expr(ctx, &x, rhs);
 	check_assignment(ctx, &x, nullptr, str_lit("type switch expression"));
@@ -1281,12 +1288,15 @@ gb_internal void check_type_switch_stmt(CheckerContext *ctx, Ast *node, u32 mod_
 			}
 		}
 
-		bool is_reference = false;
+		bool is_reference = is_addressed;
+		bool old_style = false;
 
-		if (is_ptr &&
+		if (!is_reference &&
+		    is_ptr &&
 		    cc->list.count == 1 &&
 		    case_type != nullptr) {
 			is_reference = true;
+			old_style = true;
 		}
 
 		if (cc->list.count > 1 || saw_nil) {
@@ -1305,9 +1315,12 @@ gb_internal void check_type_switch_stmt(CheckerContext *ctx, Ast *node, u32 mod_
 		{
 			Entity *tag_var = alloc_entity_variable(ctx->scope, lhs->Ident.token, case_type, EntityState_Resolved);
 			tag_var->flags |= EntityFlag_Used;
+			tag_var->flags |= EntityFlag_SwitchValue;
 			if (!is_reference) {
 				tag_var->flags |= EntityFlag_Value;
-				tag_var->flags |= EntityFlag_SwitchValue;
+			}
+			if (old_style) {
+				tag_var->flags |= EntityFlag_OldForOrSwitchValue;
 			}
 			add_entity(ctx, ctx->scope, lhs, tag_var);
 			add_entity_use(ctx, lhs, tag_var);
@@ -1683,6 +1696,7 @@ gb_internal void check_range_stmt(CheckerContext *ctx, Ast *node, u32 mod_flags)
 						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_OldForOrSwitchValue;
 					entity->flags &= ~EntityFlag_Value;
 				}
 				if (is_soa) {
diff --git a/src/entity.cpp b/src/entity.cpp
index 2dcb76482..649dd900d 100644
--- a/src/entity.cpp
+++ b/src/entity.cpp
@@ -84,7 +84,9 @@ enum EntityFlag : u64 {
 	EntityFlag_CustomLinkage_LinkOnce = 1ull<<44,
 	
 	EntityFlag_Require = 1ull<<50,
-	EntityFlag_ByPtr = 1ull<<51, // enforce parameter is passed by pointer
+	EntityFlag_ByPtr   = 1ull<<51, // enforce parameter is passed by pointer
+
+	EntityFlag_OldForOrSwitchValue = 1ull<<52,
 
 	EntityFlag_Overridden    = 1ull<<63,
 };