lower if (x) ... else |e| switch (e) to switch_block_err_union

dweiller · dweiller · commit ad85f9081245 · 2023-12-05T18:13:54.000+11:00
diff --git a/src/AstGen.zig b/src/AstGen.zig
@@ -838,7 +838,18 @@ fn expr(gz: *GenZir, scope: *Scope, ri: ResultInfo, node: Ast.Node.Index) InnerE
 
         .if_simple,
         .@"if",
-        => return ifExpr(gz, scope, ri.br(), node, tree.fullIf(node).?),
+        => {
+            const if_full = tree.fullIf(node).?;
+            if (if_full.error_token) |error_token| {
+                const tag = node_tags[if_full.ast.else_expr];
+                if ((tag == .@"switch" or tag == .switch_comma) and
+                    std.mem.eql(u8, tree.tokenSlice(error_token), tree.tokenSlice(error_token + 4)))
+                {
+                    return switchExprErrUnion(gz, scope, ri.br(), node, .@"if");
+                }
+            }
+            return ifExpr(gz, scope, ri.br(), node, if_full);
+        },
 
         .while_simple,
         .while_cont,
@@ -1019,7 +1030,7 @@ fn expr(gz: *GenZir, scope: *Scope, ri: ResultInfo, node: Ast.Node.Index) InnerE
                     token_tags[catch_token + 4] == .keyword_switch)
                 {
                     if (std.mem.eql(u8, tree.tokenSlice(catch_token + 2), tree.tokenSlice(catch_token + 6))) {
-                        return switchExprErrUnion(gz, scope, ri.br(), node);
+                        return switchExprErrUnion(gz, scope, ri.br(), node, .@"catch");
                     }
                 }
                 break :blk catch_token + 2;
