Add #+escape-analysis and #+no-escape-analysis file tags, and keep the syntactic check of unsafe returns where the escape analysis is disabled

src/check_escape.cpp

@@ -167,6 +167,7 @@ struct EscapeGraph {
 	Array<i32>                group_of;
 	Array<i32>                group_offsets;  // group -> its procedures, as `members[group_offsets[gi]..<group_offsets[gi+1]]`
 	Array<i32>                members;        // every group a group may call has a lower index
+	Array<bool>               skipped;        // per procedure, neither in a file with the escape analysis nor called from one
 	Slice<std::atomic<i32> >  pending;        // per group its calls to procedures of other groups which are not analysed yet
 	Slice<std::atomic<i32> >  cursors;        // per procedure while `targets` is filled
 	bool                      threaded;
@@ -3480,9 +3481,13 @@ gb_internal Slice<EscapeFlow> escape_analyse(EscapeGraph *g, i32 v, Array<Escape
 	ea.group    = g->group_of[v];
 	ea.reports  = reports;
 
+	// in a file without it, which only gives what flows through it to those calling it
+	bool enabled = ast_file_escape_analysis(body->file());
+	ea.muted = !enabled;
+
 	u64 vet_flags = ast_file_vet_flags(body->file());
-	ea.nil_deref    = (vet_flags & VetFlag_NilDeref) != 0;
-	ea.uninit       = (vet_flags & VetFlag_Uninitialized) != 0;
+	ea.nil_deref    = enabled && (vet_flags & VetFlag_NilDeref) != 0;
+	ea.uninit       = enabled && (vet_flags & VetFlag_Uninitialized) != 0;
 	ea.nil_uses     = array_make<EscapeNilUse>(temporary_allocator(), 0, 0);
 	ea.nil_escaped  = array_make<Entity *>(temporary_allocator(), 0, 0);
 	ea.uninit_decls = array_make<Entity *>(temporary_allocator(), 0, 0);
@@ -3498,7 +3503,7 @@ gb_internal Slice<EscapeFlow> escape_analyse(EscapeGraph *g, i32 v, Array<Escape
 	    !build_context.no_entry_point &&
 	    build_context.command_kind != Command_test) {
 		Entity *e = pi->decl->entity.load();
-		ea.muted = e != nullptr && e == global_checker_ptr.load(std::memory_order_relaxed)->info.entry_point;
+		ea.muted |= e != nullptr && e == global_checker_ptr.load(std::memory_order_relaxed)->info.entry_point;
 	}
 
 	ErrorInstantiations prev_instantiations = global_error_context.instantiations;
@@ -3596,6 +3601,9 @@ gb_internal void escape_analyse_group(EscapeGraph *g, i32 gi) {
 	i32 end   = g->group_offsets[gi+1];
 
 	i32 first = g->members[start];
+	if (g->skipped[first]) {
+		return;
+	}
 	bool recursive = end - start > 1;
 	for (i32 i = g->offsets[first]; i < g->offsets[first+1]; i++) {
 		recursive |= g->targets[i] == first;
@@ -3861,6 +3869,7 @@ gb_internal void check_escapes(Checker *c) {
 		array_free(&g->group_of);
 		array_free(&g->group_offsets);
 		array_free(&g->members);
+		array_free(&g->skipped);
 		slice_free(&g->pending, heap_allocator());
 		slice_free(&g->cursors, heap_allocator());
 		*g = {};
@@ -3893,6 +3902,29 @@ gb_internal void check_escapes(Checker *c) {
 		array_clear(&slot.array);
 	}
 
+	// when it is disabled but for some files, only their procedures and what those may call are analysed
+	array_init(&g->skipped, heap_allocator(), count);
+	if (build_context.no_escape_analysis) {
+		auto stack = array_make<i32>(temporary_allocator(), 0, count);
+		for (i32 v = 0; v < count; v++) {
+			Ast *body = g->procs[v]->body;
+			g->skipped[v] = body == nullptr || !ast_file_escape_analysis(body->file());
+			if (!g->skipped[v]) {
+				array_add(&stack, v);
+			}
+		}
+		while (stack.count > 0) {
+			i32 v = array_pop(&stack);
+			for (i32 i = g->offsets[v]; i < g->offsets[v+1]; i++) {
+				i32 w = g->targets[i];
+				if (g->skipped[w]) {
+					g->skipped[w] = false;
+					array_add(&stack, w);
+				}
+			}
+		}
+	}
+
 	// only needed to release the groups when threaded, and found meanwhile
 	TaskGroup callers = {};
 	if (g->threaded) {

src/check_expr.cpp

@@ -9562,7 +9562,7 @@ gb_internal ExprKind check_call_expr(CheckerContext *c, Operand *operand, Ast *c
 				c->decl->defer_used += 1;
 			}
 		}
-		if (c->curr_proc_decl != nullptr && !build_context.no_escape_analysis) {
+		if (c->curr_proc_decl != nullptr && escape_analysis_in_use()) {
 			per_thread_array_add(&c->info->checked_calls_queue, CheckedCall{c->curr_proc_decl, callee});
 		}
 	}

src/check_stmt.cpp

@@ -2633,6 +2633,80 @@ gb_internal void check_if_stmt(CheckerContext *ctx, Ast *node, u32 mod_flags) {
 	check_close_scope(ctx);
 }
 
+// returning stack memory made by the returned expression itself, where the escape analysis is disabled
+gb_internal void check_unsafe_return(Operand const &o, Type *type, Ast *expr) {
+	auto const unsafe_return_error = [](Operand const &o, char const *msg, Type *extra_type=nullptr) {
+		gbString s = expr_to_string(o.expr);
+		if (extra_type) {
+			gbString t = type_to_string(extra_type);
+			error(o.expr, "It is unsafe to return %s ('%s') of type ('%s') from a procedure, as it uses the current stack frame's memory", msg, s, t);
+			gb_string_free(t);
+		} else {
+			error(o.expr, "It is unsafe to return %s ('%s') from a procedure, as it uses the current stack frame's memory", msg, s);
+		}
+		gb_string_free(s);
+	};
+
+	if (type == nullptr || expr == nullptr) {
+		return;
+	}
+
+	if (expr->kind == Ast_CompoundLit && is_type_slice(type)) {
+		ast_node(cl, CompoundLit, expr);
+		if (cl->elems.count == 0) {
+			return;
+		}
+		unsafe_return_error(o, "a compound literal of a slice");
+	} else if (expr->kind == Ast_UnaryExpr && expr->UnaryExpr.op.kind == Token_And) {
+		Ast *x = unparen_expr(expr->UnaryExpr.expr);
+		Entity *e = entity_of_node(x);
+		if (is_entity_local_variable(e)) {
+			unsafe_return_error(o, "the address of a local variable");
+		} else if (x->kind == Ast_CompoundLit) {
+			unsafe_return_error(o, "the address of a compound literal");
+		} else if (x->kind == Ast_IndexExpr) {
+			Entity *f = entity_of_node(x->IndexExpr.expr);
+			if (f && (is_type_array_like(f->type) || is_type_matrix(f->type))) {
+				if (is_entity_local_variable(f)) {
+					unsafe_return_error(o, "the address of an indexed variable", f->type);
+				}
+			}
+		} else if (x->kind == Ast_MatrixIndexExpr) {
+			Entity *f = entity_of_node(x->MatrixIndexExpr.expr);
+			if (f && (is_type_matrix(f->type) && is_entity_local_variable(f))) {
+				unsafe_return_error(o, "the address of an indexed variable", f->type);
+			}
+		}
+	} else if (expr->kind == Ast_SliceExpr) {
+		Ast *x = unparen_expr(expr->SliceExpr.expr);
+		Entity *e = entity_of_node(x);
+		if (is_entity_local_variable(e) && is_type_array(e->type)) {
+			unsafe_return_error(o, "a slice of a local variable");
+		} else if (x->kind == Ast_CompoundLit) {
+			unsafe_return_error(o, "a slice of a compound literal");
+		}
+	} else if (o.mode == Addressing_Constant && is_type_slice(type)) {
+		if (is_load_directive_call(o.expr)) {
+			return;
+		}
+
+		ERROR_BLOCK();
+		unsafe_return_error(o, "a compound literal of a slice");
+		error_line("\tNote: A constant slice value will use the memory of the current stack frame\n");
+	} else if (expr->kind == Ast_CompoundLit) {
+		ast_node(cl, CompoundLit, expr);
+		for (Ast *elem : cl->elems) {
+			if (elem->kind == Ast_FieldValue) {
+				ast_node(fv, FieldValue, elem);
+				Entity *e = entity_of_node(fv->field);
+				if (e != nullptr) {
+					check_unsafe_return(o, e->type, fv->value);
+				}
+			}
+		}
+	}
+}
+
 gb_internal void check_return_stmt(CheckerContext *ctx, Ast *node) {
 	ast_node(rs, ReturnStmt, node);
 
@@ -2692,6 +2766,28 @@ gb_internal void check_return_stmt(CheckerContext *ctx, Ast *node) {
 			}
 		}
 	}
+
+	if (ast_file_escape_analysis(node->file())) {
+		return;
+	}
+	for (Operand &o : operands) {
+		if (o.expr == nullptr) {
+			continue;
+		}
+		Ast *expr = unparen_expr(o.expr);
+		while (expr->kind == Ast_CallExpr && expr->CallExpr.proc->tav.mode == Addressing_Type) {
+			if (expr->CallExpr.args.count != 1) {
+				break;
+			}
+			Ast *arg = expr->CallExpr.args[0];
+			if (arg->kind == Ast_FieldValue || !are_types_identical(arg->tav.type, expr->tav.type)) {
+				break;
+			}
+			expr = unparen_expr(arg);
+		}
+
+		check_unsafe_return(o, o.type, expr);
+	}
 }
 
 gb_internal void check_for_stmt(CheckerContext *ctx, Ast *node, u32 mod_flags) {

src/checker.cpp

@@ -6852,7 +6852,7 @@ gb_internal bool check_proc_info(Checker *c, ProcInfo *pi, UntypedExprInfoMap *u
 
 	if (body_was_checked) {
 		pi->decl->proc_info = pi;
-		if (!build_context.no_escape_analysis) {
+		if (escape_analysis_in_use()) {
 			per_thread_array_add(&c->info.checked_bodies_queue, pi);
 		}
 		pi->decl->proc_checked_state.store(ProcCheckedState_Checked);
@@ -8074,7 +8074,7 @@ gb_internal void check_parsed_files(Checker *c) {
 
 	debugf("Total Procedure Bodies Checked: %td\n", total_bodies_checked.load(std::memory_order_relaxed));
 
-	if (!build_context.no_escape_analysis) {
+	if (escape_analysis_in_use()) {
 		TIME_SECTION("check escapes");
 		check_escapes(c);
 	}

src/main.cpp

@@ -3269,7 +3269,8 @@ gb_internal int print_show_help(String const arg0, String command, String option
 
 	if (check) {
 		if (print_flag("-no-escape-analysis")) {
-			print_usage_line(2, "Disables the escape analysis of stack memory, including the error for returning the address of a local.");
+			print_usage_line(2, "Disables the escape analysis of stack memory, except in files with '#+escape-analysis'.");
+			print_usage_line(2, "Where it is disabled, by this or by '#+no-escape-analysis', only returning the address of a local or similar is an error.");
 			print_usage_line(2, "Cannot be used with -vet-nil-deref or -vet-uninitialized.");
 		}
 	}

src/parser.cpp

@@ -91,6 +91,26 @@ gb_internal u64 ast_file_vet_flags(AstFile *f) {
 	return 0;
 }
 
+// whether any file has `#+escape-analysis`, so that it runs for them when `-no-escape-analysis` disables it otherwise
+gb_global std::atomic<bool> global_escape_analysis_tagged;
+
+gb_internal bool ast_file_escape_analysis(AstFile *f) {
+	if (f == nullptr) {
+		return !build_context.no_escape_analysis;
+	}
+	if (f->flags & AstFile_EscapeAnalysis) {
+		return true;
+	}
+	if (f->flags & AstFile_NoEscapeAnalysis) {
+		return false;
+	}
+	return !build_context.no_escape_analysis;
+}
+
+gb_internal bool escape_analysis_in_use(void) {
+	return !build_context.no_escape_analysis || global_escape_analysis_tagged.load(std::memory_order_relaxed);
+}
+
 gb_internal bool ast_file_vet_style(AstFile *f) {
 	return (ast_file_vet_flags(f) & VetFlag_Style) != 0;
 }
@@ -7484,6 +7504,13 @@ gb_internal bool parse_file_tag(const String &lc, const Token &tok, AstFile *f)
 		}
 	} else if (lc == "no-instrumentation") {
 		f->flags |= AstFile_NoInstrumentation;
+	} else if (lc == "escape-analysis") {
+		f->flags |= AstFile_EscapeAnalysis;
+		f->flags &= ~AstFile_NoEscapeAnalysis;
+		global_escape_analysis_tagged.store(true, std::memory_order_relaxed);
+	} else if (lc == "no-escape-analysis") {
+		f->flags |= AstFile_NoEscapeAnalysis;
+		f->flags &= ~AstFile_EscapeAnalysis;
 	} else {
 		syntax_error(tok, "Unknown tag '%.*s'", LIT(lc));
 	}

src/parser.hpp

@@ -96,6 +96,9 @@ enum AstFileFlag : u32 {
 	AstFile_IsLazy    = 1<<4,
 
 	AstFile_NoInstrumentation = 1<<5,
+
+	AstFile_EscapeAnalysis   = 1<<6, // `#+escape-analysis`
+	AstFile_NoEscapeAnalysis = 1<<7, // `#+no-escape-analysis`
 };
 
 enum AstDelayQueueKind {