From 2820bbc269bb989052163544f2764bc4287254e4 Mon Sep 17 00:00:00 2001
From: gingerBill <bill@gingerbill.org>
Date: Fri, 5 Jan 2024 13:38:30 +0000
Subject: [PATCH] Add `@(entry_point_only)` for procedures

---
 src/check_decl.cpp | 1 +
 src/check_expr.cpp | 8 ++++++++
 src/checker.cpp    | 3 +++
 src/checker.hpp    | 1 +
 src/entity.cpp     | 1 +
 5 files changed, 14 insertions(+)

diff --git a/src/check_decl.cpp b/src/check_decl.cpp
index 71b897a84..85ba4230a 100644
--- a/src/check_decl.cpp
+++ b/src/check_decl.cpp
@@ -908,6 +908,7 @@ gb_internal void check_proc_decl(CheckerContext *ctx, Entity *e, DeclInfo *d) {
 		break;
 	}
 
+	e->Procedure.entry_point_only = ac.entry_point_only;
 	e->Procedure.is_export = ac.is_export;
 	e->deprecated_message = ac.deprecated_message;
 	e->warning_message = ac.warning_message;
diff --git a/src/check_expr.cpp b/src/check_expr.cpp
index 71accfb81..14e3bc0de 100644
--- a/src/check_expr.cpp
+++ b/src/check_expr.cpp
@@ -7195,6 +7195,14 @@ gb_internal ExprKind check_call_expr(CheckerContext *c, Operand *operand, Ast *c
 			}
 		}
 		add_entity_use(c, operand->expr, initial_entity);
+
+		if (initial_entity->Procedure.entry_point_only) {
+			if (c->curr_proc_decl && c->curr_proc_decl->entity == c->info->entry_point) {
+				// Okay
+			} else {
+				error(operand->expr, "Procedures with the attribute '@(entry_point_only)' can only be called directly from the user-level entry point procedure");
+			}
+		}
 	}
 
 	if (operand->mode != Addressing_ProcGroup) {
diff --git a/src/checker.cpp b/src/checker.cpp
index 79328d648..ac885222d 100644
--- a/src/checker.cpp
+++ b/src/checker.cpp
@@ -3413,6 +3413,9 @@ gb_internal DECL_ATTRIBUTE_PROC(proc_decl_attribute) {
 			error(elem, "Expected a string value for '%.*s'", LIT(name));
 		}
 		return true;
+	} else if (name == "entry_point_only") {
+		ac->entry_point_only = true;
+		return true;
 	}
 	return false;
 }
diff --git a/src/checker.hpp b/src/checker.hpp
index a6a5f6788..7c399e50f 100644
--- a/src/checker.hpp
+++ b/src/checker.hpp
@@ -123,6 +123,7 @@ struct AttributeContext {
 	bool    init                : 1;
 	bool    fini                : 1;
 	bool    set_cold            : 1;
+	bool    entry_point_only    : 1;
 	u32 optimization_mode; // ProcedureOptimizationMode
 	i64 foreign_import_priority_index;
 	String extra_linker_flags;
diff --git a/src/entity.cpp b/src/entity.cpp
index ce27da3f2..d0b3cf139 100644
--- a/src/entity.cpp
+++ b/src/entity.cpp
@@ -250,6 +250,7 @@ struct Entity {
 			bool    is_export                  : 1;
 			bool    generated_from_polymorphic : 1;
 			bool    target_feature_disabled    : 1;
+			bool    entry_point_only           : 1;
 			String  target_feature;
 		} Procedure;
 		struct {