stage2: "Pop" error trace for break/return within catch

This implement trace "popping" for correctly handled errors within
`catch { ... }` and `else { ... }` blocks.

When breaking from these blocks with any non-error, we pop the error
trace frames corresponding to the operand. When breaking with an error,
we preserve the frames so that error traces "chain" together as usual.

```zig
fn foo(cond1: bool, cond2: bool) !void {
    bar() catch {
    	if (cond1) {
	    // If baz() result is a non-error, pop the error trace frames from bar()
	    // If baz() result is an error, leave the bar() frames on the error trace
            return baz();
	} else if (cond2) {
	    // If we break/return an error, then leave the error frames from bar() on the error trace
	    return error.Foo;
	}
    };

    // An error returned from here does not include bar()'s error frames in the trace
    return error.Bar;
}
```

Notice that if foo() does not return an error it, it leaves no extra
frames on the error trace.

This is piece (1/3) of https://github.com/ziglang/zig/issues/1923#issuecomment-1218495574

2 files changed, 116 insertions, 29 deletions

diff --git a/src/AstGen.zig b/src/AstGen.zig
index 2ef6a0aedc..dbd65519b0 100644
--- a/src/AstGen.zig
+++ b/src/AstGen.zig
@@ -1834,6 +1834,13 @@ fn breakExpr(parent_gz: *GenZir, parent_scope: *Scope, node: Ast.Node.Index) Inn
     const break_label = node_datas[node].lhs;
     const rhs = node_datas[node].rhs;
 