@@ -6860,7 +6871,8 @@ fn switchExprErrUnion(
     parent_gz: *GenZir,
     scope: *Scope,
     ri: ResultInfo,
-    catch_node: Ast.Node.Index,
+    catch_or_if_node: Ast.Node.Index,
+    node_ty: enum { @"catch", @"if" },
 ) InnerError!Zir.Inst.Ref {
     const astgen = parent_gz.astgen;
     const gpa = astgen.gpa;
@@ -6869,21 +6881,42 @@ fn switchExprErrUnion(
     const node_tags = tree.nodes.items(.tag);
     const main_tokens = tree.nodes.items(.main_token);
     const token_tags = tree.tokens.items(.tag);
-    const operand_node = node_datas[catch_node].lhs;
-    const switch_node = node_datas[catch_node].rhs;
+
+    const if_full = switch (node_ty) {
+        .@"catch" => undefined,
+        .@"if" => tree.fullIf(catch_or_if_node).?,
+    };
+
+    const switch_node, const operand_node, const error_payload = switch (node_ty) {
+        .@"catch" => .{
+            node_datas[catch_or_if_node].rhs,
+            node_datas[catch_or_if_node].lhs,
+            main_tokens[catch_or_if_node] + 2,
+        },
+        .@"if" => .{
+            if_full.ast.else_expr,
+            if_full.ast.cond_expr,
+            if_full.error_token.?,
+        },
+    };
+    assert(node_tags[switch_node] == .@"switch" or node_tags[switch_node] == .switch_comma);
+
     const extra = tree.extraData(node_datas[switch_node].rhs, Ast.Node.SubRange);
     const case_nodes = tree.extra_data[extra.start..extra.end];
 
-    const need_rl = astgen.nodes_need_rl.contains(catch_node);
+    const need_rl = astgen.nodes_need_rl.contains(catch_or_if_node);
     const block_ri: ResultInfo = if (need_rl) ri else .{
         .rl = switch (ri.rl) {
-            .ptr => .{ .ty = (try ri.rl.resultType(parent_gz, catch_node)).? },
+            .ptr => .{ .ty = (try ri.rl.resultType(parent_gz, catch_or_if_node)).? },
             .inferred_ptr => .none,
             else => ri.rl,
         },
         .ctx = ri.ctx,
     };
 
+    const payload_is_ref = node_ty == .@"if" and
+        if_full.payload_token != null and token_tags[if_full.payload_token.?] == .asterisk;
+
     // We need to call `rvalue` to write through to the pointer only if we had a
     // result pointer and aren't forwarding it.
     const LocTag = @typeInfo(ResultInfo.Loc).Union.tag_type.?;
@@ -6951,12 +6984,15 @@ fn switchExprErrUnion(
         }
     }
 
-    const operand_ri: ResultInfo = .{ .rl = .none, .ctx = .error_handling_expr };
+    const operand_ri: ResultInfo = .{
+        .rl = if (payload_is_ref) .ref else .none,
+        .ctx = .error_handling_expr,
+    };
 
     astgen.advanceSourceCursorToNode(operand_node);
     const operand_lc = LineColumn{ astgen.source_line - parent_gz.decl_line, astgen.source_column };
 
-    const raw_operand = try reachableExpr(parent_gz, scope, operand_ri, operand_node, node_datas[catch_node].rhs);
+    const raw_operand = try reachableExpr(parent_gz, scope, operand_ri, operand_node, switch_node);
     const item_ri: ResultInfo = .{ .rl = .none };
 
     // This contains the data that goes into the `extra` array for the SwitchBlockErrUnion, except
@@ -6997,13 +7033,93 @@ fn switchExprErrUnion(
 
         try case_scope.addDbgBlockBegin();
 
-        const unwrapped_payload = try case_scope.addUnNode(.err_union_payload_unsafe, raw_operand, catch_node);
-        const case_result = switch (ri.rl) {
-            .ref, .ref_coerced_ty => unwrapped_payload,
-            else => try rvalue(&case_scope, block_scope.break_result_info, unwrapped_payload, catch_node),
-        };
-        try case_scope.addDbgBlockEnd();
-        _ = try case_scope.addBreakWithSrcNode(.@"break", switch_block, case_result, catch_node);
+        const unwrap_payload_tag: Zir.Inst.Tag = if (payload_is_ref)
+            .err_union_payload_unsafe_ptr
+        else
+            .err_union_payload_unsafe;
+
+        const unwrapped_payload = try case_scope.addUnNode(
+            unwrap_payload_tag,
+            raw_operand,
+            catch_or_if_node,
+        );
+
+        switch (node_ty) {
+            .@"catch" => {
+                const case_result = switch (ri.rl) {
+                    .ref, .ref_coerced_ty => unwrapped_payload,
+                    else => try rvalue(
+                        &case_scope,
+                        block_scope.break_result_info,
+                        unwrapped_payload,
+                        catch_or_if_node,
+                    ),
+                };
+                try case_scope.addDbgBlockEnd();
+                _ = try case_scope.addBreakWithSrcNode(
+                    .@"break",
+                    switch_block,
+                    case_result,
+                    catch_or_if_node,
+                );
+            },
+            .@"if" => {
+                var payload_val_scope: Scope.LocalVal = undefined;
+
+                try case_scope.addDbgBlockBegin();
+                const then_node = if_full.ast.then_expr;
+                const then_sub_scope = s: {
+                    assert(if_full.error_token != null);
+                    if (if_full.payload_token) |payload_token| {
+                        const token_name_index = payload_token + @intFromBool(payload_is_ref);
+                        const ident_name = try astgen.identAsString(token_name_index);
+                        const token_name_str = tree.tokenSlice(token_name_index);
+                        if (mem.eql(u8, "_", token_name_str))
+                            break :s &case_scope.base;
+                        try astgen.detectLocalShadowing(
+                            &case_scope.base,
+                            ident_name,
+                            token_name_index,
+                            token_name_str,
+                            .capture,
+                        );
+                        payload_val_scope = .{
+                            .parent = &case_scope.base,
+                            .gen_zir = &case_scope,
+                            .name = ident_name,
+                            .inst = unwrapped_payload,
+                            .token_src = payload_token,
+                            .id_cat = .capture,
+                        };
+                        try case_scope.addDbgVar(.dbg_var_val, ident_name, unwrapped_payload);
+                        break :s &payload_val_scope.base;
+                    } else {
+                        _ = try case_scope.addUnNode(
+                            .ensure_err_union_payload_void,
+                            raw_operand,
+                            catch_or_if_node,
+                        );
+                        break :s &case_scope.base;
+                    }
+                };
+                const then_result = try expr(
+                    &case_scope,
+                    then_sub_scope,
+                    block_scope.break_result_info,
+                    then_node,
+                );
+                try checkUsed(parent_gz, &case_scope.base, then_sub_scope);
+                if (!case_scope.endsWithNoReturn()) {
+                    try case_scope.addDbgBlockEnd();
+                    _ = try case_scope.addBreakWithSrcNode(
+                        .@"break",
+                        switch_block,
+                        then_result,
+                        then_node,
+                    );
+                }
+            },
+        }
 
         const case_slice = case_scope.instructionsSlice();
         // Since we use the switch_block_err_union instruction itself to refer
@@ -7020,9 +7136,18 @@ fn switchExprErrUnion(
         };
         const body_len = refs_len + astgen.countBodyLenAfterFixups(case_slice);
         try payloads.ensureUnusedCapacity(gpa, body_len);
+        const capture: Zir.Inst.SwitchBlock.ProngInfo.Capture = switch (node_ty) {
+            .@"catch" => .none,
+            .@"if" => if (if_full.payload_token == null)
+                .none
+            else if (payload_is_ref)
+                .by_ref
+            else
+                .by_val,
+        };
         payloads.items[body_len_index] = @bitCast(Zir.Inst.SwitchBlock.ProngInfo{
             .body_len = @intCast(body_len),
-            .capture = .none,
+            .capture = capture,
             .is_inline = false,
             .has_tag_capture = false,
         });
@@ -7032,16 +7157,15 @@ fn switchExprErrUnion(
         appendBodyWithFixupsArrayList(astgen, payloads, case_slice);
     }
 
-    const err_name, const error_payload = blk: {
-        const error_payload = main_tokens[catch_node] + 2;
+    const err_name = blk: {
         const err_str = tree.tokenSlice(error_payload);
         if (mem.eql(u8, err_str, "_")) {
             return astgen.failTok(error_payload, "discard of error capture; omit it instead", .{});
         }
         const err_name = try astgen.identAsString(error_payload);
         try astgen.detectLocalShadowing(scope, err_name, error_payload, err_str, .capture);
 
-        break :blk .{ err_name, error_payload };
+        break :blk err_name;
     };
 
     // allocate a shared dummy instruction for the error capture
@@ -7234,6 +7358,7 @@ fn switchExprErrUnion(
             .has_else = has_else,
             .scalar_cases_len = @intCast(scalar_cases_len),
             .any_uses_err_capture = any_uses_err_capture,
+            .payload_is_ref = payload_is_ref,
         },
     });
 
diff --git a/src/Sema.zig b/src/Sema.zig
@@ -8908,10 +8908,14 @@ fn zirErrUnionCodePtr(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileE
     const tracy = trace(@src());
     defer tracy.end();
 
-    const mod = sema.mod;
     const inst_data = sema.code.instructions.items(.data)[@intFromEnum(inst)].un_node;
     const src = inst_data.src();
     const operand = try sema.resolveInst(inst_data.operand);
+    return sema.analyzeErrUnionCodePtr(block, src, operand);
+}
+
+fn analyzeErrUnionCodePtr(sema: *Sema, block: *Block, src: LazySrcLoc, operand: Air.Inst.Ref) CompileError!Air.Inst.Ref {
+    const mod = sema.mod;
     const operand_ty = sema.typeOf(operand);
     assert(operand_ty.zigTypeTag(mod) == .Pointer);
 
@@ -8926,7 +8930,10 @@ fn zirErrUnionCodePtr(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileE
     if (try sema.resolveDefinedValue(block, src, operand)) |pointer_val| {
         if (try sema.pointerDeref(block, src, pointer_val, operand_ty)) |val| {
             assert(val.getErrorName(mod) != .none);
-            return Air.internedToRef(val.toIntern());
+            return Air.internedToRef((try mod.intern(.{ .err = .{
+                .ty = result_ty.toIntern(),
+                .name = mod.intern_pool.indexToKey(val.toIntern()).error_union.val.err_name,
+            } })));
         }
     }
 
@@ -11144,7 +11151,6 @@ fn zirSwitchBlockErrUnion(sema: *Sema, block: *Block, inst: Zir.Inst.Index) Comp
     const extra = sema.code.extraData(Zir.Inst.SwitchBlockErrUnion, inst_data.payload_index);
 
     const raw_operand_val = try sema.resolveInst(extra.data.operand);
-    assert(sema.typeOf(raw_operand_val).zigTypeTag(mod) == .ErrorUnion);
 
     // AstGen guarantees that the instruction immediately preceding
     // switch_block_err_union is a dbg_stmt
@@ -11175,6 +11181,7 @@ fn zirSwitchBlockErrUnion(sema: *Sema, block: *Block, inst: Zir.Inst.Index) Comp
     const NonError = struct {
         body: []const Zir.Inst.Index,
         end: usize,
+        capture: Zir.Inst.SwitchBlock.ProngInfo.Capture,
     };
 
     const non_error_case: NonError = non_error: {
@@ -11183,6 +11190,7 @@ fn zirSwitchBlockErrUnion(sema: *Sema, block: *Block, inst: Zir.Inst.Index) Comp
         break :non_error .{
             .body = sema.code.bodySlice(extra_body_start, info.body_len),
             .end = extra_body_start + info.body_len,
+            .capture = info.capture,
         };
     };
 
@@ -11207,7 +11215,7 @@ fn zirSwitchBlockErrUnion(sema: *Sema, block: *Block, inst: Zir.Inst.Index) Comp
             .body = sema.code.bodySlice(extra_body_start, info.body_len),
             .end = extra_body_start + info.body_len,
             .is_inline = info.is_inline,
-            .has_capture = info.capture == .by_val,
+            .has_capture = info.capture != .none,
         };
     };
 
@@ -11217,7 +11225,10 @@ fn zirSwitchBlockErrUnion(sema: *Sema, block: *Block, inst: Zir.Inst.Index) Comp
     }
 
     const operand_ty = sema.typeOf(raw_operand_val);
