From 95186e20d544dda136ad0f21ac36324418bce5e2 Mon Sep 17 00:00:00 2001
From: Brendan Hansen <brendan.f.hansen@gmail.com>
Date: Fri, 29 Sep 2023 16:58:01 -0500
Subject: [PATCH] initial implementation of switch expressions

---
 compiler/include/astnodes.h    | 13 +++--
 compiler/src/astnodes.c        | 28 +++++++++++
 compiler/src/checker.c         | 65 ++++++++++++++++++++----
 compiler/src/parser.c          | 45 +++++++++++------
 compiler/src/symres.c          | 32 ++++++++++--
 compiler/src/wasm_emit.c       | 60 ++++++++++++++++++----
 core/container/optional.onyx   | 70 ++++++++++++--------------
 core/container/result.onyx     | 92 +++++++++++++++++-----------------
 core/encoding/hex.onyx         | 13 +++--
 core/encoding/json/parser.onyx | 38 +++++++-------
 core/io/stdio.onyx             | 38 +++++++-------
 core/net/net.onyx              | 10 ++--
 core/onyx/cbindgen.onyx        | 72 +++++++++++++-------------
 core/runtime/info/helper.onyx  | 70 +++++++++++++-------------
 tests/aoc-2020/day20.onyx      | 23 ++++-----
 tests/switch_expressions       |  5 ++
 tests/switch_expressions.onyx  | 55 ++++++++++++++++++++
 tests/tagged_unions.onyx       | 16 +++---
 18 files changed, 479 insertions(+), 266 deletions(-)
 create mode 100644 tests/switch_expressions
 create mode 100644 tests/switch_expressions.onyx

