From 5ff518fbb9ea2fb5a745841731912acbe2f046d9 Mon Sep 17 00:00:00 2001
From: Veikka Tuominen <git@vexu.eu>
Date: Fri, 25 Mar 2022 13:21:22 +0200
Subject: [PATCH] Sema: implement zirSwitchCapture for error sets

---
 src/Sema.zig             | 32 ++++++++++++++++++++++++---
 src/Zir.zig              | 47 ++++++++++++++++++++++++++++++++++++++++
 test/behavior/switch.zig |  5 ++++-
 3 files changed, 80 insertions(+), 4 deletions(-)

diff --git a/src/Sema.zig b/src/Sema.zig
index 2e7bd06ada..1896ee40df 100644
--- a/src/Sema.zig
+++ b/src/Sema.zig
@@ -6938,7 +6938,6 @@ fn zirSwitchCapture(
     const switch_info = zir_datas[capture_info.switch_inst].pl_node;
     const switch_extra = sema.code.extraData(Zir.Inst.SwitchBlock, switch_info.payload_index);
     const operand_src: LazySrcLoc = .{ .node_offset_switch_operand = switch_info.src_node };
-    const switch_src = switch_info.src();
     const operand_is_ref = switch_extra.data.bits.is_ref;
     const cond_inst = Zir.refToIndex(switch_extra.data.operand).?;
     const cond_info = sema.code.instructions.items(.data)[cond_inst].un_node;
@@ -6965,7 +6964,29 @@ fn zirSwitchCapture(
     }
 
     if (is_multi) {
-        return sema.fail(block, switch_src, "TODO implement Sema for switch capture multi", .{});
+        const items = switch_extra.data.getMultiProng(sema.code, switch_extra.end, capture_info.prong_index).items;
+
+        var names: Module.ErrorSet.NameMap = .{};
+        try names.ensureUnusedCapacity(sema.arena, items.len);
+        for (items) |item| {
+            const item_ref = sema.resolveInst(item);
+            // Previous switch validation ensured this will succeed
+            const item_val = sema.resolveConstValue(block, .unneeded, item_ref) catch unreachable;
+            names.putAssumeCapacityNoClobber(
+                item_val.getError().?,
+                {},
+            );
+        }
+
+        // names must be sorted
+        Module.ErrorSet.sortNames(&names);
+        const else_error_ty = try Type.Tag.error_set_merged.create(sema.arena, names);
+
+        const operand = if (operand_is_ref)
+            try sema.analyzeLoad(block, operand_src, operand_ptr, operand_src)
+        else
+            operand_ptr;
+        return sema.bitCast(block, else_error_ty, operand, operand_src);
     }
     const scalar_prong = switch_extra.data.getScalarProng(sema.code, switch_extra.end, capture_info.prong_index);
     const item = sema.resolveInst(scalar_prong.item);
@@ -7022,7 +7043,12 @@ fn zirSwitchCapture(
             return block.addStructFieldVal(operand, field_index, field.ty);
         },
         .ErrorSet => {
-            return sema.fail(block, operand_src, "TODO implement Sema for zirSwitchCapture for error sets", .{});
+            const item_ty = try Type.Tag.error_set_single.create(sema.arena, item_val.getError().?);
+            const operand = if (operand_is_ref)
+                try sema.analyzeLoad(block, operand_src, operand_ptr, operand_src)
+            else
+                operand_ptr;
+            return sema.bitCast(block, item_ty, operand, operand_src);
         },
         else => {
             return sema.fail(block, operand_src, "switch on type '{}' provides no capture value", .{
diff --git a/src/Zir.zig b/src/Zir.zig
index f9b80c88ef..d6e8486ea7 100644
--- a/src/Zir.zig
+++ b/src/Zir.zig
@@ -2620,6 +2620,53 @@ pub const Inst = struct {
                 };
             }
         }
+
+        pub const MultiProng = struct {
+            items: []const Ref,
+            body: []const Index,
+        };
+
+        pub fn getMultiProng(
+            self: SwitchBlock,
+            zir: Zir,
+            extra_end: usize,
+            prong_index: usize,
+        ) MultiProng {
+            // +1 for self.bits.has_multi_cases == true
+            var extra_index: usize = extra_end + 1;
+
+            if (self.bits.specialProng() != .none) {
+                const body_len = zir.extra[extra_index];
+                extra_index += 1;
+                const body = zir.extra[extra_index..][0..body_len];
+                extra_index += body.len;
+            }
+
+            var scalar_i: usize = 0;
+            while (scalar_i < self.bits.scalar_cases_len) : (scalar_i += 1) {
+                extra_index += 1;
+                const body_len = zir.extra[extra_index];
+                extra_index += 1;
+                extra_index += body_len;
+            }
+            var multi_i: u32 = 0;
+            while (true) : (multi_i += 1) {
+                const items_len = zir.extra[extra_index];
+                extra_index += 2;
+                const body_len = zir.extra[extra_index];
+                extra_index += 1;
+                const items = zir.refSlice(extra_index, items_len);
+                extra_index += items_len;
+                const body = zir.extra[extra_index..][0..body_len];
+                extra_index += body_len;
+
+                if (multi_i < prong_index) continue;
+                return .{
+                    .items = items,
+                    .body = body,
+                };
+            }
+        }
     };
 
     pub const Field = struct {
diff --git a/test/behavior/switch.zig b/test/behavior/switch.zig
index db670e34a7..9634a4de37 100644
--- a/test/behavior/switch.zig
+++ b/test/behavior/switch.zig
@@ -465,7 +465,10 @@ test "else prong of switch on error set excludes other cases" {
 }
 
 test "switch prongs with error set cases make a new error set type for capture value" {
-    if (builtin.zig_backend != .stage1) return error.SkipZigTest; // TODO
+    if (builtin.zig_backend == .stage2_x86_64) return error.SkipZigTest; // TODO
+    if (builtin.zig_backend == .stage2_arm) return error.SkipZigTest; // TODO
+    if (builtin.zig_backend == .stage2_aarch64) return error.SkipZigTest; // TODO
+    if (builtin.zig_backend == .stage2_c) return error.SkipZigTest; // TODO
 
     const S = struct {
         fn doTheTest() !void {