-    const operand_err_set_ty = operand_ty.errorUnionSet(mod);
+    const operand_err_set_ty = if (extra.data.bits.payload_is_ref)
+        operand_ty.childType(mod).errorUnionSet(mod)
+    else
+        operand_ty.errorUnionSet(mod);
 
     const block_inst: Air.Inst.Index = @enumFromInt(sema.air_instructions.len);
     try sema.air_instructions.append(gpa, .{
@@ -11285,7 +11296,12 @@ fn zirSwitchBlockErrUnion(sema: *Sema, block: *Block, inst: Zir.Inst.Index) Comp
         .tag_capture_inst = undefined,
     };
 
-    if (try sema.resolveDefinedValue(&child_block, src, raw_operand_val)) |operand_val| {
+    if (try sema.resolveDefinedValue(&child_block, src, raw_operand_val)) |ov| {
+        const operand_val = if (extra.data.bits.payload_is_ref)
+            (try sema.pointerDeref(&child_block, src, ov, operand_ty)).?
+        else
+            ov;
+
         if (operand_val.errorUnionIsPayload(mod)) {
             return sema.resolveBlockBody(block, operand_src, &child_block, non_error_case.body, inst, merges);
         } else {
@@ -11295,7 +11311,10 @@ fn zirSwitchBlockErrUnion(sema: *Sema, block: *Block, inst: Zir.Inst.Index) Comp
                     .name = operand_val.getErrorName(mod).unwrap().?,
                 },
             }));