+    // Breaking out of a `catch { ... }` or `else |err| { ... }` block with a non-error value
+    // means that the corresponding error was correctly handled, and the error trace index
+    // needs to be restored so that any entries from the caught error are effectively "popped"
+    //
+    // Note: We only restore for the outermost block, since that will "pop" any nested blocks.
+    var err_trace_index_to_restore: Zir.Inst.Ref = .none;
+
     // Look for the label in the scope.
     var scope = parent_scope;
     while (true) {
@@ -1842,6 +1849,7 @@ fn breakExpr(parent_gz: *GenZir, parent_scope: *Scope, node: Ast.Node.Index) Inn
                 const block_gz = scope.cast(GenZir).?;
 
                 if (block_gz.cur_defer_node != 0) {
+                    // We are breaking out of a `defer` block.
                     return astgen.failNodeNotes(node, "cannot break out of defer expression", .{}, &.{
                         try astgen.errNoteNode(
                             block_gz.cur_defer_node,
@@ -1851,6 +1859,11 @@ fn breakExpr(parent_gz: *GenZir, parent_scope: *Scope, node: Ast.Node.Index) Inn
                     });
                 }
 
+                if (block_gz.saved_err_trace_index != .none) {
+                    // We are breaking out of a `catch { ... }` or `else |err| { ... }`.
+                    err_trace_index_to_restore = block_gz.saved_err_trace_index;
+                }
+
                 const block_inst = blk: {
                     if (break_label != 0) {
                         if (block_gz.label) |*label| {
@@ -1862,9 +1875,11 @@ fn breakExpr(parent_gz: *GenZir, parent_scope: *Scope, node: Ast.Node.Index) Inn
                     } else if (block_gz.break_block != 0) {
                         break :blk block_gz.break_block;
                     }
+                    // If not the target, start over with the parent
                     scope = block_gz.parent;
                     continue;
                 };
+                // If we made it here, this block is the target of the break expr
 
                 const break_tag: Zir.Inst.Tag = if (block_gz.is_inline or block_gz.force_comptime)
                     .break_inline
@@ -1874,6 +1889,19 @@ fn breakExpr(parent_gz: *GenZir, parent_scope: *Scope, node: Ast.Node.Index) Inn
                 if (rhs == 0) {
                     try genDefers(parent_gz, scope, parent_scope, .normal_only);
 
+                    // As our last action before the break, "pop" the error trace if needed
+                    if (err_trace_index_to_restore != .none) {
+                        // TODO: error-liveness and is_non_err
+
+                        _ = try parent_gz.add(.{
+                            .tag = .restore_err_ret_index,
+                            .data = .{ .un_node = .{
+                                .operand = err_trace_index_to_restore,
+                                .src_node = parent_gz.nodeIndexToRelative(node),
+                            } },
+                        });
+                    }
+
                     _ = try parent_gz.addBreak(break_tag, block_inst, .void_value);
                     return Zir.Inst.Ref.unreachable_value;
                 }
@@ -1884,6 +1912,19 @@ fn breakExpr(parent_gz: *GenZir, parent_scope: *Scope, node: Ast.Node.Index) Inn
 
                 try genDefers(parent_gz, scope, parent_scope, .normal_only);
 
+                // As our last action before the break, "pop" the error trace if needed
+                if (err_trace_index_to_restore != .none) {
+                    // TODO: error-liveness and is_non_err
+
+                    _ = try parent_gz.add(.{
+                        .tag = .restore_err_ret_index,
+                        .data = .{ .un_node = .{
+                            .operand = err_trace_index_to_restore,
+                            .src_node = parent_gz.nodeIndexToRelative(node),
+                        } },
+                    });
+                }
+
                 switch (block_gz.break_result_loc) {
                     .block_ptr => {
                         const br = try parent_gz.addBreak(break_tag, block_inst, operand);
@@ -5160,9 +5201,7 @@ fn orelseCatchExpr(
     block_scope.setBreakResultLoc(rl);
     defer block_scope.unstack();
 
-    if (do_err_trace) {
-        block_scope.saved_err_trace_index = try parent_gz.addNode(.save_err_ret_index, node);
-    }
+    const saved_err_trace_index = if (do_err_trace) try parent_gz.addNode(.save_err_ret_index, node) else .none;
 
     const operand_rl: ResultLoc = switch (block_scope.break_result_loc) {
         .ref => .ref,
@@ -5195,6 +5234,12 @@ fn orelseCatchExpr(
     var else_scope = block_scope.makeSubBlock(scope);
     defer else_scope.unstack();
 
+    // Any break (of a non-error value) that navigates out of this scope means
+    // that the error was handled successfully, so this index will be restored.
+    else_scope.saved_err_trace_index = saved_err_trace_index;
+    if (else_scope.outermost_err_trace_index == .none)
+        else_scope.outermost_err_trace_index = saved_err_trace_index;
+
     var err_val_scope: Scope.LocalVal = undefined;
     const else_sub_scope = blk: {
         const payload = payload_token orelse break :blk &else_scope.base;
@@ -5220,6 +5265,17 @@ fn orelseCatchExpr(
     const else_result = try expr(&else_scope, else_sub_scope, block_scope.break_result_loc, rhs);
     if (!else_scope.endsWithNoReturn()) {
         block_scope.break_count += 1;
+
+        // TODO: Add is_non_err and break check
+        if (do_err_trace) {
+            _ = try else_scope.add(.{
+                .tag = .restore_err_ret_index,
+                .data = .{ .un_node = .{
+                    .operand = saved_err_trace_index,
+                    .src_node = parent_gz.nodeIndexToRelative(node),
+                } },
+            });
+        }
     }
     try checkUsed(parent_gz, &else_scope.base, else_sub_scope);
 
@@ -5243,15 +5299,6 @@ fn orelseCatchExpr(
         block,
         break_tag,
     );
-    if (do_err_trace) {
-        _ = try parent_gz.add(.{
-            .tag = .restore_err_ret_index,
-            .data = .{ .un_node = .{
-                .operand = parent_gz.saved_err_trace_index,
-                .src_node = parent_gz.nodeIndexToRelative(node),
-            } },
-        });
-    }
     return result;
 }
 
@@ -5454,9 +5501,7 @@ fn ifExpr(
     block_scope.setBreakResultLoc(rl);
     defer block_scope.unstack();
 
-    if (do_err_trace) {
-        block_scope.saved_err_trace_index = try parent_gz.addNode(.save_err_ret_index, node);
-    }
+    const saved_err_trace_index = if (do_err_trace) try parent_gz.addNode(.save_err_ret_index, node) else .none;
 
     const payload_is_ref = if (if_full.payload_token) |payload_token|
         token_tags[payload_token] == .asterisk
@@ -5574,6 +5619,12 @@ fn ifExpr(
     var else_scope = parent_gz.makeSubBlock(scope);
     defer else_scope.unstack();
 
+    // Any break (of a non-error value) that navigates out of this scope means
+    // that the error was handled successfully, so this index will be restored.
+    else_scope.saved_err_trace_index = saved_err_trace_index;
+    if (else_scope.outermost_err_trace_index == .none)
+        else_scope.outermost_err_trace_index = saved_err_trace_index;
+
     const else_node = if_full.ast.else_expr;
     const else_info: struct {
         src: Ast.Node.Index,
@@ -5625,6 +5676,18 @@ fn ifExpr(
         },
     };
 
+    if (do_err_trace and !else_scope.endsWithNoReturn()) {
+        // TODO: is_non_err and other checks
+
+        _ = try else_scope.add(.{
+            .tag = .restore_err_ret_index,
+            .data = .{ .un_node = .{
+                .operand = saved_err_trace_index,
+                .src_node = parent_gz.nodeIndexToRelative(node),
+            } },
+        });
+    }
+
     const break_tag: Zir.Inst.Tag = if (parent_gz.force_comptime) .break_inline else .@"break";
     const result = try finishThenElseBlock(
         parent_gz,
@@ -5641,15 +5704,6 @@ fn ifExpr(
         block,
         break_tag,
     );
-    if (do_err_trace) {
-        _ = try parent_gz.add(.{
-            .tag = .restore_err_ret_index,
-            .data = .{ .un_node = .{
-                .operand = parent_gz.saved_err_trace_index,
-                .src_node = parent_gz.nodeIndexToRelative(node),
-            } },
-        });
-    }
     return result;
 }
 
@@ -6780,11 +6834,24 @@ fn ret(gz: *GenZir, scope: *Scope, node: Ast.Node.Index) InnerError!Zir.Inst.Ref
     const operand = try reachableExpr(gz, scope, rl, operand_node, node);
     gz.anon_name_strategy = prev_anon_name_strategy;
 
+    // TODO: This should be almost identical for every break/ret
     switch (nodeMayEvalToError(tree, operand_node)) {
         .never => {
             // Returning a value that cannot be an error; skip error defers.
             try genDefers(gz, defer_outer, scope, .normal_only);
             try emitDbgStmt(gz, ret_line, ret_column);
+
+            // As our last action before the return, "pop" the error trace if needed
+            if (gz.outermost_err_trace_index != .none) {
+                _ = try gz.add(.{
+                    .tag = .restore_err_ret_index,
+                    .data = .{ .un_node = .{
+                        .operand = gz.outermost_err_trace_index,
+                        .src_node = gz.nodeIndexToRelative(node),
+                    } },
+                });
+            }
+
             try gz.addRet(rl, operand, node);
             return Zir.Inst.Ref.unreachable_value;
         },
@@ -6826,6 +6893,17 @@ fn ret(gz: *GenZir, scope: *Scope, node: Ast.Node.Index) InnerError!Zir.Inst.Ref
             };
             try genDefers(&else_scope, defer_outer, scope, which_ones);
             try emitDbgStmt(&else_scope, ret_line, ret_column);
+
+            // As our last action before the return, "pop" the error trace if needed
+            if (else_scope.outermost_err_trace_index != .none) {
+                _ = try else_scope.add(.{
+                    .tag = .restore_err_ret_index,
+                    .data = .{ .un_node = .{
+                        .operand = else_scope.outermost_err_trace_index,
+                        .src_node = else_scope.nodeIndexToRelative(node),
+                    } },
+                });
+            }
             try else_scope.addRet(rl, operand, node);
 
             try setCondBrPayload(condbr, is_non_err, &then_scope, 0, &else_scope, 0);
@@ -10334,7 +10412,12 @@ const GenZir = struct {
     /// Keys are the raw instruction index, values are the closure_capture instruction.
     captures: std.AutoHashMapUnmanaged(Zir.Inst.Index, Zir.Inst.Index) = .{},
 
+    /// If this GenZir corresponds to a `catch { ... }` or `else |err| { ... }` block,
+    /// this err_trace_index can be restored to "pop" the trace entries for the block.
     saved_err_trace_index: Zir.Inst.Ref = .none,
+    /// When returning from a function with a non-error, we must pop all trace entries
+    /// from any containing `catch { ... }` or `else |err| { ... }` blocks.
+    outermost_err_trace_index: Zir.Inst.Ref = .none,
 
     const unstacked_top = std.math.maxInt(usize);
     /// Call unstack before adding any new instructions to containing GenZir.
@@ -10380,7 +10463,7 @@ const GenZir = struct {
             .any_defer_node = gz.any_defer_node,
             .instructions = gz.instructions,
             .instructions_top = gz.instructions.items.len,
-            .saved_err_trace_index = gz.saved_err_trace_index,
+            .outermost_err_trace_index = gz.outermost_err_trace_index,
         };
     }
 
diff --git a/src/Sema.zig b/src/Sema.zig
index 28a012bfb1..ce3a10fd38 100644
--- a/src/Sema.zig
+++ b/src/Sema.zig
@@ -16190,9 +16190,14 @@ fn zirSaveErrRetIndex(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileE
     // This is only relevant at runtime.
     if (block.is_comptime) return Air.Inst.Ref.zero_usize;
 
+    // In the corner case that `catch { ... }` or `else |err| { ... }` is used in a function
+    // that does *not* make any errorable calls, we still need an error trace to interact with
+    // the AIR instructions we've already emitted.
+    if (sema.owner_func != null)
+        sema.owner_func.?.calls_or_awaits_errorable_fn = true;
+
     const backend_supports_error_return_tracing = sema.mod.comp.bin_file.options.use_llvm;
-    const ok = sema.owner_func.?.calls_or_awaits_errorable_fn and
-        sema.mod.comp.bin_file.options.error_return_tracing and
+    const ok = sema.mod.comp.bin_file.options.error_return_tracing and
         backend_supports_error_return_tracing;
     if (!ok) return Air.Inst.Ref.zero_usize;
 
@@ -16211,8 +16216,7 @@ fn zirRestoreErrRetIndex(sema: *Sema, block: *Block, inst: Zir.Inst.Index) Compi
     if (block.is_comptime) return;
 
     const backend_supports_error_return_tracing = sema.mod.comp.bin_file.options.use_llvm;
-    const ok = sema.owner_func.?.calls_or_awaits_errorable_fn and
-        sema.mod.comp.bin_file.options.error_return_tracing and
+    const ok = sema.mod.comp.bin_file.options.error_return_tracing and
         backend_supports_error_return_tracing;
     if (!ok) return;