implement @expect builtin (#19658)

* implement `@expect`

* add docs

* add a second arg for expected bool

* fix typo

* move `expect` to use BinOp

* update to newer langref format

doc/langref.html.in

@@ -4799,6 +4799,14 @@ fn cmpxchgWeakButNotAtomic(comptime T: type, ptr: *T, expected_value: T, new_val
       {#see_also|@export#}
       {#header_close#}
 
+      {#header_open|@expect#}
+      <pre>{#syntax#}@expect(operand: bool, expected: bool) bool{#endsyntax#}</pre>
+      <p>
+      Informs the optimizer that {#syntax#}operand{#endsyntax#} will likely be {#syntax#}expected{#endsyntax#}, which influences branch compilation to prefer generating the true branch first.
+      </p>
+      {#code|expect_if.zig#}
+      {#header_close#}
+
       {#header_open|@fence#}
       <pre>{#syntax#}@fence(order: AtomicOrder) void{#endsyntax#}</pre>
       <p>

doc/langref/expect_if.zig

@@ -0,0 +1,15 @@
+pub fn a(x: u32) void {
+    if (@expect(x == 0, false)) {
+        // condition check falls through at code generation
+        return;
+    } else {
+        // condition is branched to at code generation
+        return;
+    }
+}
+
+test "expect" {
+    a(10);
+}
+
+// test

lib/std/zig/AstGen.zig

@@ -2823,6 +2823,7 @@ fn addEnsureResult(gz: *GenZir, maybe_unused_result: Zir.Inst.Ref, statement: As
                 .set_float_mode,
                 .set_align_stack,
                 .set_cold,
+                .expect,
                 => break :b true,
                 else => break :b false,
             },
@@ -9292,7 +9293,14 @@ fn builtinCall(
             });
             return rvalue(gz, ri, .void_value, node);
         },
-
+        .expect => {
+            const val = try gz.addExtendedPayload(.expect, Zir.Inst.BinNode{
+                .node = gz.nodeIndexToRelative(node),
+                .lhs = try expr(gz, scope, .{ .rl = .{ .ty = .bool_type } }, params[0]),
+                .rhs = try expr(gz, scope, .{ .rl = .{ .ty = .bool_type } }, params[1]),
+            });
+            return rvalue(gz, ri, val, node);
+        },
         .src => {
             const token_starts = tree.tokens.items(.start);
             const node_start = token_starts[tree.firstToken(node)];

lib/std/zig/AstRlAnnotate.zig

@@ -1100,5 +1100,10 @@ fn builtinCall(astrl: *AstRlAnnotate, block: ?*Block, ri: ResultInfo, node: Ast.
             _ = try astrl.expr(args[4], block, ResultInfo.type_only);
             return false;
         },
+        .expect => {
+            _ = try astrl.expr(args[0], block, ResultInfo.none);
+            _ = try astrl.expr(args[1], block, ResultInfo.none);
+            return false;
+        },
     }
 }

lib/std/zig/BuiltinFn.zig

@@ -82,6 +82,7 @@ pub const Tag = enum {
     select,
     set_align_stack,
     set_cold,
+    expect,
     set_eval_branch_quota,
     set_float_mode,
     set_runtime_safety,
@@ -743,6 +744,13 @@ pub const list = list: {
                 .illegal_outside_function = true,
             },
         },
+        .{
+            "@expect",
+            .{
+                .tag = .expect,
+                .param_count = 2,
+            },
+        },
         .{
             "@setEvalBranchQuota",
             .{

lib/std/zig/Zir.zig

@@ -2060,6 +2060,9 @@ pub const Inst = struct {
         /// Guaranteed to not have the `ptr_cast` flag.
         /// Uses the `pl_node` union field with payload `FieldParentPtr`.
         field_parent_ptr,
+        /// Implements the `@expect` builtin.
+        /// `operand` is BinOp
+        expect,
 
         pub const InstData = struct {
             opcode: Extended,

lib/zig.h

@@ -318,6 +318,12 @@ typedef char bool;
 #define zig_noreturn
 #endif
 
+#if defined(__GNUC__) || defined(__clang__)
+#define zig_expect(op, exp) __builtin_expect(op, exp)
+#else
+#define zig_expect(op, exp) (op)
+#endif
+
 #define zig_bitSizeOf(T) (CHAR_BIT * sizeof(T))
 
 #define zig_compiler_rt_abbrev_uint32_t si

src/Air.zig

@@ -848,6 +848,10 @@ pub const Inst = struct {
         /// Operand is unused and set to Ref.none
         work_group_id,
 
+        /// Implements @expect builtin.
+        /// Uses the `bin_op` field.
+        expect,
+
         pub fn fromCmpOp(op: std.math.CompareOperator, optimized: bool) Tag {
             switch (op) {
                 .lt => return if (optimized) .cmp_lt_optimized else .cmp_lt,
@@ -1517,6 +1521,8 @@ pub fn typeOfIndex(air: *const Air, inst: Air.Inst.Index, ip: *const InternPool)
         .work_group_id,
         => return Type.u32,
 
+        .expect => return Type.bool,
+
         .inferred_alloc => unreachable,
         .inferred_alloc_comptime => unreachable,
     }
@@ -1634,6 +1640,7 @@ pub fn mustLower(air: Air, inst: Air.Inst.Index, ip: *const InternPool) bool {
         .add_safe,
         .sub_safe,
         .mul_safe,
+        .expect,
         => true,
 
         .add,

src/Liveness.zig

@@ -286,6 +286,7 @@ pub fn categorizeOperand(
         .cmp_gte_optimized,
         .cmp_gt_optimized,
         .cmp_neq_optimized,
+        .expect,
         => {
             const o = air_datas[@intFromEnum(inst)].bin_op;
             if (o.lhs == operand_ref) return matchOperandSmallIndex(l, inst, 0, .none);
@@ -955,6 +956,7 @@ fn analyzeInst(
         .memset,
         .memset_safe,
         .memcpy,
+        .expect,
         => {
             const o = inst_datas[@intFromEnum(inst)].bin_op;
             return analyzeOperands(a, pass, data, inst, .{ o.lhs, o.rhs, .none });

src/Liveness/Verify.zig

@@ -257,6 +257,7 @@ fn verifyBody(self: *Verify, body: []const Air.Inst.Index) Error!void {
             .memset,
             .memset_safe,
             .memcpy,
+            .expect,
             => {
                 const bin_op = data[@intFromEnum(inst)].bin_op;
                 try self.verifyInstOperands(inst, .{ bin_op.lhs, bin_op.rhs, .none });

src/Module.zig

@@ -5546,6 +5546,7 @@ pub const Feature = enum {
     /// to generate better machine code in the backends. All backends should migrate to
     /// enabling this feature.
     safety_checked_instructions,
+    can_expect,
 };
 
 pub fn backendSupportsFeature(zcu: Module, feature: Feature) bool {

src/Sema.zig

@@ -1258,6 +1258,7 @@ fn analyzeBodyInner(
                     .work_group_size    => try sema.zirWorkItem(          block, extended, extended.opcode),
                     .work_group_id      => try sema.zirWorkItem(          block, extended, extended.opcode),
                     .in_comptime        => try sema.zirInComptime(        block),
+                    .expect             => try sema.zirExpect(            block, extended),
                     .closure_get        => try sema.zirClosureGet(        block, extended),
                     // zig fmt: on
 
@@ -17553,6 +17554,34 @@ fn zirThis(
     return sema.analyzeDeclVal(block, src, this_decl_index);
 }
 
+fn zirExpect(sema: *Sema, block: *Block, inst: Zir.Inst.Extended.InstData) CompileError!Air.Inst.Ref {
+    const bin_op = sema.code.extraData(Zir.Inst.BinNode, inst.operand).data;
+    const operand = try sema.resolveInst(bin_op.lhs);
+    const expected = try sema.resolveInst(bin_op.rhs);
+
+    const expected_src = LazySrcLoc{ .node_offset_builtin_call_arg1 = bin_op.node };
+
+    if (!try sema.isComptimeKnown(expected)) {
+        return sema.fail(block, expected_src, "@expect 'expected' must be comptime-known", .{});
+    }
+
+    if (try sema.resolveValue(operand)) |op| {
+        return Air.internedToRef(op.toIntern());
+    }
+
+    if (sema.mod.backendSupportsFeature(.can_expect) and sema.mod.optimizeMode() != .Debug) {
+        return try block.addInst(.{
+            .tag = .expect,
+            .data = .{ .bin_op = .{
+                .lhs = operand,
+                .rhs = expected,
+            } },
+        });
+    } else {
+        return operand;
+    }
+}
+
 fn zirClosureGet(sema: *Sema, block: *Block, extended: Zir.Inst.Extended.InstData) CompileError!Air.Inst.Ref {
     const mod = sema.mod;
     const ip = &mod.intern_pool;

src/arch/aarch64/CodeGen.zig

@@ -803,6 +803,8 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void {
             .@"try"          => try self.airTry(inst),
             .try_ptr         => try self.airTryPtr(inst),
 
+            .expect => unreachable,
+
             .dbg_stmt         => try self.airDbgStmt(inst),
             .dbg_inline_block => try self.airDbgInlineBlock(inst),
             .dbg_var_ptr,

src/arch/arm/CodeGen.zig

@@ -844,6 +844,8 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void {
             .wrap_errunion_payload => try self.airWrapErrUnionPayload(inst),
             .wrap_errunion_err     => try self.airWrapErrUnionErr(inst),
 
+            .expect => unreachable,
+
             .add_optimized,
             .sub_optimized,
             .mul_optimized,

src/arch/riscv64/CodeGen.zig

@@ -1200,6 +1200,8 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void {
             .@"try"          =>  try self.airTry(inst),
             .try_ptr         =>  return self.fail("TODO: try_ptr", .{}),
 
+            .expect => unreachable,
+            
             .dbg_var_ptr,
             .dbg_var_val,
             => try self.airDbgVar(inst),

src/arch/sparc64/CodeGen.zig

@@ -636,6 +636,8 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void {
             .@"try"          => try self.airTry(inst),
             .try_ptr         => @panic("TODO try self.airTryPtr(inst)"),
 
+            .expect => unreachable,
+
             .dbg_stmt         => try self.airDbgStmt(inst),
             .dbg_inline_block => try self.airDbgInlineBlock(inst),
             .dbg_var_ptr,

src/arch/wasm/CodeGen.zig

@@ -2016,6 +2016,8 @@ fn genInst(func: *CodeGen, inst: Air.Inst.Index) InnerError!void {
         .c_va_start,
         => |tag| return func.fail("TODO: Implement wasm inst: {s}", .{@tagName(tag)}),
 
+        .expect => unreachable,
+
         .atomic_load => func.airAtomicLoad(inst),
         .atomic_store_unordered,
         .atomic_store_monotonic,

src/arch/x86_64/CodeGen.zig

@@ -2014,6 +2014,8 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void {
 
             .abs => try self.airAbs(inst),
 
+            .expect => unreachable,
+
             .add_with_overflow => try self.airAddSubWithOverflow(inst),
             .sub_with_overflow => try self.airAddSubWithOverflow(inst),
             .mul_with_overflow => try self.airMulWithOverflow(inst),

src/codegen/c.zig

@@ -3343,6 +3343,8 @@ fn genBodyInner(f: *Function, body: []const Air.Inst.Index) error{ AnalysisFail,
 
             .@"try"  => try airTry(f, inst),
             .try_ptr => try airTryPtr(f, inst),
+            
+            .expect => try airExpect(f, inst),
 
             .dbg_stmt => try airDbgStmt(f, inst),
             .dbg_inline_block => try airDbgInlineBlock(f, inst),
@@ -4704,6 +4706,27 @@ fn airTryPtr(f: *Function, inst: Air.Inst.Index) !CValue {
     return lowerTry(f, inst, extra.data.ptr, body, err_union_ty, true);
 }
 
+fn airExpect(f: *Function, inst: Air.Inst.Index) !CValue {
+    const bin_op = f.air.instructions.items(.data)[@intFromEnum(inst)].bin_op;
+    const operand = try f.resolveInst(bin_op.lhs);
+    const expected = try f.resolveInst(bin_op.rhs);
+
+    const writer = f.object.writer();
+    const local = try f.allocLocal(inst, Type.bool);
+    const a = try Assignment.start(f, writer, CType.bool);
+    try f.writeCValue(writer, local, .Other);
+    try a.assign(f, writer);
+
+    try writer.writeAll("zig_expect(");
+    try f.writeCValue(writer, operand, .FunctionArgument);
+    try writer.writeAll(", ");
+    try f.writeCValue(writer, expected, .FunctionArgument);
+    try writer.writeAll(")");
+
+    try a.end(f, writer);
+    return local;
+}
+
 fn lowerTry(
     f: *Function,
     inst: Air.Inst.Index,

src/codegen/llvm.zig

@@ -5038,6 +5038,8 @@ pub const FuncGen = struct {
                 .slice_ptr      => try self.airSliceField(inst, 0),
                 .slice_len      => try self.airSliceField(inst, 1),
 
+                .expect => try self.airExpect(inst),
+
                 .call              => try self.airCall(inst, .auto),
                 .call_always_tail  => try self.airCall(inst, .always_tail),
                 .call_never_tail   => try self.airCall(inst, .never_tail),
@@ -6365,6 +6367,26 @@ pub const FuncGen = struct {
         return result;
     }
 
+    // Note that the LowerExpectPass only runs in Release modes
+    fn airExpect(self: *FuncGen, inst: Air.Inst.Index) !Builder.Value {
+        const bin_op = self.air.instructions.items(.data)[@intFromEnum(inst)].bin_op;
+
+        const operand = try self.resolveInst(bin_op.lhs);
+        const expected = try self.resolveInst(bin_op.rhs);
+
+        return try self.wip.callIntrinsic(
+            .normal,
+            .none,
+            .expect,
+            &.{operand.typeOfWip(&self.wip)},
+            &.{
+                operand,
+                expected,
+            },
+            "",
+        );
+    }
+
     fn sliceOrArrayPtr(fg: *FuncGen, ptr: Builder.Value, ty: Type) Allocator.Error!Builder.Value {
         const o = fg.dg.object;
         const mod = o.module;

src/print_air.zig

@@ -162,6 +162,7 @@ const Writer = struct {
             .memcpy,
             .memset,
             .memset_safe,
+            .expect,
             => try w.writeBinOp(s, inst),
 
             .is_null,

src/print_zir.zig

@@ -591,6 +591,7 @@ const Writer = struct {
             .wasm_memory_grow,
             .prefetch,
             .c_va_arg,
+            .expect,
             => {
                 const inst_data = self.code.extraData(Zir.Inst.BinNode, extended.operand).data;
                 const src = LazySrcLoc.nodeOffset(inst_data.node);

src/target.zig

@@ -535,5 +535,6 @@ pub fn backendSupportsFeature(
         .error_set_has_value => use_llvm or cpu_arch.isWasm(),
         .field_reordering => ofmt == .c or use_llvm,
         .safety_checked_instructions => use_llvm,
+        .can_expect => use_llvm or ofmt == .c,
     };
 }

test/behavior/expect.zig

@@ -0,0 +1,37 @@
+const std = @import("std");
+const expect = std.testing.expect;
+
+test "@expect if-statement" {
+    const x: u32 = 10;
+    _ = &x;
+    if (@expect(x == 20, true)) {}
+}
+
+test "@expect runtime if-statement" {
+    var x: u32 = 10;
+    var y: u32 = 20;
+    _ = &x;
+    _ = &y;
+    if (@expect(x != y, false)) {}
+}
+
+test "@expect bool input/output" {
+    const b: bool = true;
+    try expect(@TypeOf(@expect(b, false)) == bool);
+}
+
+test "@expect bool is transitive" {
+    const a: bool = true;
+    const b = @expect(a, false);
+
+    const c = @intFromBool(!b);
+    std.mem.doNotOptimizeAway(c);
+
+    try expect(c == 0);
+    try expect(@expect(c != 0, false) == false);
+}
+
+test "@expect at comptime" {
+    const a: bool = true;
+    comptime try expect(@expect(a, true) == true);
+}

test/cases/compile_errors/@expect_non_bool.zig

@@ -0,0 +1,11 @@
+export fn a() void {
+    var x: u32 = 10;
+    _ = &x;
+    _ = @expect(x, true);
+}
+
+// error
+// backend=stage2
+// target=native
+//
+// :4:17: error: expected type 'bool', found 'u32'