Potentially fix a race condition with parapoly types (related to #3328)

4 files changed, 104 insertions, 97 deletions

diff --git a/src/check_expr.cpp b/src/check_expr.cpp
index fd10374c1..44d65e376 100644
--- a/src/check_expr.cpp
+++ b/src/check_expr.cpp
@@ -79,7 +79,6 @@ gb_internal Type *   check_type_expr                (CheckerContext *c, Ast *exp
 gb_internal Type *   make_optional_ok_type          (Type *value, bool typed=true);
 gb_internal Entity * check_selector                 (CheckerContext *c, Operand *operand, Ast *node, Type *type_hint);
 gb_internal Entity * check_ident                    (CheckerContext *c, Operand *o, Ast *n, Type *named_type, Type *type_hint, bool allow_import_name);
-gb_internal Entity * find_polymorphic_record_entity (CheckerContext *c, Type *original_type, isize param_count, Array<Operand> const &ordered_operands, bool *failure);
 gb_internal void     check_not_tuple                (CheckerContext *c, Operand *operand);
 gb_internal void     convert_to_typed               (CheckerContext *c, Operand *operand, Type *target_type);
 gb_internal gbString expr_to_string                 (Ast *expression);
@@ -121,6 +120,8 @@ gb_internal isize get_procedure_param_count_excluding_defaults(Type *pt, isize *
 
 gb_internal bool is_expr_inferred_fixed_array(Ast *type_expr);
 
+gb_internal Entity *find_polymorphic_record_entity(GenTypesData *found_gen_types, isize param_count, Array<Operand> const &ordered_operands);
+
 enum LoadDirectiveResult {
 	LoadDirective_Success  = 0,
 	LoadDirective_Error    = 1,
@@ -7171,8 +7172,12 @@ gb_internal CallArgumentError check_polymorphic_record_type(CheckerContext *c, O
 	}
 
 	{
-		bool failure = false;
-		Entity *found_entity = find_polymorphic_record_entity(c, original_type, param_count, ordered_operands, &failure);
+		GenTypesData *found_gen_types = ensure_polymorphic_record_entity_has_gen_types(c, original_type);
+
+		mutex_lock(&found_gen_types->mutex);
+		defer (mutex_unlock(&found_gen_types->mutex));
+		Entity *found_entity = find_polymorphic_record_entity(found_gen_types, param_count, ordered_operands);
+
 		if (found_entity) {
 			operand->mode = Addressing_Type;
 			operand->type = found_entity->type;
diff --git a/src/check_type.cpp b/src/check_type.cpp
index 0b9042905..e22d4b62b 100644
--- a/src/check_type.cpp
+++ b/src/check_type.cpp
@@ -260,84 +260,21 @@ gb_internal bool check_custom_align(CheckerContext *ctx, Ast *node, i64 *align_,
 }
 
 
-gb_internal Entity *find_polymorphic_record_entity(CheckerContext *ctx, Type *original_type, isize param_count, Array<Operand> const &ordered_operands, bool *failure) {
-	rw_mutex_shared_lock(&ctx->info->gen_types_mutex); // @@global
-
-	auto *found_gen_types = map_get(&ctx->info->gen_types, original_type);
-	if (found_gen_types == nullptr) {
-		rw_mutex_shared_unlock(&ctx->info->gen_types_mutex); // @@global
-		return nullptr;
-	}
-
-	rw_mutex_shared_lock(&found_gen_types->mutex); // @@local
-	defer (rw_mutex_shared_unlock(&found_gen_types->mutex)); // @@local
-
-	rw_mutex_shared_unlock(&ctx->info->gen_types_mutex); // @@global
-
-	for (Entity *e : found_gen_types->types) {
-		Type *t = base_type(e->type);
-		TypeTuple *tuple = nullptr;
-		switch (t->kind) {
-		case Type_Struct:
-			if (t->Struct.polymorphic_params) {
-				tuple = &t->Struct.polymorphic_params->Tuple;
-			}
-			break;
-		case Type_Union:
-			if (t->Union.polymorphic_params) {
-				tuple = &t->Union.polymorphic_params->Tuple;
-			}
-			break;
-		}
-		GB_ASSERT_MSG(tuple != nullptr, "%s :: %s", type_to_string(e->type), type_to_string(t));
-		GB_ASSERT(param_count == tuple->variables.count);
-
-		bool skip = false;
-
-		for (isize j = 0; j < param_count; j++) {
-			Entity *p = tuple->variables[j];
-			Operand o = {};
-			if (j < ordered_operands.count) {
-				o = ordered_operands[j];
-			}
-			if (o.expr == nullptr) {
-				continue;
-			}
-			Entity *oe = entity_of_node(o.expr);
-			if (p == oe) {
-				// NOTE(bill): This is the same type, make sure that it will be be same thing and use that
-				// Saves on a lot of checking too below
-				continue;
-			}
-
-			if (p->kind == Entity_TypeName) {
-				if (is_type_polymorphic(o.type)) {
-					// NOTE(bill): Do not add polymorphic version to the gen_types
-					skip = true;
-					break;
-				}
-				if (!are_types_identical(o.type, p->type)) {
-					skip = true;
-					break;
-				}
-			} else if (p->kind == Entity_Constant) {
-				if (!compare_exact_values(Token_CmpEq, o.value, p->Constant.value)) {
-					skip = true;
-					break;
-				}
-				if (!are_types_identical(o.type, p->type)) {
-					skip = true;
-					break;
-				}
-			} else {
-				GB_PANIC("Unknown entity kind");
-			}
-		}
-		if (!skip) {
-			return e;
-		}
+gb_internal GenTypesData *ensure_polymorphic_record_entity_has_gen_types(CheckerContext *ctx, Type *original_type) {
+	mutex_lock(&ctx->info->gen_types_mutex); // @@global
+
+	GenTypesData *found_gen_types = nullptr;
+	auto *found_gen_types_ptr = map_get(&ctx->info->gen_types, original_type);
+	if (found_gen_types_ptr == nullptr) {
+		GenTypesData *gen_types = gb_alloc_item(permanent_allocator(), GenTypesData);
+		gen_types->types = array_make<Entity *>(heap_allocator());
+		map_set(&ctx->info->gen_types, original_type, gen_types);
+		found_gen_types_ptr = map_get(&ctx->info->gen_types, original_type);
 	}
-	return nullptr;
+	found_gen_types = *found_gen_types_ptr;
+	GB_ASSERT(found_gen_types != nullptr);
+	mutex_unlock(&ctx->info->gen_types_mutex); // @@global
+	return found_gen_types;
 }
 
 
@@ -367,19 +304,16 @@ gb_internal void add_polymorphic_record_entity(CheckerContext *ctx, Ast *node, T
 	// TODO(bill): Is this even correct? Or should the metadata be copied?
 	e->TypeName.objc_metadata = original_type->Named.type_name->TypeName.objc_metadata;
 
-	rw_mutex_lock(&ctx->info->gen_types_mutex);
-	auto *found_gen_types = map_get(&ctx->info->gen_types, original_type);
-	if (found_gen_types) {
-		rw_mutex_lock(&found_gen_types->mutex);
-		array_add(&found_gen_types->types, e);
-		rw_mutex_unlock(&found_gen_types->mutex);
-	} else {
-		GenTypesData gen_types = {};
-		gen_types.types = array_make<Entity *>(heap_allocator());
-		array_add(&gen_types.types, e);
-		map_set(&ctx->info->gen_types, original_type, gen_types);
+	auto *found_gen_types = ensure_polymorphic_record_entity_has_gen_types(ctx, original_type);
+	mutex_lock(&found_gen_types->mutex);
+	defer (mutex_unlock(&found_gen_types->mutex));
+
+	for (Entity *prev : found_gen_types->types) {
+		if (prev == e) {
+			return;
+		}
 	}
-	rw_mutex_unlock(&ctx->info->gen_types_mutex);
+	array_add(&found_gen_types->types, e);
 }
 
 
@@ -612,6 +546,73 @@ gb_internal bool check_record_poly_operand_specialization(CheckerContext *ctx, T
 	return true;
 }
 
+gb_internal Entity *find_polymorphic_record_entity(GenTypesData *found_gen_types, isize param_count, Array<Operand> const &ordered_operands) {
+	for (Entity *e : found_gen_types->types) {
+		Type *t = base_type(e->type);
+		TypeTuple *tuple = nullptr;
+		switch (t->kind) {
+		case Type_Struct:
+			if (t->Struct.polymorphic_params) {
+				tuple = &t->Struct.polymorphic_params->Tuple;
+			}
+			break;
+		case Type_Union:
+			if (t->Union.polymorphic_params) {
+				tuple = &t->Union.polymorphic_params->Tuple;
+			}
+			break;
+		}
+		GB_ASSERT_MSG(tuple != nullptr, "%s :: %s", type_to_string(e->type), type_to_string(t));
+		GB_ASSERT(param_count == tuple->variables.count);
+
+		bool skip = false;
+
+		for (isize j = 0; j < param_count; j++) {
+			Entity *p = tuple->variables[j];
+			Operand o = {};
+			if (j < ordered_operands.count) {
+				o = ordered_operands[j];
+			}
+			if (o.expr == nullptr) {
+				continue;
+			}
+			Entity *oe = entity_of_node(o.expr);
+			if (p == oe) {
+				// NOTE(bill): This is the same type, make sure that it will be be same thing and use that
+				// Saves on a lot of checking too below
+				continue;
+			}
+
+			if (p->kind == Entity_TypeName) {
+				if (is_type_polymorphic(o.type)) {
+					// NOTE(bill): Do not add polymorphic version to the gen_types
+					skip = true;
+					break;
+				}
+				if (!are_types_identical(o.type, p->type)) {
+					skip = true;
+					break;
+				}
+			} else if (p->kind == Entity_Constant) {
+				if (!compare_exact_values(Token_CmpEq, o.value, p->Constant.value)) {
+					skip = true;
+					break;
+				}
+				if (!are_types_identical(o.type, p->type)) {
+					skip = true;
+					break;
+				}
+			} else {
+				GB_PANIC("Unknown entity kind");
+			}
+		}
+		if (!skip) {
+			return e;
+		}
+	}
+	return nullptr;
+};
+
 
 gb_internal void check_struct_type(CheckerContext *ctx, Type *struct_type, Ast *node, Array<Operand> *poly_operands, Type *named_type, Type *original_type_for_poly) {
 	GB_ASSERT(is_type_struct(struct_type));
diff --git a/src/checker.hpp b/src/checker.hpp
index eea99578e..e0dc54a87 100644
--- a/src/checker.hpp
+++ b/src/checker.hpp
@@ -360,7 +360,7 @@ struct GenProcsData {
 
 struct GenTypesData {
 	Array<Entity *> types;
-	RwMutex         mutex;
+	RecursiveMutex  mutex;
 };
 
 // CheckerInfo stores all the symbol information for a type-checked program
@@ -400,8 +400,8 @@ struct CheckerInfo {
 
 	RecursiveMutex lazy_mutex; // Mutex required for lazy type checking of specific files
 
-	RwMutex       gen_types_mutex;
-	PtrMap<Type *, GenTypesData > gen_types;
+	BlockingMutex                  gen_types_mutex;
+	PtrMap<Type *, GenTypesData *> gen_types;
 
 	BlockingMutex type_info_mutex; // NOT recursive
 	Array<Type *> type_info_types;
@@ -560,3 +560,6 @@ gb_internal void init_core_context(Checker *c);
 gb_internal void init_mem_allocator(Checker *c);
 
 gb_internal void add_untyped_expressions(CheckerInfo *cinfo, UntypedExprInfoMap *untyped);
+
+
+gb_internal GenTypesData *ensure_polymorphic_record_entity_has_gen_types(CheckerContext *ctx, Type *original_type);
\ No newline at end of file
diff --git a/src/threading.cpp b/src/threading.cpp
index a469435d2..d9538f66e 100644
--- a/src/threading.cpp
+++ b/src/threading.cpp
@@ -122,8 +122,6 @@ gb_internal void wait_signal_set(Wait_Signal *ws) {
 	futex_broadcast(&ws->futex);
 }
 
-
-
 struct MutexGuard {
 	MutexGuard()                   = delete;
 	MutexGuard(MutexGuard const &) = delete;