diff --git a/compiler/include/astnodes.h b/compiler/include/astnodes.h
index c88f5bdf..1bf31450 100644
--- a/compiler/include/astnodes.h
+++ b/compiler/include/astnodes.h
@@ -881,7 +881,7 @@ typedef enum SwitchKind {
 typedef struct CaseToBlock {
     AstTyped *original_value;
     AstBinaryOp *comparison;
-    AstBlock *block;
+    AstSwitchCase *casestmt;
 } CaseToBlock;
 
 struct AstSwitchCase {
@@ -890,16 +890,21 @@ struct AstSwitchCase {
     // NOTE: All expressions that end up in this block
     bh_arr(AstTyped *) values;
 
-    AstBlock *block;
+    union {
+        AstBlock *block;
+        AstTyped *expr;
+    };
 
     AstLocal *capture;
+    Scope *scope; // Scope for the capture
 
     b32 is_default: 1; // Could this be inferred by the values array being null?
     b32 capture_is_by_pointer: 1;
+    b32 body_is_expr : 1;
 };
 
 struct AstSwitch {
-    AstNode_base;
+    AstTyped_base;
 
     Scope *scope;
     AstNode* initialization;
@@ -918,6 +923,8 @@ struct AstSwitch {
     // been handled.
     u8 *union_variants_handled;
 
+    b32 is_expr;
+
     union {
         struct {
             // NOTE: This is a mapping from the compile time known case value
diff --git a/compiler/src/astnodes.c b/compiler/src/astnodes.c
index f8f43626..1c4583f7 100644
--- a/compiler/src/astnodes.c
+++ b/compiler/src/astnodes.c
@@ -770,6 +770,34 @@ TypeMatch unify_node_and_type_(AstTyped** pnode, Type* type, b32 permanent) {
         }
     }
 
+    if (node->kind == Ast_Kind_Switch) {
+        AstSwitch *switchnode = (AstSwitch *) node;
+        if (!switchnode->is_expr) return TYPE_MATCH_FAILED;
+
+        if (switchnode->cases == NULL) return TYPE_MATCH_YIELD;
+
+        bh_arr_each(AstSwitchCase *, pcasestmt, switchnode->cases) {
+            AstSwitchCase *casestmt = *pcasestmt;
+            if (!casestmt->body_is_expr) continue;
+
+            switch (unify_node_and_type_(&casestmt->expr, type, permanent)) {
+                case TYPE_MATCH_SUCCESS: break;
+                case TYPE_MATCH_FAILED: return TYPE_MATCH_FAILED;
+                case TYPE_MATCH_YIELD: return TYPE_MATCH_YIELD;
+            }
+        }
+
+        if (switchnode->default_case) {
+            switch (unify_node_and_type_((AstTyped **) &switchnode->default_case, type, permanent)) {
+                case TYPE_MATCH_SUCCESS: break;
+                case TYPE_MATCH_FAILED: return TYPE_MATCH_FAILED;
+                case TYPE_MATCH_YIELD: return TYPE_MATCH_YIELD;
+            }
+        }
+
+        if (permanent) switchnode->type = type;
+        return TYPE_MATCH_SUCCESS;
+    }
 
     // If the destination type is an optional, and the node's type is a value of
     // the same underlying type, then we can construct an optional with a value
diff --git a/compiler/src/checker.c b/compiler/src/checker.c
index 6c8aff42..01f4fc28 100644
--- a/compiler/src/checker.c
+++ b/compiler/src/checker.c
@@ -362,7 +362,7 @@ fornode_expr_checked:
     return Check_Success;
 }
 
-static b32 add_case_to_switch_statement(AstSwitch* switchnode, u64 case_value, AstBlock* block, OnyxFilePos pos) {
+static b32 add_case_to_switch_statement(AstSwitch* switchnode, u64 case_value, AstSwitchCase* casestmt, OnyxFilePos pos) {
     assert(switchnode->switch_kind == Switch_Kind_Integer || switchnode->switch_kind == Switch_Kind_Union);
 
     switchnode->min_case = bh_min(switchnode->min_case, case_value);
@@ -373,7 +373,7 @@ static b32 add_case_to_switch_statement(AstSwitch* switchnode, u64 case_value, A
         return 1;
     }
 
-    bh_imap_put(&switchnode->case_map, case_value, (u64) block);
+    bh_imap_put(&switchnode->case_map, case_value, (u64) casestmt);
     return 0;
 }
 
@@ -488,7 +488,7 @@ CheckStatus check_switch(AstSwitch* switchnode) {
         AstSwitchCase *sc = switchnode->cases[i];
 
         if (sc->capture && bh_arr_length(sc->values) != 1) {
-            ERROR(sc->token->pos, "Expected exactly one value in switch-case when using a capture, i.e. `case X => Y { ... }`.");
+            ERROR(sc->token->pos, "Expected exactly one value in switch-case when using a capture, i.e. `case value: X { ... }`.");
         }
 
         if (sc->capture && switchnode->switch_kind != Switch_Kind_Union) {
@@ -519,7 +519,7 @@ CheckStatus check_switch(AstSwitch* switchnode) {
 
                 // NOTE: This is inclusive!!!!
                 fori (case_value, lower, upper + 1) {
-                    if (add_case_to_switch_statement(switchnode, case_value, sc->block, rl->token->pos))
+                    if (add_case_to_switch_statement(switchnode, case_value, sc, rl->token->pos))
                         return Check_Error;
                 }
 
@@ -573,7 +573,7 @@ CheckStatus check_switch(AstSwitch* switchnode) {
                     if (!is_valid)
                         ERROR_((*value)->token->pos, "Case statement expected compile time known integer. Got '%s'.", onyx_ast_node_kind_string((*value)->kind));
 
-                    if (add_case_to_switch_statement(switchnode, integer_value, sc->block, sc->block->token->pos))
+                    if (add_case_to_switch_statement(switchnode, integer_value, sc, sc->block->token->pos))
                         return Check_Error;
 
                     break;
@@ -592,7 +592,7 @@ CheckStatus check_switch(AstSwitch* switchnode) {
                     if (found) break;
 
                     CaseToBlock ctb;
-                    ctb.block = sc->block;
+                    ctb.casestmt = sc;
                     ctb.original_value = *value;
                     ctb.comparison = make_binary_op(context.ast_alloc, Binary_Op_Equal, switchnode->expr, *value);
                     ctb.comparison->token = (*value)->token;
@@ -607,13 +607,52 @@ CheckStatus check_switch(AstSwitch* switchnode) {
         sc->flags |= Ast_Flag_Has_Been_Checked;
 
       check_switch_case_block:
-        CHECK(block, sc->block);
+        if (switchnode->is_expr) {
+            if (!sc->body_is_expr) {
+                onyx_report_error(sc->token->pos, Error_Critical, "Inside a switch expression, all cases must return a value.");
+                ERROR(sc->token->pos, "Change the case statement to look like 'case X => expr'.");
+            }
+        } else {
+            if (sc->body_is_expr) {
+                ERROR(sc->token->pos, "This kind of case statement is only allowed in switch expressions, not switch statements.");
+            }
+        }
+
+        if (sc->body_is_expr) {
+            CHECK(expression, &sc->expr);
+            if (switchnode->type == NULL) {
+                switchnode->type = resolve_expression_type(sc->expr);
+            } else {
+                TYPE_CHECK(&sc->expr, switchnode->type) { 
+                    ERROR_(sc->token->pos, "Expected case expression to be of type '%s', got '%s'.",
+                        type_get_name(switchnode->type),
+                        type_get_name(sc->expr->type));
+                }
+            }
+
+        } else {
+            CHECK(block, sc->block);
+        }
 
         switchnode->yield_return_index += 1;
     }
 
     if (switchnode->default_case) {
-        CHECK(block, switchnode->default_case);
+        if (switchnode->is_expr) {
+            AstTyped **default_case = (AstTyped **) &switchnode->default_case;
+            CHECK(expression, default_case);
+
+            if (switchnode->type) {
+                TYPE_CHECK(default_case, switchnode->type) { 
+                    ERROR_((*default_case)->token->pos, "Expected case expression to be of type '%s', got '%s'.",
+                        type_get_name(switchnode->type),
+                        type_get_name((*default_case)->type));
+                }
+            }
+
+        } else {
+            CHECK(block, switchnode->default_case);
+        }
 
     } else if (switchnode->switch_kind == Switch_Kind_Union) {
         // If there is no default case, and this is a union switch,
@@ -2411,6 +2450,15 @@ CheckStatus check_expression(AstTyped** pexpr) {
             break;
         }
 
+        case Ast_Kind_Switch: {
+            AstSwitch* switch_node = (AstSwitch *) expr;
+            assert(switch_node->is_expr);
+
+            CHECK(switch, switch_node);
+            break;
+        }
+
+        case Ast_Kind_Switch_Case: break;
         case Ast_Kind_File_Contents: break;
         case Ast_Kind_Overloaded_Function: break;
         case Ast_Kind_Enum_Value: break;
@@ -2418,7 +2466,6 @@ CheckStatus check_expression(AstTyped** pexpr) {
         case Ast_Kind_Package: break;
         case Ast_Kind_Error: break;
         case Ast_Kind_Unary_Field_Access: break;
-        case Ast_Kind_Switch_Case: break;
         case Ast_Kind_Foreign_Block: break;
         case Ast_Kind_Zero_Value: break;
         case Ast_Kind_Interface: break;
diff --git a/compiler/src/parser.c b/compiler/src/parser.c
index c06c2ac9..123cc074 100644
--- a/compiler/src/parser.c
+++ b/compiler/src/parser.c
@@ -50,7 +50,7 @@ static AstIfWhile*    parse_if_stmt(OnyxParser* parser);
 static AstIfWhile*    parse_while_stmt(OnyxParser* parser);
 static AstFor*        parse_for_stmt(OnyxParser* parser);
 static AstSwitchCase* parse_case_stmt(OnyxParser* parser);
-static AstSwitch*     parse_switch_stmt(OnyxParser* parser);
+static AstSwitch*     parse_switch_stmt(OnyxParser* parser, b32 switch_is_expr);
 static i32            parse_possible_symbol_declaration(OnyxParser* parser, AstNode** ret);
 static AstReturn*     parse_return_stmt(OnyxParser* parser);
 static AstBlock*      parse_block(OnyxParser* parser, b32 make_a_new_scope, char* block_name);
@@ -631,6 +631,11 @@ static AstTyped* parse_factor(OnyxParser* parser) {
             break;
         }
 
+        case Token_Type_Keyword_Switch: {
+            retval = (AstTyped *) parse_switch_stmt(parser, 1);
+            break;
+        }
+
         case '[': {
             // HACK CLEANUP
             // :LinearTokenDependent
@@ -1336,6 +1341,23 @@ static AstSwitchCase* parse_case_stmt(OnyxParser* parser) {
         sc_node->is_default = 1;
 
     } else {
+        if (   next_tokens_are(parser, 3, '&', Token_Type_Symbol, ':')
+            || next_tokens_are(parser, 3, '^', Token_Type_Symbol, ':')
+            || next_tokens_are(parser, 2, Token_Type_Symbol, ':')
+       ) {
+            b32 is_pointer = 0;
+            if (consume_token_if_next(parser, '&') || consume_token_if_next(parser, '^'))
+                is_pointer = 1;
+            
+            OnyxToken *capture_symbol = expect_token(parser, Token_Type_Symbol);
+            AstLocal *capture = make_local(parser->allocator, capture_symbol, NULL);
+            
+            sc_node->capture = capture;
+            sc_node->capture_is_by_pointer = is_pointer;
+
+            expect_token(parser, ':');
+        }
+
         bh_arr_new(global_heap_allocator, sc_node->values, 1);
 
         parser->parse_quick_functions = 0;
@@ -1352,26 +1374,19 @@ static AstSwitchCase* parse_case_stmt(OnyxParser* parser) {
     }
 
     if (consume_token_if_next(parser, Token_Type_Fat_Right_Arrow)) {
-        // Captured value for union switching
-        b32 is_pointer = 0;
-        if (consume_token_if_next(parser, '&') || consume_token_if_next(parser, '^'))
-            is_pointer = 1;
-        
-        OnyxToken *capture_symbol = expect_token(parser, Token_Type_Symbol);
-        AstLocal *capture = make_local(parser->allocator, capture_symbol, NULL);
-        
-        sc_node->capture = capture;
-        sc_node->capture_is_by_pointer = is_pointer;
+        sc_node->expr = parse_expression(parser, 0);
+        sc_node->body_is_expr = 1;
+    } else {
+        sc_node->block = parse_block(parser, 1, NULL);
     }
 
-    sc_node->block = parse_block(parser, 1, NULL);
-
     return sc_node;
 }
 
-static AstSwitch* parse_switch_stmt(OnyxParser* parser) {
+static AstSwitch* parse_switch_stmt(OnyxParser* parser, b32 switch_is_expr) {
     AstSwitch* switch_node = make_node(AstSwitch, Ast_Kind_Switch);
     switch_node->token = expect_token(parser, Token_Type_Keyword_Switch);
+    switch_node->is_expr = switch_is_expr;
 
     AstTyped* expr;
     AstNode* initialization_or_expr=NULL;
@@ -1642,7 +1657,7 @@ static AstNode* parse_statement(OnyxParser* parser) {
 
         case Token_Type_Keyword_Switch:
             needs_semicolon = 0;
-            retval = (AstNode *) parse_switch_stmt(parser);
+            retval = (AstNode *) parse_switch_stmt(parser, 0);
             break;
 
         case Token_Type_Keyword_Case:
diff --git a/compiler/src/symres.c b/compiler/src/symres.c
index 6f934039..9ec553ca 100644
--- a/compiler/src/symres.c
+++ b/compiler/src/symres.c
@@ -744,6 +744,11 @@ static SymresStatus symres_expression(AstTyped** expr) {
             break;
         }
 
+        case Ast_Kind_Switch: {
+            SYMRES(switch, (AstSwitch *) *expr);
+            break;
+        }
+
         default: break;
     }
 
@@ -832,11 +837,24 @@ static SymresStatus symres_case(AstSwitchCase *casenode) {
     }
 
     if (casenode->capture) {
-        casenode->block->scope = scope_create(context.ast_alloc, current_scope, casenode->block->token->pos);
-        symbol_introduce(casenode->block->scope, casenode->capture->token, (AstNode *) casenode->capture);
+        if (casenode->scope == NULL) {
+            casenode->scope = scope_create(context.ast_alloc, current_scope, casenode->token->pos);
+            symbol_introduce(casenode->scope, casenode->capture->token, (AstNode *) casenode->capture);
+        }
+
+        scope_enter(casenode->scope);
+    }
+
+    if (casenode->body_is_expr) {
+        SYMRES(expression, &casenode->expr);
+    } else {
+        SYMRES(block, casenode->block);
+    }
+
+    if (casenode->capture) {
+        scope_leave();
     }
 
-    SYMRES(block, casenode->block);
     return Symres_Success;
 }
 
@@ -857,7 +875,13 @@ static SymresStatus symres_switch(AstSwitch* switchnode) {
             SYMRES(case, *pcase);
         }
 
-        if (switchnode->default_case) SYMRES(block, switchnode->default_case);
+        if (switchnode->default_case) {
+            if (switchnode->is_expr) {
+                SYMRES(expression, (AstTyped **) &switchnode->default_case);
+            } else {
+                SYMRES(block, switchnode->default_case);
+            }
+        }
     }
 
     if (switchnode->switch_kind == Switch_Kind_Use_Equals && switchnode->case_exprs) {
diff --git a/compiler/src/wasm_emit.c b/compiler/src/wasm_emit.c
index 7afd7c6d..93d1f395 100644
--- a/compiler/src/wasm_emit.c
+++ b/compiler/src/wasm_emit.c
@@ -1612,6 +1612,12 @@ EMIT_FUNC(switch, AstSwitch* switch_node) {
     bh_imap block_map;
     bh_imap_init(&block_map, global_heap_allocator, bh_arr_length(switch_node->cases));
 
+    u64 expr_result_local = 0;
+    if (switch_node->is_expr) {
+        expr_result_local = local_allocate(mod->local_alloc, (AstTyped *) switch_node);
+        bh_imap_put(&mod->local_map, (u64) switch_node, expr_result_local);
+    }
+
     if (switch_node->initialization != NULL) {
         forll (AstNode, stmt, switch_node->initialization, next) {
             emit_statement(mod, &code, stmt);
@@ -1622,11 +1628,11 @@ EMIT_FUNC(switch, AstSwitch* switch_node) {
 
     u64 block_num = 0;
     bh_arr_each(AstSwitchCase *, sc, switch_node->cases) {
-        if (bh_imap_has(&block_map, (u64) (*sc)->block)) continue;
+        if (bh_imap_has(&block_map, (u64) *sc)) continue;
 
         emit_enter_structured_block(mod, &code, SBT_Fallthrough_Block, (*sc)->block->token);
 
-        bh_imap_put(&block_map, (u64) (*sc)->block, block_num);
+        bh_imap_put(&block_map, (u64) *sc, block_num);
         block_num++;
     }
 
@@ -1642,6 +1648,7 @@ EMIT_FUNC(switch, AstSwitch* switch_node) {
             fori (i, 0, bt->count) bt->cases[i] = bt->default_case;
 
             bh_arr_each(bh__imap_entry, sc, switch_node->case_map.entries) {
+                assert(bh_imap_has(&block_map, (u64) sc->value));
                 bt->cases[sc->key - switch_node->min_case] = bh_imap_get(&block_map, (u64) sc->value);
             }
 
@@ -1691,7 +1698,7 @@ EMIT_FUNC(switch, AstSwitch* switch_node) {
             bh_arr_each(CaseToBlock, ctb, switch_node->case_exprs) {
                 emit_expression(mod, &code, (AstTyped *) ctb->comparison);
 
-                u64 bn = bh_imap_get(&block_map, (u64) ctb->block);
+                u64 bn = bh_imap_get(&block_map, (u64) ctb->casestmt);
                 WID(switch_node->expr->token, WI_IF_START, 0x40);
                 WID(switch_node->expr->token, WI_JUMP, bn + 1);
                 WI(switch_node->expr->token, WI_IF_END);
@@ -1707,7 +1714,7 @@ EMIT_FUNC(switch, AstSwitch* switch_node) {
         AstSwitchCase *sc = *psc;
         if (bh_imap_get(&block_map, (u64) sc->block) == 0xdeadbeef) continue;
 
-        u64 bn = bh_imap_get(&block_map, (u64) sc->block);
+        u64 bn = bh_imap_get(&block_map, (u64) sc);
 
         if (sc->capture) {
             assert(union_capture_idx != 0);
@@ -1744,7 +1751,14 @@ EMIT_FUNC(switch, AstSwitch* switch_node) {
 
         // Maybe the Symbol Frame idea should be controlled as a block_flag?
         debug_enter_symbol_frame(mod);
-        emit_block(mod, &code, sc->block, 0);
+        if (sc->body_is_expr) {
+            emit_expression(mod, &code, sc->expr);
+            emit_generic_store_instruction(mod, &code, (AstTyped *) switch_node, switch_node->token);
+
+        } else {
+            emit_block(mod, &code, sc->block, 0);
+        }
+
         debug_leave_symbol_frame(mod);
 
         if (bh_arr_last(code).type != WI_JUMP)
@@ -1752,16 +1766,35 @@ EMIT_FUNC(switch, AstSwitch* switch_node) {
 
         emit_leave_structured_block(mod, &code);
 
-        bh_imap_put(&block_map, (u64) sc->block, 0xdeadbeef);
+        bh_imap_put(&block_map, (u64) sc, 0xdeadbeef);
     }
 
     if (switch_node->default_case != NULL) {
-        emit_block(mod, &code, switch_node->default_case, 0);
+        if (switch_node->is_expr) {
+            emit_expression(mod, &code, (AstTyped *) switch_node->default_case);
+            emit_generic_store_instruction(mod, &code, (AstTyped *) switch_node, switch_node->token);
+
+        } else {
+            emit_block(mod, &code, switch_node->default_case, 0);
+        }
     }
 
     if (union_capture_idx != 0) local_raw_free(mod->local_alloc, WASM_TYPE_PTR);
     emit_leave_structured_block(mod, &code);
 
+    if (switch_node->is_expr) {
+        if ((expr_result_local & LOCAL_IS_WASM) == 0) {
+            u64 offset = 0;
+            emit_local_location(mod, &code, (AstLocal *) switch_node, &offset);
+            emit_load_instruction(mod, &code, switch_node->type, offset);
+
+        } else {
+            WIL(switch_node->token, WI_LOCAL_GET, expr_result_local);
+        }
+
+        local_free(mod->local_alloc, (AstTyped *) switch_node);
+    }
+
     bh_imap_free(&block_map);
     *pcode = code;
 }
@@ -3279,7 +3312,8 @@ EMIT_FUNC(location_return_offset, AstTyped* expr, u64* offset_return) {
         case Ast_Kind_Do_Block:
         case Ast_Kind_If_Expression:
         case Ast_Kind_Call_Site:
-        case Ast_Kind_Zero_Value: {
+        case Ast_Kind_Zero_Value:
+        case Ast_Kind_Switch: {
             emit_local_location(mod, &code, (AstLocal *) expr, offset_return);
             break;
         }
@@ -3822,8 +3856,16 @@ EMIT_FUNC(expression, AstTyped* expr) {
             break;
         }
 
+        case Ast_Kind_Switch: {
+            AstSwitch* switchnode = (AstSwitch *) expr;
+            assert(switchnode->is_expr);
+
+            emit_switch(mod, &code, switchnode);
+            break;
+        }
+
         default:
-            bh_printf("Unhandled case: %d\n", expr->kind);
+            bh_printf("Unhandled case: %s\n", onyx_ast_node_kind_string(expr->kind));
             DEBUG_HERE;
             assert(0);
     }
diff --git a/core/container/optional.onyx b/core/container/optional.onyx
index fa7bad94..c9c7a992 100644
--- a/core/container/optional.onyx
+++ b/core/container/optional.onyx
@@ -46,11 +46,9 @@ use core
         Extracts the value from the Optional, or uses a default if
         no value is present.
     """
-    value_or :: (o: ?$T, default: T) -> T {
-        switch o {
-            case .Some => v do return v;
-            case #default do return default;
-        }
+    value_or :: (o: ?$T, default: T) => switch o {
+        case v: .Some => v;
+        case #default => default;
     }
 
     #doc "Clears the value in the Optional, zeroing the memory of the value."
@@ -65,16 +63,16 @@ use core
 
     #doc "Monadic chaining operation."
     and_then :: (o: ?$T, transform: (T) -> ?$R) -> ?R {
-        switch o {
-            case .Some => v do return transform(v);
-            case #default do return .{ None = .{} };
-        }
+        return switch o {
+            case v: .Some => transform(v);
+            case #default => .{ None = .{} };
+        };
     }
 
     #doc "Changes the value inside the optional, if present."
     transform :: (o: ?$T, transform: (T) -> $R) -> ?R {
         switch o {
-            case .Some => v do return .{ Some = transform(v) };
+            case v: .Some do return .{ Some = transform(v) };
             case #default do return .{ None = .{} };
         }
     }
@@ -84,10 +82,10 @@ use core
         provide a function to generate a value.
     """
     or_else :: (o: ?$T, generate: () -> ?T) -> ?T {
-        switch o {
-            case .Some => v do return o;
-            case #default do return generate();
-        }
+        return switch o {
+            case .Some => o;
+            case #default => generate();
+        };
     }
 
     #doc """
@@ -97,7 +95,7 @@ use core
     """
     unwrap :: (o: ?$T) -> T {
         switch o {
-            case .Some => v do return v;
+            case v: .Some do return v;
             case #default {
                 assert(false, "Unwrapping empty Optional.");
             }
@@ -111,7 +109,7 @@ use core
     """
     unwrap_ptr :: (o: & ?$T) -> &T {
         switch o {
-            case .Some => &v do return v;
+            case &v: .Some do return v;
             case #default {
                 assert(false, "Unwrapping empty Optional.");
             }
@@ -121,7 +119,7 @@ use core
     or_return :: #match {
         macro (o: ?$T) -> T {
             switch value := o; value {
-                case .Some => v do return v;
+                case v: .Some do return v;
                 case #default {
                     return return .{};
                 }
@@ -129,7 +127,7 @@ use core
         },
         macro (o: ?$T, return_value: $R) -> T {
             switch value := o; value {
-                case .Some => v do return v;
+                case v: .Some do return v;
                 case #default {
                     return return return_value;
                 }
@@ -139,7 +137,7 @@ use core
 
     catch :: macro (o: ?$T, body: Code) -> T {
         switch value := o; value {
-            case .Some => v do return v;
+            case v: .Some do return v;
             case .None {
                 #unquote body;
             }
@@ -149,7 +147,7 @@ use core
     with :: macro (o: ?$T, body: Code) {
         switch o {
             case .None ---;
-            case .Some => it {
+            case it: .Some {
                 #unquote body(it);
             }
         }
@@ -207,35 +205,31 @@ use core
     }
 
 
-    hash :: (o: ?$T/core.hash.Hashable) -> u32 {
-        switch o {
-            case .Some => v do return core.hash.hash(v);
-            case #default do return 0;
-        }
+    hash :: (o: ?$T/core.hash.Hashable) => switch o {
+        case v: .Some => core.hash.hash(v);
+        case #default => 0;
     }
 }
 
 #operator == (o1, o2: ?$T) -> bool {
     if cast(Optional(T).tag_enum, o1) != cast(Optional(T).tag_enum, o2) do return false;
-    switch o1 {
-        case .None do return true;
-        case .Some => v1 {
-            v2 := o2->unwrap();
-            return v1 == v2;
-        }
-    }
+
+    return switch o1 {
+        case .None => true;
+        case v1: .Some => v1 == o2->unwrap();
+    };
 }
 
 #operator ?? macro (opt: ?$T, default: T) -> T {
-    switch value := opt; value {
-        case .Some => v do return v;
-        case #default do return default;
-    }
+    return switch value := opt; value {
+        case v: .Some => v;
+        case #default => default;
+    };
 }
 
 #operator ?? macro (opt: ?$T, catch: Code) -> T {
     switch value := opt; value {
-        case .Some => v do return v;
+        case v: .Some do return v;
         case #default ---
     }
 
@@ -244,7 +238,7 @@ use core
 
 #operator ? macro (opt: ?$T) -> T {
     switch value := opt; value {
-        case .Some => v do return v;
+        case v: .Some do return v;
         case #default do return return .{};
     }
 }
diff --git a/core/container/result.onyx b/core/container/result.onyx
index 7183b8f6..a266d062 100644
--- a/core/container/result.onyx
+++ b/core/container/result.onyx
@@ -27,34 +27,34 @@ Result :: union (Ok_Type: type_expr, Err_Type: type_expr) {
 #inject Result {
     #doc "Returns true if the result contains an Ok value."
     is_ok :: (r: #Self) -> bool {
-        switch r {
-            case .Ok do return true;
-            case #default do return false;
-        }
+        return switch r {
+            case .Ok => true;
+            case #default => false;
+        };
     }
 
     #doc "Returns true if the result contains an Err value."
     is_err :: (r: #Self) -> bool {
-        switch r {
-            case .Err do return true;
-            case #default do return false;
-        }
+        return switch r {
+            case .Err => true;
+            case #default => false;
+        };
     }
 
     #doc "Returns an Optional of the Ok type."
     ok :: (r: #Self) -> Optional(r.Ok_Type) {
-        switch r {
-            case .Ok => v do return v;
-            case #default do return .{};
-        }
+        return switch r {
+            case v: .Ok => Optional.make(v);
+            case #default => .{};
+        };
     }
 
     #doc "Returns an Optional of the Err type."
     err :: (r: #Self) -> Optional(r.Err_Type) {
-        switch r {
-            case .Err => v do return v;
-            case #default do return .{};
-        }
+        return switch r {
+            case v: .Err => Optional.make(v);
+            case #default => .{};
+        };
     }
 
     #doc """
@@ -63,8 +63,8 @@ Result :: union (Ok_Type: type_expr, Err_Type: type_expr) {
     """
     unwrap :: (r: #Self) -> r.Ok_Type {
         switch r {
-            case .Ok => v do return v;
-            case .Err => err {
+            case v: .Ok do return v;
+            case err: .Err {
                 msg := tprintf("Unwrapping Result with error '{}'.", err);
                 assert(false, msg);
                 return .{};
@@ -77,10 +77,10 @@ Result :: union (Ok_Type: type_expr, Err_Type: type_expr) {
         result contains an Err, the empty .{} value is returned.
     """
     unwrap_or_default :: (r: #Self) -> r.Ok_Type {
-        switch r {
-            case .Ok => v do return v;
-            case #default do return .{};
-        }
+        return switch r {
+            case v: .Ok => v;
+            case #default => .{};
+        };
     }
 
     #doc """
@@ -89,7 +89,7 @@ Result :: union (Ok_Type: type_expr, Err_Type: type_expr) {
     """
     expect :: (r: #Self, msg: str) -> r.Ok_Type {
         switch r {
-            case .Ok => v do return v;
+            case v: .Ok do return v;
             case #default {
                 assert(false, msg);
             }
@@ -102,26 +102,26 @@ Result :: union (Ok_Type: type_expr, Err_Type: type_expr) {
             Err(e) => Err(e)
     """
     transform :: (r: Result($T, $E), f: (T) -> $R) -> Result(R, E) {
-        switch r {
-            case .Ok => v  do return .{ Ok = f(v) };
-            case .Err => v do return .{ Err = v };
-        }
+        return switch r {
+            case v: .Ok => .{ Ok = f(v) };
+            case v: .Err => .{ Err = v };
+        };
     }
 
     #doc "Monadic chaining operation."
     and_then :: (r: #Self, f: (r.Ok_Type) -> Result($R, r.Err_Type)) -> Result(R, r.Err_Type) {
-        switch r {
-            case .Ok => v  do return f(v);
-            case .Err => v do return .{ Err = v };
-        }
+        return switch r {
+            case v: .Ok  => f(v);
+            case v: .Err => .{ Err = v };
+        };
     }
 
     #doc "If the Result contains Err, generate is called to make a value"
     or_else :: (r: #Self, generate: () -> typeof r) => {
-        switch r {
-            case .Ok => v do return v;
-            case #default do return generate();
-        }
+        return switch r {
+            case v: .Ok   => v;
+            case #default => generate();
+        };
     }
 
     #doc """
@@ -142,8 +142,8 @@ Result :: union (Ok_Type: type_expr, Err_Type: type_expr) {
     """
     forward_err :: macro (r: Result($T, $E)) -> T {
         switch res := r; res {
-            case .Ok => v  do return v;
-            case .Err => v do return return .{ Err = v };
+            case v: .Ok  do return v;
+            case v: .Err do return return .{ Err = v };
         }
     }
 
@@ -153,7 +153,7 @@ Result :: union (Ok_Type: type_expr, Err_Type: type_expr) {
     """
     or_return :: macro (r: Result($T, $E), v: $V) -> T {
         switch res := r; res {
-            case .Ok => v  do return v;
+            case v: .Ok  do return v;
             case .Err do return return v;
         }
     }
@@ -168,8 +168,8 @@ Result :: union (Ok_Type: type_expr, Err_Type: type_expr) {
     """
     catch :: macro (r: Result($T, $E), on_err: Code) -> T {
         switch res := r; res {
-            case .Ok => v  do return v;
-            case .Err => err {
+            case v: .Ok  do return v;
+            case err: .Err {
                 #unquote on_err(err);
             }
         }
@@ -182,16 +182,16 @@ __implicit_bool_cast :: macro (r: Result($O, $E)) => cast(Result(O, E).tag_enum,
 
 #operator ? macro (r: Result($T, $E)) -> T {
     switch res := r; res {
-        case .Ok => v  do return v;
-        case .Err => v do return return .{ Err = v };
+        case v: .Ok do return v;
+        case v: .Err do return return .{ Err = v };
     }
 }
 
 #operator ?? macro (r: Result($T, $E), v: T) -> T {
-    switch res := r; res {
-        case .Ok => v do return v;
-        case .Err do return v;
-    }
+    return switch res := r; res {
+        case val: .Ok => val;
+        case .Err => v;
+    };
 }
 
 
diff --git a/core/encoding/hex.onyx b/core/encoding/hex.onyx
index bb1c4158..2feb1991 100644
--- a/core/encoding/hex.onyx
+++ b/core/encoding/hex.onyx
@@ -23,12 +23,11 @@ decode :: (s: str, allocator := context.allocator) -> str {
     return new_str;
 
     digit_to_value :: (it: u8) -> u32 {
-        switch it {
-            case #char "0" .. #char "9" do return ~~(it - #char "0");
-            case #char "a" .. #char "f" do return ~~(it - #char "a" + 10);
-            case #char "A" .. #char "F" do return ~~(it - #char "A" + 10);
-        }
-
-        return 0;
+        return ~~ switch it {
+            case #char "0" .. #char "9" => it - #char "0";
+            case #char "a" .. #char "f" => it - #char "a" + 10;
+            case #char "A" .. #char "F" => it - #char "A" + 10;
+            case #default => 0;
+        };
     }
 }
\ No newline at end of file
diff --git a/core/encoding/json/parser.onyx b/core/encoding/json/parser.onyx
index 5ebb2f01..72a73048 100644
--- a/core/encoding/json/parser.onyx
+++ b/core/encoding/json/parser.onyx
@@ -405,25 +405,23 @@ parse_and_write_utf8_character :: (s: str, out: [&] u8) -> (i32, i32) {
     return 0, 0;
 
     digit_to_hex :: (c: u8) -> i32 {
-        switch c {
-            case #char "0" do return 0;
-            case #char "1" do return 1;
-            case #char "2" do return 2;
-            case #char "3" do return 3;
-            case #char "4" do return 4;
-            case #char "5" do return 5;
-            case #char "6" do return 6;
-            case #char "7" do return 7;
-            case #char "8" do return 8;
-            case #char "9" do return 9;
-            case #char "A", #char "a" do return 10;
-            case #char "B", #char "b" do return 11;
-            case #char "C", #char "c" do return 12;
-            case #char "D", #char "d" do return 13;
-            case #char "E", #char "e" do return 14;
-            case #char "F", #char "f" do return 15;
-        }
-
-        return 0;
+        return switch c {
+            case '0' => 0;
+            case '1' => 1;
+            case '2' => 2;
+            case '3' => 3;
+            case '4' => 4;
+            case '5' => 5;
+            case '6' => 6;
+            case '7' => 7;
+            case '8' => 8;
+            case '9' => 9;
+            case 'A', 'a' => 10;
+            case 'B', 'b' => 11;
+            case 'C', 'c' => 12;
+            case 'D', 'd' => 13;
+            case 'E', 'e' => 14;
+            case 'F', 'f' => 15;
+        };
     }
 }
diff --git a/core/io/stdio.onyx b/core/io/stdio.onyx
index 14317c9b..f9c83204 100644
--- a/core/io/stdio.onyx
+++ b/core/io/stdio.onyx
@@ -156,25 +156,25 @@ __byte_dump :: (ptr: rawptr, byte_count: u32, bytes_per_line := 8) {
 
 
     map_to_ascii :: (x: u8) -> u8 {
-        switch x {
-            case 0 do return  #char "0";
-            case 1 do return  #char "1";
-            case 2 do return  #char "2";
-            case 3 do return  #char "3";
-            case 4 do return  #char "4";
-            case 5 do return  #char "5";
-            case 6 do return  #char "6";
-            case 7 do return  #char "7";
-            case 8 do return  #char "8";
-            case 9 do return  #char "9";
-            case 10 do return #char "A";
-            case 11 do return #char "B";
-            case 12 do return #char "C";
-            case 13 do return #char "D";
-            case 14 do return #char "E";
-            case 15 do return #char "F";
-            case #default do return #char "X";
-        }
+        return switch x {
+            case 0 =>  #char "0";
+            case 1 =>  #char "1";
+            case 2 =>  #char "2";
+            case 3 =>  #char "3";
+            case 4 =>  #char "4";
+            case 5 =>  #char "5";
+            case 6 =>  #char "6";
+            case 7 =>  #char "7";
+            case 8 =>  #char "8";
+            case 9 =>  #char "9";
+            case 10 => #char "A";
+            case 11 => #char "B";
+            case 12 => #char "C";
+            case 13 => #char "D";
+            case 14 => #char "E";
+            case 15 => #char "F";
+            case #default => #char "X";
+        };
     }
 }
 
diff --git a/core/net/net.onyx b/core/net/net.onyx
index ae8238bd..4c3cdf44 100644
--- a/core/net/net.onyx
+++ b/core/net/net.onyx
@@ -137,11 +137,11 @@ socket_is_alive :: (s: &Socket) -> bool {
 }
 
 socket_connect :: (s: &Socket, host: str, port: u16 = 0) -> SocketError {
-    switch s.family {
-        case .Inet do return __net_connect_ipv4(s.handle, host, port);
-        case .Unix do return __net_connect_unix(s.handle, host);
-        case #default do return .BadSettings;
-    }
+    return switch s.family {
+        case .Inet => __net_connect_ipv4(s.handle, host, port);
+        case .Unix => __net_connect_unix(s.handle, host);
+        case #default => .BadSettings;
+    };
 }
 
 socket_bind :: (s: &Socket, bind_address: &Socket_Address) -> bool {
diff --git a/core/onyx/cbindgen.onyx b/core/onyx/cbindgen.onyx
index b2678a88..b351145a 100644
--- a/core/onyx/cbindgen.onyx
+++ b/core/onyx/cbindgen.onyx
@@ -293,26 +293,26 @@ compile_c_file :: (
 
         param_info := get_type_info(t);
         switch param_info.kind {
-            case .Basic do switch t {
-                case bool do return "i32";
-                case i8   do return "i32";
-                case u8   do return "i32";
-                case i16  do return "i32";
-                case u16  do return "i32";
-                case i32  do return "i32";
-                case u32  do return "i32";
-                case i64  do return "i64";
-                case u64  do return "i64";
-
-                case f32  do return "f32";
-                case f64  do return "f64";
-
-                case rawptr do return "ptr"; 
+            case .Basic do return switch t {
+                case bool => "i32";
+                case i8   => "i32";
+                case u8   => "i32";
+                case i16  => "i32";
+                case u16  => "i32";
+                case i32  => "i32";
+                case u32  => "i32";
+                case i64  => "i64";
+                case u64  => "i64";
+
+                case f32  => "f32";
+                case f64  => "f64";
+
+                case rawptr => "ptr"; 
                 
-                case i8x16, i16x8, i32x4, i64x2, f32x4, f64x2, v128 do return "v128";
+                case i8x16, i16x8, i32x4, i64x2, f32x4, f64x2, v128 => "v128";
 
-                case type_expr do return "i32";
-            }
+                case type_expr => "i32";
+            };
 
             case .Pointer do return "ptr"; 
             case .Multi_Pointer do return "ptr"; 
@@ -350,26 +350,26 @@ compile_c_file :: (
 
         param_info := get_type_info(t);
         switch param_info.kind {
-            case .Basic do switch t {
-                case bool do return "WASM_I32";
-                case i8   do return "WASM_I32";
-                case u8   do return "WASM_I32";
-                case i16  do return "WASM_I32";
-                case u16  do return "WASM_I32";
-                case i32  do return "WASM_I32";
-                case u32  do return "WASM_I32";
-                case i64  do return "WASM_I64";
-                case u64  do return "WASM_I64";
-
-                case f32  do return "WASM_F32";
-                case f64  do return "WASM_F64";
-
-                case rawptr do return "WASM_I32"; // This will have to depend on the pointer size...
+            case .Basic do return switch t {
+                case bool => "WASM_I32";
+                case i8   => "WASM_I32";
+                case u8   => "WASM_I32";
+                case i16  => "WASM_I32";
+                case u16  => "WASM_I32";
+                case i32  => "WASM_I32";
+                case u32  => "WASM_I32";
+                case i64  => "WASM_I64";
+                case u64  => "WASM_I64";
+
+                case f32  => "WASM_F32";
+                case f64  => "WASM_F64";
+
+                case rawptr => "WASM_I32"; // This will have to depend on the pointer size...
                 
-                case i8x16, i16x8, i32x4, i64x2, f32x4, f64x2, v128 do return "WASM_V128";
+                case i8x16, i16x8, i32x4, i64x2, f32x4, f64x2, v128 => "WASM_V128";
 
-                case type_expr do return "WASM_I32";
-            }
+                case type_expr => "WASM_I32";
+            };
 
             case .Pointer do return "WASM_I32"; // This will also have to depend on the pointer size...
             case .Multi_Pointer do return "WASM_I32"; // This will also have to depend on the pointer size...
diff --git a/core/runtime/info/helper.onyx b/core/runtime/info/helper.onyx
index 74fb82ae..ac8a6c41 100644
--- a/core/runtime/info/helper.onyx
+++ b/core/runtime/info/helper.onyx
@@ -26,15 +26,15 @@ write_type_name :: (writer: &io.Writer, t: type_expr) {
                 case .F32  do io.write_str(writer, "f32");
                 case .F64  do io.write_str(writer, "f64");
 
-                case .Rawptr do return io.write_str(writer, "rawptr");
+                case .Rawptr do io.write_str(writer, "rawptr");
 
-                case .I8X16 do return io.write_str(writer, "i8x16");
-                case .I16X8 do return io.write_str(writer, "i16x8");
-                case .I32X4 do return io.write_str(writer, "i32x4");
-                case .I64X2 do return io.write_str(writer, "i64x2");
-                case .F32X4 do return io.write_str(writer, "f32x4");
-                case .F64X2 do return io.write_str(writer, "f64x2");
-                case .V128  do return io.write_str(writer, "v128");
+                case .I8X16 do io.write_str(writer, "i8x16");
+                case .I16X8 do io.write_str(writer, "i16x8");
+                case .I32X4 do io.write_str(writer, "i32x4");
+                case .I64X2 do io.write_str(writer, "i64x2");
+                case .F32X4 do io.write_str(writer, "f32x4");
+                case .F64X2 do io.write_str(writer, "f64x2");
+                case .V128  do io.write_str(writer, "v128");
 
                 case .Type_Index do io.write_str(writer, "type_expr");
                 case .Unsized_Int do io.write_str(writer, "<unsized int>");
@@ -154,57 +154,57 @@ size_of :: (t: type_expr) -> u32 {
     info := get_type_info(t);
     if info == null do return 0;
 
-    switch info.kind {
-        case .Basic {
+    return switch info.kind {
+        case .Basic => do {
             basic := cast(&Type_Info_Basic) info;
 
-            switch basic.basic_kind {
-                case .Void do return 0;
-                case .Bool, .U8, .I8 do return 1;
-                case .U16, .I16 do return 2;
-                case .U32, .I32, .F32, .Type_Index do return 4;
-                case .U64, .I64, .F64 do return 8;
-                case .I8X16, .I16X8, .I32X4, .I64X2, .F32X4, .F64X2, .V128 do return 16;
-                case .Rawptr do return sizeof rawptr;
-
-                case .Unsized_Int do return 0;
-                case .Unsized_Float do return 0;
-            }
+            return switch basic.basic_kind {
+                case .Void => 0;
+                case .Bool, .U8, .I8 => 1;
+                case .U16, .I16 => 2;
+                case .U32, .I32, .F32, .Type_Index => 4;
+                case .U64, .I64, .F64 => 8;
+                case .I8X16, .I16X8, .I32X4, .I64X2, .F32X4, .F64X2, .V128 => 16;
+                case .Rawptr => sizeof rawptr;
+
+                case .Unsized_Int => 0;
+                case .Unsized_Float => 0;
+            };
         }
 
-        case .Pointer do return sizeof rawptr;
+        case .Pointer => sizeof rawptr
 
-        case .Array {
+        case .Array => do {
             arr := cast(&Type_Info_Array) info;
             return size_of(arr.of) * arr.count;
         }
 
-        case .Slice do return sizeof str;
-        case .Dynamic_Array do return sizeof [..] void;
-        case .Variadic_Argument do return sizeof str;
-        case .Enum {
+        case .Slice => sizeof str;
+        case .Dynamic_Array => sizeof [..] void;
+        case .Variadic_Argument => sizeof str;
+        case .Enum => do {
             e := cast(&Type_Info_Enum) info;
             return e.size;
         }
 
-        case .Struct {
+        case .Struct => do {
             s := cast(&Type_Info_Struct) info;
             return s.size;
         }
 
-        case .Polymorphic_Struct do return 0;
+        case .Polymorphic_Struct => 0;
 
-        case .Compound do return 0;
+        case .Compound => 0;
 
-        case .Function do return 4;
+        case .Function => 4;
 
-        case .Distinct {
+        case .Distinct => do {
             d := cast(&Type_Info_Distinct) info;
             return size_of(d.base_type);
         }
-    }
 
-    return 0;
+        case #default => 0;
+    };
 }
 
 offset_of :: (T: type_expr, member_name: str) -> u32 {
diff --git a/tests/aoc-2020/day20.onyx b/tests/aoc-2020/day20.onyx
index 9ae73e7c..a4ed4858 100644
--- a/tests/aoc-2020/day20.onyx
+++ b/tests/aoc-2020/day20.onyx
@@ -180,18 +180,17 @@ apply_orientation :: (t: &Tile, ori: TO) {
 }
 
 index_square_with_orientation :: (data: [&] $T, ori: TO, size: i32, x: i32, y: i32) -> &T {
-    switch ori {
-        case TO.N     do return &data[x + y * size];
-        case TO.R90   do return &data[y + (size - 1 - x) * size];
-        case TO.R180  do return &data[(size - 1 - x) + (size - 1 - y) * size];
-        case TO.R270  do return &data[(size - 1 - y) + x * size];
-        case TO.F     do return &data[x + (size - 1 - y) * size];
-        case TO.FR90  do return &data[y + x * size];
-        case TO.FR180 do return &data[(size - 1 - x) + y * size];
-        case TO.FR270 do return &data[(size - 1 - y) + (size - 1 - x) * size];
-    }
-
-    return null;
+    return switch ori {
+        case TO.N     => &data[x + y * size];
+        case TO.R90   => &data[y + (size - 1 - x) * size];
+        case TO.R180  => &data[(size - 1 - x) + (size - 1 - y) * size];
+        case TO.R270  => &data[(size - 1 - y) + x * size];
+        case TO.F     => &data[x + (size - 1 - y) * size];
+        case TO.FR90  => &data[y + x * size];
+        case TO.FR180 => &data[(size - 1 - x) + y * size];
+        case TO.FR270 => &data[(size - 1 - y) + (size - 1 - x) * size];
+        case #default => null
+    };
 }
 
 sea_monster_width  := 20;
diff --git a/tests/switch_expressions b/tests/switch_expressions
new file mode 100644
index 00000000..d11c0e3b
--- /dev/null
+++ b/tests/switch_expressions
@@ -0,0 +1,5 @@
+Val3
+Val1
+123
+12.3400
+A string
diff --git a/tests/switch_expressions.onyx b/tests/switch_expressions.onyx
new file mode 100644
index 00000000..65f9e9d0
--- /dev/null
+++ b/tests/switch_expressions.onyx
@@ -0,0 +1,55 @@
+use core {*}
+
+Foo :: enum {
+    Val1;
+    Val2;
+    Val3;
+}
+
+assignment_test :: () {
+    v: Foo = switch 3 {
+        case 1 => .Val1;
+        case 2 => .Val2;
+        case 3 => .Val3;
+    };
+
+    printf("{}\n", v);
+}
+
+quick_map_test :: () {
+    println(quick_map(3));
+
+    quick_map :: (n: i32) -> Foo {
+        return switch n {
+            case 4 => .Val1;
+            case 5 => .Val2;
+            case 6 => .Val3;
+            case #default => .Val1;
+        };
+    }
+}
+
+quick_union_map :: () { 
+
+    quick_map(SomeUnion.{ x = 123 }) |> println();
+    quick_map(SomeUnion.{ y = 12.34 }) |> println();
+    quick_map(SomeUnion.{ z = "A string" }) |> println();
+
+    SomeUnion :: union {
+        x: i32;
+        y: f32;
+        z: str;
+    }
+
+    quick_map :: (v: $T) => switch v {
+        case val: .z => val;
+        case val: .y => conv.format("{}", val);
+        case val: .x => conv.format("{}", val);
+    }
+}
+
+main :: () {
+    assignment_test();
+    quick_map_test();
+    quick_union_map();
+}
\ No newline at end of file
diff --git a/tests/tagged_unions.onyx b/tests/tagged_unions.onyx
index c0ad0c3a..96218748 100644
--- a/tests/tagged_unions.onyx
+++ b/tests/tagged_unions.onyx
@@ -7,7 +7,7 @@ union_is :: macro (u: $U, $variant: U.tag_enum) -> bool {
 
 extract_variant :: macro (u: $U, $variant: U.tag_enum) => {
     switch u {
-        case variant => v {
+        case v: variant {
             return Optional.make(v);
         }
     }
@@ -44,15 +44,15 @@ extraction_example :: () {
     value := Extraction.{ string = "This works" };
 
     switch value {
-        case .int => int_value {
+        case int_value: .int {
             printf("This is an integer: {}\n", int_value);
         }
 
-        case .float => float_value {
+        case float_value: .float {
             printf("This is a float: {}\n", float_value);
         }
 
-        case .string => string_value {
+        case string_value: .string {
             printf("This is a string: {}\n", string_value);
         }
     }
@@ -66,15 +66,15 @@ method_example :: () {
 
         do_the_thing :: (value: &Methoded) {
             switch *value {
-                case .int => int_value {
+                case int_value: .int {
                     printf("This is an integer: {}\n", int_value);
                 }
 
-                case .float => float_value {
+                case float_value: .float {
                     printf("This is a float: {}\n", float_value);
                 }
 
-                case .string => string_value {
+                case string_value: .string {
                     printf("This is a string: {}\n", string_value);
                 }
             }
@@ -100,7 +100,7 @@ linked_list_example :: () {
             switch walker {
                 case .End do break break;
 
-                case .Next => &next {
+                case &next: .Next {
                     printf("{}\n", next.data);
                     walker = next.next;
                 }
-- 
2.25.1