IR: support inline switch

5 files changed, 68 insertions, 11 deletions

diff --git a/doc/langref.md b/doc/langref.md
index c209b3dbfc..4844a25f31 100644
--- a/doc/langref.md
+++ b/doc/langref.md
@@ -73,7 +73,7 @@ AssignmentOperator = "=" | "*=" | "/=" | "%=" | "+=" | "-=" | "<<=" | ">>=" | "&
 
 BlockExpression = IfExpression | Block | WhileExpression | ForExpression | SwitchExpression
 
-SwitchExpression = "switch" "(" Expression ")" "{" many(SwitchProng) "}"
+SwitchExpression = option("inline") "switch" "(" Expression ")" "{" many(SwitchProng) "}"
 
 SwitchProng = (list(SwitchItem, ",") | "else") "=>" option("|" option("*") Symbol "|") Expression ","
 
diff --git a/src/all_types.hpp b/src/all_types.hpp
index 223842261c..7c585175e5 100644
--- a/src/all_types.hpp
+++ b/src/all_types.hpp
@@ -557,6 +557,7 @@ struct AstNodeForExpr {
 struct AstNodeSwitchExpr {
     AstNode *expr;
     ZigList<AstNode *> prongs;
+    bool is_inline;
 
     // populated by semantic analyzer
     Expr resolved_expr;
diff --git a/src/ir.cpp b/src/ir.cpp
index f380077037..71f70391e0 100644
--- a/src/ir.cpp
+++ b/src/ir.cpp
@@ -2264,7 +2264,7 @@ static IrInstruction *ir_gen_switch_expr(IrBuilder *irb, AstNode *node) {
 
     size_t prong_count = node->data.switch_expr.prongs.length;
     ZigList<IrInstructionSwitchBrCase> cases = {0};
-    bool is_inline = (node->block_context->fn_entry == nullptr);
+    bool is_inline = node->data.switch_expr.is_inline || (node->block_context->fn_entry == nullptr);
 
     ZigList<IrInstruction *> incoming_values = {0};
     ZigList<IrBasicBlock *> incoming_blocks = {0};
@@ -5249,7 +5249,37 @@ static TypeTableEntry *ir_analyze_instruction_switch_br(IrAnalyze *ira,
     bool is_inline = switch_br_instruction->is_inline;
 
     if (is_inline || target_value->static_value.special != ConstValSpecialRuntime) {
-        zig_panic("TODO compile time switch br");
+        ConstExprValue *target_val = ir_resolve_const(ira, target_value);
+        if (!target_val)
+            return ir_finish_anal(ira, ira->codegen->builtin_types.entry_unreachable);
+
+        for (size_t i = 0; i < case_count; i += 1) {
+            IrInstructionSwitchBrCase *old_case = &switch_br_instruction->cases[i];
+            IrInstruction *case_value = old_case->value->other;
+            if (case_value->type_entry->id == TypeTableEntryIdInvalid)
+                return ir_finish_anal(ira, ira->codegen->builtin_types.entry_unreachable);
+
+            IrInstruction *casted_case_value = ir_get_casted_value(ira, case_value, target_value->type_entry);
+            if (casted_case_value->type_entry->id == TypeTableEntryIdInvalid)
+                return ir_finish_anal(ira, ira->codegen->builtin_types.entry_unreachable);
+
+            ConstExprValue *case_val = ir_resolve_const(ira, casted_case_value);
+            if (!case_val)
+                return ir_finish_anal(ira, ira->codegen->builtin_types.entry_unreachable);
+
+            if (const_values_equal(target_val, case_val, target_value->type_entry)) {
+                IrBasicBlock *old_dest_block = old_case->block;
+                if (is_inline || old_dest_block->ref_count == 1) {
+                    ir_inline_bb(ira, old_dest_block);
+                    return ira->codegen->builtin_types.entry_unreachable;
+                } else {
+                    IrBasicBlock *new_dest_block = ir_get_new_bb(ira, old_dest_block);
+                    ir_build_br_from(&ira->new_irb, &switch_br_instruction->base, new_dest_block);
+                    return ir_finish_anal(ira, ira->codegen->builtin_types.entry_unreachable);
+                }
+            }
+        }
+
     }
 
     IrInstructionSwitchBrCase *cases = allocate<IrInstructionSwitchBrCase>(case_count);
diff --git a/src/parser.cpp b/src/parser.cpp
index 8713c6d529..c316dada4c 100644
--- a/src/parser.cpp
+++ b/src/parser.cpp
@@ -1642,23 +1642,36 @@ static AstNode *ast_parse_for_expr(ParseContext *pc, size_t *token_index, bool m
 }
 
 /*
-SwitchExpression : "switch" "(" Expression ")" "{" many(SwitchProng) "}"
+SwitchExpression = option("inline") "switch" "(" Expression ")" "{" many(SwitchProng) "}"
 SwitchProng = (list(SwitchItem, ",") | "else") "=>" option("|" option("*") Symbol "|") Expression ","
 SwitchItem : Expression | (Expression "..." Expression)
 */
 static AstNode *ast_parse_switch_expr(ParseContext *pc, size_t *token_index, bool mandatory) {
-    Token *token = &pc->tokens->at(*token_index);
-
-    if (token->id != TokenIdKeywordSwitch) {
-        if (mandatory) {
-            ast_expect_token(pc, token, TokenIdKeywordSwitch);
+    Token *first_token = &pc->tokens->at(*token_index);
+    Token *switch_token;
+    bool is_inline;
+    if (first_token->id == TokenIdKeywordInline) {
+        is_inline = true;
+        switch_token = &pc->tokens->at(*token_index + 1);
+        if (switch_token->id == TokenIdKeywordSwitch) {
+            *token_index += 2;
+        } else if (mandatory) {
+            ast_expect_token(pc, first_token, TokenIdKeywordSwitch);
         } else {
             return nullptr;
         }
+    } else if (first_token->id == TokenIdKeywordSwitch) {
+        is_inline = false;
+        switch_token = first_token;
+        *token_index += 1;
+    } else if (mandatory) {
+        ast_expect_token(pc, first_token, TokenIdKeywordSwitch);
+    } else {
+        return nullptr;
     }
-    *token_index += 1;
 
-    AstNode *node = ast_create_node(pc, NodeTypeSwitchExpr, token);
+    AstNode *node = ast_create_node(pc, NodeTypeSwitchExpr, switch_token);
+    node->data.switch_expr.is_inline = is_inline;
 
     ast_eat_token(pc, token_index, TokenIdLParen);
     node->data.switch_expr.expr = ast_parse_expression(pc, token_index, true);
diff --git a/test/self_hosted2.zig b/test/self_hosted2.zig
index d3dfb6811f..0ed6bca373 100644
--- a/test/self_hosted2.zig
+++ b/test/self_hosted2.zig
@@ -49,6 +49,18 @@ fn testSwitchWithAllRanges(x: u32, y: u32) -> u32 {
     }
 }
 
+fn testInlineSwitch() {
+    const x = 3 + 4;
+    const result = inline switch (x) {
+        3 => 10,
+        4 => 11,
+        5, 6 => 12,
+        7, 8 => 13,
+        else => 14,
+    };
+    assert(result + 1 == 14);
+}
+
 fn assert(ok: bool) {
     if (!ok)
         @unreachable();
@@ -60,6 +72,7 @@ fn runAllTests() {
     inlinedLoop();
     switchWithNumbers();
     switchWithAllRanges();
+    testInlineSwitch();
 }
 
 export nakedcc fn _start() -> unreachable {