Allow `x :: y when cond else proc(...){...}`

3 files changed, 189 insertions, 79 deletions

diff --git a/src/check_decl.cpp b/src/check_decl.cpp
index 3c4a4b3de..0f48c28e3 100644
--- a/src/check_decl.cpp
+++ b/src/check_decl.cpp
@@ -155,6 +155,154 @@ gb_internal void check_init_variables(CheckerContext *ctx, Entity **lhs, isize l
 	}
 }
 
+
+gb_internal void override_entity_in_scope(Entity *original_entity, Entity *new_entity) {
+	// NOTE(bill): The original_entity's scope may not be same scope that it was inserted into
+	// e.g. file entity inserted into its package scope
+	String original_name = original_entity->token.string;
+	Scope *found_scope = nullptr;
+	Entity *found_entity = nullptr;
+	scope_lookup_parent(original_entity->scope, original_name, &found_scope, &found_entity);
+	if (found_scope == nullptr) {
+		return;
+	}
+	rw_mutex_lock(&found_scope->mutex);
+	defer (rw_mutex_unlock(&found_scope->mutex));
+
+	// IMPORTANT NOTE(bill, 2021-04-10): Overriding behaviour was flawed in that the
+	// original entity was still used check checked, but the checking was only
+	// relying on "constant" data such as the Entity.type and Entity.Constant.value
+	//
+	// Therefore two things can be done: the type can be assigned to state that it
+	// has been "evaluated" and the variant data can be copied across
+
+	string_map_set(&found_scope->elements, original_name, new_entity);
+
+	original_entity->flags |= EntityFlag_Overridden;
+	original_entity->type = new_entity->type;
+	original_entity->aliased_of = 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
+	// This is most likely NEVER required, but it does not at all hurt to keep
+	isize offset = cast(u8 *)&original_entity->Dummy.start - cast(u8 *)original_entity;
+	isize size = gb_size_of(*original_entity) - offset;
+	gb_memmove(cast(u8 *)original_entity, cast(u8 *)new_entity, size);
+}
+
+
+gb_internal void check_proc_decl(CheckerContext *ctx, Entity *e, DeclInfo *d);
+
+gb_internal bool check_try_override_const_decl(CheckerContext *ctx, Entity *e, Entity *entity, Ast *init, Type *named_type) {
+	if (entity == nullptr) {
+	retry_proc_lit:;
+		init = unparen_expr(init);
+		if (init == nullptr) {
+			return false;
+		}
+		if (init->kind == Ast_TernaryWhenExpr) {
+			ast_node(we, TernaryWhenExpr, init);
+			if (we->cond == nullptr) {
+				return false;
+			}
+			if (we->cond->tav.value.kind != ExactValue_Bool) {
+				return false;
+			}
+			init = we->cond->tav.value.value_bool ? we->x : we->y;
+			goto retry_proc_lit;
+		} if (init->kind == Ast_ProcLit) {
+			// NOTE(bill, 2024-07-04): Override as a procedure entity because this could be within a `when` statement
+			e->kind = Entity_Procedure;
+			e->type = nullptr;
+			DeclInfo *d = decl_info_of_entity(e);
+			d->proc_lit = init;
+			check_proc_decl(ctx, e, d);
+			return true;
+		}
+
+		return false;
+	}
+	switch (entity->kind) {
+	case Entity_TypeName:
+		// @TypeAliasingProblem
+		// NOTE(bill, 2022-02-03): This is used to solve the problem caused by type aliases
+		// being "confused" as constants
+		//
+		//         A :: B
+		//         C :: proc "c" (^A)
+		//         B :: struct {x: C}
+		//
+		//     A gets evaluated first, and then checks B.
+		//     B then checks C.
+		//     C then tries to check A which is unresolved but thought to be a constant.
+		//     Therefore within C's check, A errs as "not a type".
+		//
+		// This is because a const declaration may or may not be a type and this cannot
+		// be determined from a syntactical standpoint.
+		// This check allows the compiler to override the entity to be checked as a type.
+		//
+		// There is no problem if B is prefixed with the `#type` helper enforcing at
+		// both a syntax and semantic level that B must be a type.
+		//
+		//         A :: #type B
+		//
+		// This approach is not fool proof and can fail in case such as:
+		//
+		//         X :: type_of(x)
+		//         X :: Foo(int).Type
+		//
+		// Since even these kind of declarations may cause weird checking cycles.
+		// For the time being, these are going to be treated as an unfortunate error
+		// until there is a proper delaying system to try declaration again if they
+		// have failed.
+
+		e->kind = Entity_TypeName;
+		check_type_decl(ctx, e, init, named_type);
+		return true;
+	case Entity_Builtin:
+		if (e->type != nullptr) {
+			return false;
+		}
+		e->kind = Entity_Builtin;
+		e->Builtin.id = entity->Builtin.id;
+		e->type = t_invalid;
+		return true;
+	case Entity_ProcGroup:
+		// NOTE(bill, 2020-06-10): It is better to just clone the contents than overriding the entity in the scope
+		// Thank goodness I made entities a tagged union to allow for this implace patching
+		e->kind = Entity_ProcGroup;
+		e->ProcGroup.entities = array_clone(heap_allocator(), entity->ProcGroup.entities);
+		return true;
+	}
+
+	if (e->type != nullptr && entity->type != nullptr) {
+		Operand x = {};
+		x.type = entity->type;
+		x.mode = Addressing_Variable;
+		if (!check_is_assignable_to(ctx, &x, e->type)) {
+			return false;
+		}
+	}
+
+	// NOTE(bill): Override aliased entity
+	switch (entity->kind) {
+	case Entity_ProcGroup:
+	case Entity_Procedure:
+	case Entity_LibraryName:
+	case Entity_ImportName:
+		override_entity_in_scope(e, entity);
+		return true;
+	}
+	return false;
+}
+
 gb_internal void check_init_constant(CheckerContext *ctx, Entity *e, Operand *operand) {
 	if (operand->mode == Addressing_Invalid ||
 		operand->type == t_invalid ||
@@ -166,6 +314,13 @@ gb_internal void check_init_constant(CheckerContext *ctx, Entity *e, Operand *op
 	}
 
 	if (operand->mode != Addressing_Constant) {
+		Entity *entity = entity_of_node(operand->expr);
+		if (check_try_override_const_decl(ctx, e, entity, operand->expr, nullptr)) {
+			return;
+		}
+	}
+
+	if (operand->mode != Addressing_Constant) {
 		gbString str = expr_to_string(operand->expr);
 		error(operand->expr, "'%s' is not a compile-time known constant", str);
 		gb_string_free(str);
@@ -373,49 +528,6 @@ gb_internal void check_type_decl(CheckerContext *ctx, Entity *e, Ast *init_expr,
 }
 
 
-gb_internal void override_entity_in_scope(Entity *original_entity, Entity *new_entity) {
-	// NOTE(bill): The original_entity's scope may not be same scope that it was inserted into
-	// e.g. file entity inserted into its package scope
-	String original_name = original_entity->token.string;
-	Scope *found_scope = nullptr;
-	Entity *found_entity = nullptr;
-	scope_lookup_parent(original_entity->scope, original_name, &found_scope, &found_entity);
-	if (found_scope == nullptr) {
-		return;
-	}
-	rw_mutex_lock(&found_scope->mutex);
-	defer (rw_mutex_unlock(&found_scope->mutex));
-
-	// IMPORTANT NOTE(bill, 2021-04-10): Overriding behaviour was flawed in that the
-	// original entity was still used check checked, but the checking was only
-	// relying on "constant" data such as the Entity.type and Entity.Constant.value
-	//
-	// Therefore two things can be done: the type can be assigned to state that it
-	// has been "evaluated" and the variant data can be copied across
-
-	string_map_set(&found_scope->elements, original_name, new_entity);
-
-	original_entity->flags |= EntityFlag_Overridden;
-	original_entity->type = new_entity->type;
-	original_entity->aliased_of = 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
-	// This is most likely NEVER required, but it does not at all hurt to keep
-	isize offset = cast(u8 *)&original_entity->Dummy.start - cast(u8 *)original_entity;
-	isize size = gb_size_of(*original_entity) - offset;
-	gb_memmove(cast(u8 *)original_entity, cast(u8 *)new_entity, size);
-}
-
-
-
 gb_internal void check_const_decl(CheckerContext *ctx, Entity *e, Ast *type_expr, Ast *init, Type *named_type) {
 	GB_ASSERT(e->type == nullptr);
 	GB_ASSERT(e->kind == Entity_Constant);
@@ -441,41 +553,7 @@ gb_internal void check_const_decl(CheckerContext *ctx, Entity *e, Ast *type_expr
 
 	if (init != nullptr) {
 		Entity *entity = check_entity_from_ident_or_selector(ctx, init, false);
-		if (entity != nullptr && entity->kind == Entity_TypeName) {
-			// @TypeAliasingProblem
-			// NOTE(bill, 2022-02-03): This is used to solve the problem caused by type aliases
-			// being "confused" as constants
-			//
-			//         A :: B
-			//         C :: proc "c" (^A)
-			//         B :: struct {x: C}
-			//
-			//     A gets evaluated first, and then checks B.
-			//     B then checks C.
-			//     C then tries to check A which is unresolved but thought to be a constant.
-			//     Therefore within C's check, A errs as "not a type".
-			//
-			// This is because a const declaration may or may not be a type and this cannot
-			// be determined from a syntactical standpoint.
-			// This check allows the compiler to override the entity to be checked as a type.
-			//
-			// There is no problem if B is prefixed with the `#type` helper enforcing at
-			// both a syntax and semantic level that B must be a type.
-			//
-			//         A :: #type B
-			//
-			// This approach is not fool proof and can fail in case such as:
-			//
-			//         X :: type_of(x)
-			//         X :: Foo(int).Type
-			//
-			// Since even these kind of declarations may cause weird checking cycles.
-			// For the time being, these are going to be treated as an unfortunate error
-			// until there is a proper delaying system to try declaration again if they
-			// have failed.
-
-			e->kind = Entity_TypeName;
-			check_type_decl(ctx, e, init, named_type);
+		if (check_try_override_const_decl(ctx, e, entity, init, named_type)) {
 			return;
 		}
 		entity = nullptr;
diff --git a/src/check_expr.cpp b/src/check_expr.cpp
index 2d0b9dfa9..af5a9c521 100644
--- a/src/check_expr.cpp
+++ b/src/check_expr.cpp
@@ -4968,7 +4968,27 @@ gb_internal bool is_entity_declared_for_selector(Entity *entity, Scope *import_s
 
 // NOTE(bill, 2022-02-03): see `check_const_decl` for why it exists reasoning
 gb_internal Entity *check_entity_from_ident_or_selector(CheckerContext *c, Ast *node, bool ident_only) {
-	if (node->kind == Ast_Ident) {
+	if (node == nullptr) {
+		return nullptr;
+	}
+	if (node->kind == Ast_TernaryWhenExpr) {
+		ast_node(we, TernaryWhenExpr, node);
+		if (we->cond == nullptr) {
+			return nullptr;
+		}
+		if (we->cond->tav.mode != Addressing_Constant) {
+			return nullptr;
+		}
+		if (we->cond->tav.value.kind != ExactValue_Bool) {
+			return nullptr;
+		}
+		if (we->cond->tav.value.value_bool) {
+			return check_entity_from_ident_or_selector(c, we->x, ident_only);
+		} else {
+			Entity *e = check_entity_from_ident_or_selector(c, we->y, ident_only);
+			return e;
+		}
+	} else if (node->kind == Ast_Ident) {
 		String name = node->Ident.token.string;
 		return scope_lookup(c->scope, name);
 	} else if (!ident_only) if (node->kind == Ast_SelectorExpr) {
diff --git a/src/checker.cpp b/src/checker.cpp
index deee93dcc..503494c76 100644
--- a/src/checker.cpp
+++ b/src/checker.cpp
@@ -1470,6 +1470,7 @@ gb_internal Entity *implicit_entity_of_node(Ast *clause) {
 }
 
 gb_internal Entity *entity_of_node(Ast *expr) {
+retry:;
 	expr = unparen_expr(expr);
 	switch (expr->kind) {
 	case_ast_node(ident, Ident, expr);
@@ -1490,6 +1491,17 @@ gb_internal Entity *entity_of_node(Ast *expr) {
 	case_ast_node(ce, CallExpr, expr);
 		return ce->entity_procedure_of;
 	case_end;
+
+	case_ast_node(we, TernaryWhenExpr, expr);
+		if (we->cond == nullptr) {
+			break;
+		}
+		if (we->cond->tav.value.kind != ExactValue_Bool) {
+			break;
+		}
+		expr = we->cond->tav.value.value_bool ? we->x : we->y;
+		goto retry;
+	case_end;
 	}
 	return nullptr;
 }