Fix race condition from `add_entity_use` due to Entity.identifier

6 files changed, 17 insertions, 15 deletions

diff --git a/src/check_decl.cpp b/src/check_decl.cpp
index c8bacb9b2..4cce25097 100644
--- a/src/check_decl.cpp
+++ b/src/check_decl.cpp
@@ -350,12 +350,12 @@ void override_entity_in_scope(Entity *original_entity, Entity *new_entity) {
 	original_entity->type = new_entity->type;
 	original_entity->aliased_of = new_entity;
 
-	if (original_entity->identifier == nullptr) {
-		original_entity->identifier = new_entity->identifier;
-	}
-	if (original_entity->identifier != nullptr &&
-	    original_entity->identifier->kind == Ast_Ident) {
-		original_entity->identifier->Ident.entity = new_entity;
+	Ast *empty_ident = nullptr;
+	original_entity->identifier.compare_exchange_strong(empty_ident, new_entity->identifier);
+
+	if (original_entity->identifier.load() != nullptr &&
+	    original_entity->identifier.load()->kind == Ast_Ident) {
+		original_entity->identifier.load()->Ident.entity = new_entity;
 	}
 
 	// IMPORTANT NOTE(bill, 2021-04-10): copy only the variants
diff --git a/src/check_type.cpp b/src/check_type.cpp
index d3bcb5bb6..867b4d7b2 100644
--- a/src/check_type.cpp
+++ b/src/check_type.cpp
@@ -1565,7 +1565,7 @@ Type *check_get_params(CheckerContext *ctx, Scope *scope, Ast *_params, bool *is
 						if (is_type_proc(op.type)) {
 							Entity *proc_entity = entity_from_expr(op.expr);
 							valid = proc_entity != nullptr;
-							poly_const = exact_value_procedure(proc_entity->identifier ? proc_entity->identifier : op.expr);
+							poly_const = exact_value_procedure(proc_entity->identifier.load() ? proc_entity->identifier.load() : op.expr);
 						}
 						if (!valid) {
 							if (op.mode == Addressing_Constant) {
diff --git a/src/checker.cpp b/src/checker.cpp
index 4ae5f5b1c..cd023998c 100644
--- a/src/checker.cpp
+++ b/src/checker.cpp
@@ -1297,9 +1297,9 @@ void add_entity_use(CheckerContext *c, Ast *identifier, Entity *entity) {
 		if (identifier->kind != Ast_Ident) {
 			return;
 		}
-		if (entity->identifier == nullptr) {
-			entity->identifier = identifier;
-		}
+		Ast *empty_ident = nullptr;
+		entity->identifier.compare_exchange_strong(empty_ident, identifier);
+
 		identifier->Ident.entity = entity;
 
 		if (c->info->allow_identifier_uses) {
diff --git a/src/entity.cpp b/src/entity.cpp
index fb65bc3a8..9f3b8f84d 100644
--- a/src/entity.cpp
+++ b/src/entity.cpp
@@ -120,7 +120,7 @@ struct Entity {
 	Token       token;
 	Scope *     scope;
 	Type *      type;
-	Ast *       identifier; // Can be nullptr
+	std::atomic<Ast *> identifier; // Can be nullptr
 	DeclInfo *  decl_info;
 	DeclInfo *  parent_proc_decl; // nullptr if in file/global scope
 	AstFile *   file;
diff --git a/src/llvm_backend.cpp b/src/llvm_backend.cpp
index 05b10ffd7..d50a512c7 100644
--- a/src/llvm_backend.cpp
+++ b/src/llvm_backend.cpp
@@ -1204,8 +1204,9 @@ void lb_generate_code(lbGenerator *gen) {
 
 			LLVMBool is_optimized = build_context.optimization_level > 0;
 			AstFile *init_file = m->info->init_package->files[0];
-			if (m->info->entry_point && m->info->entry_point->identifier && m->info->entry_point->identifier->file) {
-				init_file = m->info->entry_point->identifier->file;
+			Ast *ident = m->info->entry_point->identifier.load();
+			if (m->info->entry_point && ident && ident->file) {
+				init_file = ident->file;
 			}
 
 			LLVMBool split_debug_inlining = false;
diff --git a/src/llvm_backend_proc.cpp b/src/llvm_backend_proc.cpp
index 9516f2e3d..366fba780 100644
--- a/src/llvm_backend_proc.cpp
+++ b/src/llvm_backend_proc.cpp
@@ -211,11 +211,12 @@ lbProcedure *lb_create_procedure(lbModule *m, Entity *entity, bool ignore_body)
 		scope = p->module->debug_compile_unit;
 		type = lb_debug_type_internal_proc(m, bt);
 
+		Ast *ident = entity->identifier.load();
 		if (entity->file != nullptr) {
 			file = lb_get_llvm_metadata(m, entity->file);
 			scope = file;
-		} else if (entity->identifier != nullptr && entity->identifier->file != nullptr) {
-			file = lb_get_llvm_metadata(m, entity->identifier->file);
+		} else if (ident != nullptr && ident->file != nullptr) {
+			file = lb_get_llvm_metadata(m, ident->file);
 			scope = file;
 		} else if (entity->scope != nullptr) {
 			file = lb_get_llvm_metadata(m, entity->scope->file);