Merge pull request #11289 from schmee/stage2-select

stage2: implement `@select`

11 files changed, 184 insertions, 3 deletions

diff --git a/src/Air.zig b/src/Air.zig
index 404ee8f9b7..5120e0fd67 100644
--- a/src/Air.zig
+++ b/src/Air.zig
@@ -344,7 +344,7 @@ pub const Inst = struct {
         /// to the storage for the variable. The local may be a const or a var.
         /// Result type is always void.
         /// Uses `pl_op`. The payload index is the variable name. It points to the extra
-        /// array, reinterpreting the bytes there as a null-terminated string. 
+        /// array, reinterpreting the bytes there as a null-terminated string.
         dbg_var_ptr,
         /// Same as `dbg_var_ptr` except the local is a const, not a var, and the
         /// operand is the local's value.
@@ -553,6 +553,9 @@ pub const Inst = struct {
         /// Constructs a vector by selecting elements from `a` and `b` based on `mask`.
         /// Uses the `ty_pl` field with payload `Shuffle`.
         shuffle,
+        /// Constructs a vector element-wise from `a` or `b` based on `pred`.
+        /// Uses the `pl_op` field with `pred` as operand, and payload `Bin`.
+        select,
 
         /// Given dest ptr, value, and len, set all elements at dest to value.
         /// Result type is always void.
@@ -1067,6 +1070,10 @@ pub fn typeOfIndex(air: Air, inst: Air.Inst.Index) Type {
         .reduce => return air.typeOf(datas[inst].reduce.operand).childType(),
 
         .mul_add => return air.typeOf(datas[inst].pl_op.operand),
+        .select => {
+            const extra = air.extraData(Air.Bin, datas[inst].pl_op.payload).data;
+            return air.typeOf(extra.lhs);
+        },
 
         .add_with_overflow,
         .sub_with_overflow,
diff --git a/src/Liveness.zig b/src/Liveness.zig
index 79521e7a94..b9f4e6b33a 100644
--- a/src/Liveness.zig
+++ b/src/Liveness.zig
@@ -433,6 +433,11 @@ fn analyzeInst(
             }
             return extra_tombs.finish();
         },
+        .select => {
+            const pl_op = inst_datas[inst].pl_op;
+            const extra = a.air.extraData(Air.Bin, pl_op.payload).data;
+            return trackOperands(a, new_set, inst, main_tomb, .{ pl_op.operand, extra.lhs, extra.rhs });
+        },
         .shuffle => {
             const extra = a.air.extraData(Air.Shuffle, inst_datas[inst].ty_pl.payload).data;
             return trackOperands(a, new_set, inst, main_tomb, .{ extra.a, extra.b, .none });
diff --git a/src/Sema.zig b/src/Sema.zig
index 428ad71da1..80227f81af 100644
--- a/src/Sema.zig
+++ b/src/Sema.zig
@@ -14890,8 +14890,91 @@ fn analyzeShuffle(
 
 fn zirSelect(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileError!Air.Inst.Ref {
     const inst_data = sema.code.instructions.items(.data)[inst].pl_node;
-    const src = inst_data.src();
-    return sema.fail(block, src, "TODO: Sema.zirSelect", .{});
+    const extra = sema.code.extraData(Zir.Inst.Select, inst_data.payload_index).data;
+    const target = sema.mod.getTarget();
+
+    const elem_ty_src: LazySrcLoc = .{ .node_offset_builtin_call_arg0 = inst_data.src_node };
+    const pred_src: LazySrcLoc = .{ .node_offset_builtin_call_arg1 = inst_data.src_node };
+    const a_src: LazySrcLoc = .{ .node_offset_builtin_call_arg2 = inst_data.src_node };
+    const b_src: LazySrcLoc = .{ .node_offset_builtin_call_arg3 = inst_data.src_node };
+
+    const elem_ty = try sema.resolveType(block, elem_ty_src, extra.elem_type);
+    try sema.checkVectorElemType(block, elem_ty_src, elem_ty);
+    const pred_uncoerced = sema.resolveInst(extra.pred);
+    const pred_ty = sema.typeOf(pred_uncoerced);
+
+    const vec_len_u64 = switch (try pred_ty.zigTypeTagOrPoison()) {
+        .Vector, .Array => pred_ty.arrayLen(),
+        else => return sema.fail(block, pred_src, "expected vector or array, found '{}'", .{pred_ty.fmt(target)}),
+    };
+    const vec_len = try sema.usizeCast(block, pred_src, vec_len_u64);
+
+    const bool_vec_ty = try Type.vector(sema.arena, vec_len, Type.bool);
+    const pred = try sema.coerce(block, bool_vec_ty, pred_uncoerced, pred_src);
+
+    const vec_ty = try Type.vector(sema.arena, vec_len, elem_ty);
+    const a = try sema.coerce(block, vec_ty, sema.resolveInst(extra.a), a_src);
+    const b = try sema.coerce(block, vec_ty, sema.resolveInst(extra.b), b_src);
+
+    const maybe_pred = try sema.resolveMaybeUndefVal(block, pred_src, pred);
+    const maybe_a = try sema.resolveMaybeUndefVal(block, a_src, a);
+    const maybe_b = try sema.resolveMaybeUndefVal(block, b_src, b);
+
+    const runtime_src = if (maybe_pred) |pred_val| rs: {
+        if (pred_val.isUndef()) return sema.addConstUndef(vec_ty);
+
+        if (maybe_a) |a_val| {
+            if (a_val.isUndef()) return sema.addConstUndef(vec_ty);
+
+            if (maybe_b) |b_val| {
+                if (b_val.isUndef()) return sema.addConstUndef(vec_ty);
+
+                var buf: Value.ElemValueBuffer = undefined;
+                const elems = try sema.gpa.alloc(Value, vec_len);
+                for (elems) |*elem, i| {
+                    const pred_elem_val = pred_val.elemValueBuffer(i, &buf);
+                    const should_choose_a = pred_elem_val.toBool();
+                    if (should_choose_a) {
+                        elem.* = a_val.elemValueBuffer(i, &buf);
+                    } else {
+                        elem.* = b_val.elemValueBuffer(i, &buf);
+                    }
+                }
+
+                return sema.addConstant(
+                    vec_ty,
+                    try Value.Tag.aggregate.create(sema.arena, elems),
+                );
+            } else {
+                break :rs b_src;
+            }
+        } else {
+            if (maybe_b) |b_val| {
+                if (b_val.isUndef()) return sema.addConstUndef(vec_ty);
+            }
+            break :rs a_src;
+        }
+    } else rs: {
+        if (maybe_a) |a_val| {
+            if (a_val.isUndef()) return sema.addConstUndef(vec_ty);
+        }
+        if (maybe_b) |b_val| {
+            if (b_val.isUndef()) return sema.addConstUndef(vec_ty);
+        }
+        break :rs pred_src;
+    };
+
+    try sema.requireRuntimeBlock(block, runtime_src);
+    return block.addInst(.{
+        .tag = .select,
+        .data = .{ .pl_op = .{
+            .operand = pred,
+            .payload = try block.sema.addExtra(Air.Bin{
+                .lhs = a,
+                .rhs = b,
+            }),
+        } },
+    });
 }
 
 fn zirAtomicLoad(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileError!Air.Inst.Ref {
diff --git a/src/arch/aarch64/CodeGen.zig b/src/arch/aarch64/CodeGen.zig
index 09219f74ba..da6da9c877 100644
--- a/src/arch/aarch64/CodeGen.zig
+++ b/src/arch/aarch64/CodeGen.zig
@@ -633,6 +633,7 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void {
             .tag_name        => try self.airTagName(inst),
             .error_name      => try self.airErrorName(inst),
             .splat           => try self.airSplat(inst),
+            .select          => try self.airSelect(inst),
             .shuffle         => try self.airShuffle(inst),
             .reduce          => try self.airReduce(inst),
             .aggregate_init  => try self.airAggregateInit(inst),
@@ -3666,6 +3667,13 @@ fn airSplat(self: *Self, inst: Air.Inst.Index) !void {
     return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
 }
 
+fn airSelect(self: *Self, inst: Air.Inst.Index) !void {
+    const pl_op = self.air.instructions.items(.data)[inst].pl_op;
+    const extra = self.air.extraData(Air.Bin, pl_op.payload).data;
+    const result: MCValue = if (self.liveness.isUnused(inst)) .dead else return self.fail("TODO implement airSelect for {}", .{self.target.cpu.arch});
+    return self.finishAir(inst, result, .{ pl_op.operand, extra.lhs, extra.rhs });
+}
+
 fn airShuffle(self: *Self, inst: Air.Inst.Index) !void {
     const ty_op = self.air.instructions.items(.data)[inst].ty_op;
     const result: MCValue = if (self.liveness.isUnused(inst)) .dead else return self.fail("TODO implement airShuffle for {}", .{self.target.cpu.arch});
diff --git a/src/arch/arm/CodeGen.zig b/src/arch/arm/CodeGen.zig
index 5bbacb8999..15d0e47d4a 100644
--- a/src/arch/arm/CodeGen.zig
+++ b/src/arch/arm/CodeGen.zig
@@ -630,6 +630,7 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void {
             .tag_name        => try self.airTagName(inst),
             .error_name      => try self.airErrorName(inst),
             .splat           => try self.airSplat(inst),
+            .select          => try self.airSelect(inst),
             .shuffle         => try self.airShuffle(inst),
             .reduce          => try self.airReduce(inst),
             .aggregate_init  => try self.airAggregateInit(inst),
@@ -4323,6 +4324,13 @@ fn airSplat(self: *Self, inst: Air.Inst.Index) !void {
     return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
 }
 
+fn airSelect(self: *Self, inst: Air.Inst.Index) !void {
+    const pl_op = self.air.instructions.items(.data)[inst].pl_op;
+    const extra = self.air.extraData(Air.Bin, pl_op.payload).data;
+    const result: MCValue = if (self.liveness.isUnused(inst)) .dead else return self.fail("TODO implement airSelect for arm", .{});
+    return self.finishAir(inst, result, .{ pl_op.operand, extra.lhs, extra.rhs });
+}
+
 fn airShuffle(self: *Self, inst: Air.Inst.Index) !void {
     const ty_op = self.air.instructions.items(.data)[inst].ty_op;
     const result: MCValue = if (self.liveness.isUnused(inst)) .dead else return self.fail("TODO implement airShuffle for arm", .{});
diff --git a/src/arch/riscv64/CodeGen.zig b/src/arch/riscv64/CodeGen.zig
index 8cbd26b7a6..8b697677c1 100644
--- a/src/arch/riscv64/CodeGen.zig
+++ b/src/arch/riscv64/CodeGen.zig
@@ -600,6 +600,7 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void {
             .tag_name        => try self.airTagName(inst),
             .error_name      => try self.airErrorName(inst),
             .splat           => try self.airSplat(inst),
+            .select          => try self.airSelect(inst),
             .shuffle         => try self.airShuffle(inst),
             .reduce          => try self.airReduce(inst),
             .aggregate_init  => try self.airAggregateInit(inst),
@@ -2396,6 +2397,13 @@ fn airSplat(self: *Self, inst: Air.Inst.Index) !void {
     return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
 }
 
+fn airSelect(self: *Self, inst: Air.Inst.Index) !void {
+    const pl_op = self.air.instructions.items(.data)[inst].pl_op;
+    const extra = self.air.extraData(Air.Bin, pl_op.payload).data;
+    const result: MCValue = if (self.liveness.isUnused(inst)) .dead else return self.fail("TODO implement airSelect for riscv64", .{});
+    return self.finishAir(inst, result, .{ pl_op.operand, extra.lhs, extra.rhs });
+}
+
 fn airShuffle(self: *Self, inst: Air.Inst.Index) !void {
     const ty_op = self.air.instructions.items(.data)[inst].ty_op;
     const result: MCValue = if (self.liveness.isUnused(inst)) .dead else return self.fail("TODO implement airShuffle for riscv64", .{});
diff --git a/src/arch/wasm/CodeGen.zig b/src/arch/wasm/CodeGen.zig
index f2979d96b1..7ab4a33be9 100644
--- a/src/arch/wasm/CodeGen.zig
+++ b/src/arch/wasm/CodeGen.zig
@@ -1371,6 +1371,7 @@ fn genInst(self: *Self, inst: Air.Inst.Index) !WValue {
         .ret_ptr => self.airRetPtr(inst),
         .ret_load => self.airRetLoad(inst),
         .splat => self.airSplat(inst),
+        .select => self.airSelect(inst),
         .shuffle => self.airShuffle(inst),
         .reduce => self.airReduce(inst),
         .aggregate_init => self.airAggregateInit(inst),
@@ -3265,6 +3266,16 @@ fn airSplat(self: *Self, inst: Air.Inst.Index) InnerError!WValue {
     return self.fail("TODO: Implement wasm airSplat", .{});
 }
 
+fn airSelect(self: *Self, inst: Air.Inst.Index) InnerError!WValue {
+    if (self.liveness.isUnused(inst)) return WValue{ .none = {} };
+
+    const pl_op = self.air.instructions.items(.data)[inst].pl_op;
+    const operand = try self.resolveInst(pl_op.operand);
+
+    _ = operand;
+    return self.fail("TODO: Implement wasm airSelect", .{});
+}
+
 fn airShuffle(self: *Self, inst: Air.Inst.Index) InnerError!WValue {
     if (self.liveness.isUnused(inst)) return WValue{ .none = {} };
 
diff --git a/src/arch/x86_64/CodeGen.zig b/src/arch/x86_64/CodeGen.zig
index 17a113bc81..5b252b6c0a 100644
--- a/src/arch/x86_64/CodeGen.zig
+++ b/src/arch/x86_64/CodeGen.zig
@@ -714,6 +714,7 @@ fn genBody(self: *Self, body: []const Air.Inst.Index) InnerError!void {
             .tag_name        => try self.airTagName(inst),
             .error_name      => try self.airErrorName(inst),
             .splat           => try self.airSplat(inst),
+            .select          => try self.airSelect(inst),
             .shuffle         => try self.airShuffle(inst),
             .reduce          => try self.airReduce(inst),
             .aggregate_init  => try self.airAggregateInit(inst),
@@ -5624,6 +5625,13 @@ fn airSplat(self: *Self, inst: Air.Inst.Index) !void {
     return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
 }
 
+fn airSelect(self: *Self, inst: Air.Inst.Index) !void {
+    const pl_op = self.air.instructions.items(.data)[inst].pl_op;
+    const extra = self.air.extraData(Air.Bin, pl_op.payload).data;
+    const result: MCValue = if (self.liveness.isUnused(inst)) .dead else return self.fail("TODO implement airSelect for x86_64", .{});
+    return self.finishAir(inst, result, .{ pl_op.operand, extra.lhs, extra.rhs });
+}
+
 fn airShuffle(self: *Self, inst: Air.Inst.Index) !void {
     const ty_op = self.air.instructions.items(.data)[inst].ty_op;
     const result: MCValue = if (self.liveness.isUnused(inst)) .dead else return self.fail("TODO implement airShuffle for x86_64", .{});
diff --git a/src/codegen/c.zig b/src/codegen/c.zig
index 16558c6e04..84bb3a211b 100644
--- a/src/codegen/c.zig
+++ b/src/codegen/c.zig
@@ -1825,6 +1825,7 @@ fn genBody(f: *Function, body: []const Air.Inst.Index) error{ AnalysisFail, OutO
             .tag_name         => try airTagName(f, inst),
             .error_name       => try airErrorName(f, inst),
             .splat            => try airSplat(f, inst),
+            .select           => try airSelect(f, inst),
             .shuffle          => try airShuffle(f, inst),
             .reduce           => try airReduce(f, inst),
             .aggregate_init   => try airAggregateInit(f, inst),
@@ -3794,6 +3795,21 @@ fn airSplat(f: *Function, inst: Air.Inst.Index) !CValue {
     return f.fail("TODO: C backend: implement airSplat", .{});
 }
 
+fn airSelect(f: *Function, inst: Air.Inst.Index) !CValue {
+    if (f.liveness.isUnused(inst)) return CValue.none;
+
+    const inst_ty = f.air.typeOfIndex(inst);
+    const ty_pl = f.air.instructions.items(.data)[inst].ty_pl;
+
+    const writer = f.object.writer();
+    const local = try f.allocLocal(inst_ty, .Const);
+    try writer.writeAll(" = ");
+
+    _ = local;
+    _ = ty_pl;
+    return f.fail("TODO: C backend: implement airSelect", .{});
+}
+
 fn airShuffle(f: *Function, inst: Air.Inst.Index) !CValue {
     if (f.liveness.isUnused(inst)) return CValue.none;
 
diff --git a/src/codegen/llvm.zig b/src/codegen/llvm.zig
index c8ca540f59..00074d69d1 100644
--- a/src/codegen/llvm.zig
+++ b/src/codegen/llvm.zig
@@ -3444,6 +3444,7 @@ pub const FuncGen = struct {
                 .tag_name       => try self.airTagName(inst),
                 .error_name     => try self.airErrorName(inst),
                 .splat          => try self.airSplat(inst),
+                .select         => try self.airSelect(inst),
                 .shuffle        => try self.airShuffle(inst),
                 .reduce         => try self.airReduce(inst),
                 .aggregate_init => try self.airAggregateInit(inst),
@@ -6355,6 +6356,18 @@ pub const FuncGen = struct {
         return self.builder.buildShuffleVector(op_vector, undef_vector, mask_llvm_ty.constNull(), "");
     }
 
+    fn airSelect(self: *FuncGen, inst: Air.Inst.Index) !?*const llvm.Value {
+        if (self.liveness.isUnused(inst)) return null;
+
+        const pl_op = self.air.instructions.items(.data)[inst].pl_op;
+        const extra = self.air.extraData(Air.Bin, pl_op.payload).data;
+        const pred = try self.resolveInst(pl_op.operand);
+        const a = try self.resolveInst(extra.lhs);
+        const b = try self.resolveInst(extra.rhs);
+
+        return self.builder.buildSelect(pred, a, b, "");
+    }
+
     fn airShuffle(self: *FuncGen, inst: Air.Inst.Index) !?*const llvm.Value {
         if (self.liveness.isUnused(inst)) return null;
 
diff --git a/src/print_air.zig b/src/print_air.zig
index 5e0179e3cc..f1e51150a6 100644
--- a/src/print_air.zig
+++ b/src/print_air.zig
@@ -264,6 +264,7 @@ const Writer = struct {
             .wasm_memory_size => try w.writeWasmMemorySize(s, inst),
             .wasm_memory_grow => try w.writeWasmMemoryGrow(s, inst),
             .mul_add => try w.writeMulAdd(s, inst),
+            .select => try w.writeSelect(s, inst),
             .shuffle => try w.writeShuffle(s, inst),
             .reduce => try w.writeReduce(s, inst),
             .cmp_vector => try w.writeCmpVector(s, inst),
@@ -396,6 +397,19 @@ const Writer = struct {
         try s.print(", mask {d}, len {d}", .{ extra.mask, extra.mask_len });
     }
 
+    fn writeSelect(w: *Writer, s: anytype, inst: Air.Inst.Index) @TypeOf(s).Error!void {
+        const pl_op = w.air.instructions.items(.data)[inst].pl_op;
+        const extra = w.air.extraData(Air.Bin, pl_op.payload).data;
+
+        const elem_ty = w.air.typeOfIndex(inst).childType();
+        try s.print("{}, ", .{elem_ty.fmtDebug()});
+        try w.writeOperand(s, inst, 0, pl_op.operand);
+        try s.writeAll(", ");
+        try w.writeOperand(s, inst, 1, extra.lhs);
+        try s.writeAll(", ");
+        try w.writeOperand(s, inst, 2, extra.rhs);
+    }
+
     fn writeReduce(w: *Writer, s: anytype, inst: Air.Inst.Index) @TypeOf(s).Error!void {
         const reduce = w.air.instructions.items(.data)[inst].reduce;