-            spa.operand = try sema.analyzeErrUnionCode(block, operand_src, raw_operand_val);
+            spa.operand = if (extra.data.bits.payload_is_ref)
+                try sema.analyzeErrUnionCodePtr(block, operand_src, raw_operand_val)
+            else
+                try sema.analyzeErrUnionCode(block, operand_src, raw_operand_val);
 
             if (extra.data.bits.any_uses_err_capture) {
                 sema.inst_map.putAssumeCapacity(err_capture_inst, spa.operand);
@@ -11339,7 +11358,14 @@ fn zirSwitchBlockErrUnion(sema: *Sema, block: *Block, inst: Zir.Inst.Index) Comp
         unreachable;
     }
 
-    const cond = try sema.analyzeIsNonErr(block, src, raw_operand_val);
+    const cond = if (extra.data.bits.payload_is_ref) blk: {
+        try sema.checkErrorType(block, src, sema.typeOf(raw_operand_val).elemType2(mod));
+        const loaded = try sema.analyzeLoad(block, src, raw_operand_val, src);
+        break :blk try sema.analyzeIsNonErr(block, src, loaded);
+    } else blk: {
+        try sema.checkErrorType(block, src, sema.typeOf(raw_operand_val));
+        break :blk try sema.analyzeIsNonErr(block, src, raw_operand_val);
+    };
 
     var sub_block = child_block.makeSubBlock();
     sub_block.runtime_loop = null;
@@ -11351,7 +11377,11 @@ fn zirSwitchBlockErrUnion(sema: *Sema, block: *Block, inst: Zir.Inst.Index) Comp
     const true_instructions = try sub_block.instructions.toOwnedSlice(gpa);
     defer gpa.free(true_instructions);
 
-    spa.operand = try sema.analyzeErrUnionCode(&sub_block, operand_src, raw_operand_val);
+    spa.operand = if (extra.data.bits.payload_is_ref)
+        try sema.analyzeErrUnionCodePtr(&sub_block, operand_src, raw_operand_val)
+    else
+        try sema.analyzeErrUnionCode(&sub_block, operand_src, raw_operand_val);
+
     if (extra.data.bits.any_uses_err_capture) {
         sema.inst_map.putAssumeCapacity(err_capture_inst, spa.operand);
     }
diff --git a/src/Zir.zig b/src/Zir.zig
@@ -2794,9 +2794,10 @@ pub const Inst = struct {
             /// If true, there is an else prong. This is mutually exclusive with `has_under`.
             has_else: bool,
             any_uses_err_capture: bool,
+            payload_is_ref: bool,
             scalar_cases_len: ScalarCasesLen,
 
-            pub const ScalarCasesLen = u29;
+            pub const ScalarCasesLen = u28;
         };
 
         pub const MultiProng = struct {
diff --git a/test/behavior/switch_on_captured_error.zig b/test/behavior/switch_on_captured_error.zig
