From af1d4d6e72eeb75b32c40f3d4ca7bf6a78e8a043 Mon Sep 17 00:00:00 2001
From: gingerBill <bill@gingerbill.org>
Date: Thu, 14 May 2020 00:00:10 +0100
Subject: [PATCH] Implement Allow `.?` operator to unwrap any union #549

---
 src/check_expr.cpp | 21 ++++++++++++++++++++-
 src/parser.hpp     |  2 +-
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/src/check_expr.cpp b/src/check_expr.cpp
index f6fbd25c4..d55692a41 100644
--- a/src/check_expr.cpp
+++ b/src/check_expr.cpp
@@ -8897,12 +8897,31 @@ ExprKind check_expr_base_internal(CheckerContext *c, Operand *o, Ast *node, Type
 		if (ta->type != nullptr && ta->type->kind == Ast_UnaryExpr && ta->type->UnaryExpr.op.kind == Token_Question) {
 			if (!is_type_union(src)) {
 				gbString str = type_to_string(o->type);
-				error(o->expr, "Type assertions with .? can only operate on unions with 1 variant, got %s", str);
+				error(o->expr, "Type assertions with .? can only operate on unions, got %s", str);
 				gb_string_free(str);
 				o->mode = Addressing_Invalid;
 				o->expr = node;
 				return kind;
 			}
+
+			if (bsrc->Union.variants.count != 1 && type_hint != nullptr) {
+				bool allowed = false;
+				for_array(i, bsrc->Union.variants) {
+					Type *vt = bsrc->Union.variants[i];
+					if (are_types_identical(vt, type_hint)) {
+						allowed = true;
+						add_type_info_type(c, vt);
+						break;
+					}
+				}
+				if (allowed) {
+					add_type_info_type(c, o->type);
+					o->type = type_hint;
+					o->mode = Addressing_OptionalOk;
+					return kind;
+				}
+			}
+
 			if (bsrc->Union.variants.count != 1) {
 				error(o->expr, "Type assertions with .? can only operate on unions with 1 variant, got %lld", cast(long long)bsrc->Union.variants.count);
 				o->mode = Addressing_Invalid;
diff --git a/src/parser.hpp b/src/parser.hpp
index ca480f403..61be8c54a 100644
--- a/src/parser.hpp
+++ b/src/parser.hpp
@@ -292,7 +292,7 @@ AST_KIND(_ExprBegin,  "",  bool) \
 	AST_KIND(TernaryExpr,     "ternary expression",       struct { Ast *cond, *x, *y; }) \
 	AST_KIND(TernaryIfExpr,   "ternary if expression",    struct { Ast *x, *cond, *y; }) \
 	AST_KIND(TernaryWhenExpr, "ternary when expression",  struct { Ast *x, *cond, *y; }) \
-	AST_KIND(TypeAssertion, "type assertion",      struct { Ast *expr; Token dot; Ast *type; }) \
+	AST_KIND(TypeAssertion, "type assertion",      struct { Ast *expr; Token dot; Ast *type; Type *type_hint; }) \
 	AST_KIND(TypeCast,      "type cast",           struct { Token token; Ast *type, *expr; }) \
 	AST_KIND(AutoCast,      "auto_cast",           struct { Token token; Ast *expr; }) \
 AST_KIND(_ExprEnd,       "", bool) \