Merge pull request 'Sema: allow splatting to scalar types' (#37349) from masonremaley/generic-splat into master

Reviewed-on: https://codeberg.org/ziglang/zig/pulls/37349
Reviewed-by: Andrew Kelley <andrew@ziglang.org>

doc/langref.html.in

@@ -5591,7 +5591,7 @@ fn cmpxchgWeakButNotAtomic(comptime T: type, ptr: *T, expected_value: T, new_val
       <p>
       Produces an array or vector where each element is the value
       {#syntax#}element{#endsyntax#}. The return type, including the number of
-      elements, is inferred.
+      elements, is inferred. If the result type is a scalar, `element` is returned.
       </p>
       {#code|test_splat_builtin.zig#}
       {#see_also|Vectors|Arrays|@shuffle#}

doc/langref/test_splat_builtin.zig

@@ -1,5 +1,6 @@
 const std = @import("std");
 const expectEqualSlices = std.testing.expectEqualSlices;
+const expectEqual = std.testing.expectEqual;
 
 test "vector @splat" {
     const scalar: u32 = 5;
@@ -13,4 +14,10 @@ test "array @splat" {
     try expectEqualSlices(u32, &[_]u32{ 5, 5, 5, 5 }, &@as([4]u32, result));
 }
 
+test "scalar @splat" {
+    const scalar: u32 = 5;
+    const result: u32 = @splat(scalar);
+    try expectEqual(5, result);
+}
+
 // test

src/Sema.zig

@@ -7515,13 +7515,12 @@ fn zirSplatOpResultType(sema: *Sema, block: *Block, inst: Zir.Inst.Index) Compil
     const un_node = sema.code.instructions.items(.data)[@backingInt(inst)].un_node;
 
     const raw_ty = try sema.resolveTypeOrPoison(block, LazySrcLoc.unneeded, un_node.operand) orelse return .generic_poison_type;
-    const vec_ty = raw_ty.optEuBaseType(zcu);
+    const base_ty = raw_ty.optEuBaseType(zcu);
 
-    switch (vec_ty.zigTypeTag(zcu)) {
-        .array, .vector => {},
-        else => return sema.fail(block, block.nodeOffset(un_node.src_node), "expected array or vector type, found '{f}'", .{vec_ty.fmt(zcu)}),
+    switch (base_ty.zigTypeTag(zcu)) {
+        .array, .vector => return Air.internedToRef(base_ty.childType(zcu).toIntern()),
+        else => return Air.internedToRef(base_ty.toIntern()),
     }
-    return Air.internedToRef(vec_ty.childType(zcu).toIntern());
 }
 
 fn zirVectorType(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileError!Air.Inst.Ref {
@@ -23314,17 +23313,17 @@ fn zirSplat(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileError!Air.I
     const inst_data = sema.code.instructions.items(.data)[@backingInt(inst)].pl_node;
     const extra = sema.code.extraData(Zir.Inst.Bin, inst_data.payload_index).data;
     const src = block.nodeOffset(inst_data.src_node);
-    const scalar_src = block.builtinCallArgSrc(inst_data.src_node, 0);
+    const operand_src = block.builtinCallArgSrc(inst_data.src_node, 0);
     const dest_ty = try sema.resolveDestType(block, src, extra.lhs, .remove_eu_opt, "@splat");
+    const operand = sema.resolveInst(extra.rhs);
 
     switch (dest_ty.zigTypeTag(zcu)) {
         .array, .vector => {},
-        else => return sema.fail(block, src, "expected array or vector type, found '{f}'", .{dest_ty.fmt(zcu)}),
+        else => return sema.coerce(block, dest_ty, operand, operand_src),
     }
 
-    const operand = sema.resolveInst(extra.rhs);
     const scalar_ty = dest_ty.childType(zcu);
-    const scalar = try sema.coerce(block, scalar_ty, operand, scalar_src);
+    const scalar = try sema.coerce(block, scalar_ty, operand, operand_src);
 
     const len = try sema.usizeCast(block, src, dest_ty.arrayLen(zcu));
 
@@ -23348,7 +23347,7 @@ fn zirSplat(sema: *Sema, block: *Block, inst: Zir.Inst.Index) CompileError!Air.I
         return .fromValue(try pt.aggregateSplatValue(dest_ty, scalar_val));
     }
 
-    try sema.requireRuntimeBlock(block, src, scalar_src);
+    try sema.requireRuntimeBlock(block, src, operand_src);
 
     switch (dest_ty.zigTypeTag(zcu)) {
         .vector, .array => return block.addTyOp(.splat, dest_ty, scalar),

test/behavior/splat.zig

@@ -160,3 +160,22 @@ test "vector @splat" {
     try S.doTheTest();
     try comptime S.doTheTest();
 }
+
+test "scalar @splat" {
+    const S = struct {
+        fn testForTypes(R: type, result: R, Op: type, op: Op) !void {
+            try expect(@as(R, result) == @as(Op, @splat(op)));
+        }
+        fn doTheTest() !void {
+            try testForTypes(comptime_float, 3.1415, comptime_float, 3.1415);
+            try testForTypes(f32, 3.1415, comptime_float, 3.1415);
+            try testForTypes(f32, 3.1415, f32, 3.1415);
+            try testForTypes(u8, 31, u8, 31);
+            try testForTypes(u8, 31, u16, 31);
+            try testForTypes(i16, 31, i16, 31);
+            try testForTypes(i16, 31, i8, 31);
+        }
+    };
+    try S.doTheTest();
+    try comptime S.doTheTest();
+}

test/behavior/x86_64/binary.zig

@@ -16,7 +16,6 @@ const nan = math.nan;
 const Scalar = math.Scalar;
 const select = math.select;
 const sign = math.sign;
-const splat = math.splat;
 const Sse = math.Sse;
 const tmin = math.tmin;
 
@@ -5250,7 +5249,7 @@ test divTruncUnoptimized {
 
 inline fn divTruncOptimized(comptime Type: type, lhs: Type, rhs: Type) Type {
     @setFloatMode(.optimized);
-    return @divTrunc(lhs, select(@abs(rhs) > splat(Type, 0.0), rhs, splat(Type, 1.0)));
+    return @divTrunc(lhs, select(@abs(rhs) > @as(Type, @splat(0.0)), rhs, @as(Type, @splat(1.0))));
 }
 test divTruncOptimized {
     const test_div_trunc_optimized = binary(divTruncOptimized, .{ .compare = .approx_int });
@@ -5271,7 +5270,7 @@ test divFloorUnoptimized {
 
 inline fn divFloorOptimized(comptime Type: type, lhs: Type, rhs: Type) Type {
     @setFloatMode(.optimized);
-    return @divFloor(lhs, select(@abs(rhs) > splat(Type, 0.0), rhs, splat(Type, 1.0)));
+    return @divFloor(lhs, select(@abs(rhs) > @as(Type, @splat(0.0)), rhs, @as(Type, @splat(1.0))));
 }
 test divFloorOptimized {
     const test_div_floor_optimized = binary(divFloorOptimized, .{ .compare = .approx_int });
@@ -5292,7 +5291,7 @@ test divCeilUnoptimized {
 
 inline fn divCeilOptimized(comptime Type: type, lhs: Type, rhs: Type) Type {
     @setFloatMode(.optimized);
-    return @divCeil(lhs, select(@abs(rhs) > splat(Type, 0.0), rhs, splat(Type, 1.0)));
+    return @divCeil(lhs, select(@abs(rhs) > @as(Type, @splat(0.0)), rhs, @as(Type, @splat(1.0))));
 }
 test divCeilOptimized {
     const test_div_ceil_optimized = binary(divCeilOptimized, .{ .compare = .approx_int });
@@ -5393,7 +5392,7 @@ test mulWithOverflow {
 inline fn shlWithOverflow(comptime Type: type, lhs: Type, rhs: Type) struct { Type, ChangeScalar(Type, u1) } {
     const bit_cast_rhs: AsSignedness(Type, .unsigned) = @bitCast(rhs);
     const truncate_rhs: Log2Int(Type) = @truncate(bit_cast_rhs);
-    return @shlWithOverflow(lhs, if (comptime cast(Log2Int(Scalar(Type)), @bitSizeOf(Scalar(Type)))) |bits| truncate_rhs % splat(Log2Int(Type), bits) else truncate_rhs);
+    return @shlWithOverflow(lhs, if (comptime cast(Log2Int(Scalar(Type)), @bitSizeOf(Scalar(Type)))) |bits| truncate_rhs % @as(Log2Int(Type), @splat(bits)) else truncate_rhs);
 }
 test shlWithOverflow {
     const test_shl_with_overflow = binary(shlWithOverflow, .{});
@@ -5492,7 +5491,7 @@ test bitOr {
 inline fn shr(comptime Type: type, lhs: Type, rhs: Type) Type {
     const bit_cast_rhs: AsSignedness(Type, .unsigned) = @bitCast(rhs);
     const truncate_rhs: Log2Int(Type) = @truncate(bit_cast_rhs);
-    return lhs >> if (comptime cast(Log2Int(Scalar(Type)), @bitSizeOf(Scalar(Type)))) |bits| truncate_rhs % splat(Log2Int(Type), bits) else truncate_rhs;
+    return lhs >> if (comptime cast(Log2Int(Scalar(Type)), @bitSizeOf(Scalar(Type)))) |bits| truncate_rhs % @as(Log2Int(Type), @splat(bits)) else truncate_rhs;
 }
 test shr {
     const test_shr = binary(shr, .{});
@@ -5503,7 +5502,7 @@ test shr {
 inline fn shrExact(comptime Type: type, lhs: Type, rhs: Type) Type {
     const bit_cast_rhs: AsSignedness(Type, .unsigned) = @bitCast(rhs);
     const truncate_rhs: Log2Int(Type) = @truncate(bit_cast_rhs);
-    const final_rhs = if (comptime cast(Log2Int(Scalar(Type)), @bitSizeOf(Scalar(Type)))) |bits| truncate_rhs % splat(Log2Int(Type), bits) else truncate_rhs;
+    const final_rhs = if (comptime cast(Log2Int(Scalar(Type)), @bitSizeOf(Scalar(Type)))) |bits| truncate_rhs % @as(Log2Int(Type), @splat(bits)) else truncate_rhs;
     return @shrExact(lhs >> final_rhs << final_rhs, final_rhs);
 }
 test shrExact {
@@ -5515,7 +5514,7 @@ test shrExact {
 inline fn shl(comptime Type: type, lhs: Type, rhs: Type) Type {
     const bit_cast_rhs: AsSignedness(Type, .unsigned) = @bitCast(rhs);
     const truncate_rhs: Log2Int(Type) = @truncate(bit_cast_rhs);
-    return lhs << if (comptime cast(Log2Int(Scalar(Type)), @bitSizeOf(Scalar(Type)))) |bits| truncate_rhs % splat(Log2Int(Type), bits) else truncate_rhs;
+    return lhs << if (comptime cast(Log2Int(Scalar(Type)), @bitSizeOf(Scalar(Type)))) |bits| truncate_rhs % @as(Log2Int(Type), @splat(bits)) else truncate_rhs;
 }
 test shl {
     const test_shl = binary(shl, .{});
@@ -5527,7 +5526,7 @@ inline fn shlExactUnsafe(comptime Type: type, lhs: Type, rhs: Type) Type {
     @setRuntimeSafety(false);
     const bit_cast_rhs: AsSignedness(Type, .unsigned) = @bitCast(rhs);
     const truncate_rhs: Log2Int(Type) = @truncate(bit_cast_rhs);
-    const final_rhs = if (comptime cast(Log2Int(Scalar(Type)), @bitSizeOf(Scalar(Type)))) |bits| truncate_rhs % splat(Log2Int(Type), bits) else truncate_rhs;
+    const final_rhs = if (comptime cast(Log2Int(Scalar(Type)), @bitSizeOf(Scalar(Type)))) |bits| truncate_rhs % @as(Log2Int(Type), @splat(bits)) else truncate_rhs;
     return @shlExact(lhs << final_rhs >> final_rhs, final_rhs);
 }
 test shlExactUnsafe {
@@ -5540,7 +5539,7 @@ inline fn shlSaturate(comptime Type: type, lhs: Type, rhs: Type) Type {
     // workaround https://github.com/ziglang/zig/issues/23034
     if (@inComptime()) {
         // workaround https://github.com/ziglang/zig/issues/23139
-        return lhs <<| @min(@abs(rhs), splat(ChangeScalar(Type, u64), imax(u64)));
+        return lhs <<| @min(@abs(rhs), @as(ChangeScalar(Type, u64), @splat(imax(u64))));
     }
     return lhs <<| @abs(rhs);
 }

test/behavior/x86_64/cast.zig

@@ -12,7 +12,6 @@ const math = @import("math.zig");
 const nan = math.nan;
 const next = math.next;
 const Scalar = math.Scalar;
-const splat = math.splat;
 const Sse = math.Sse;
 const tmin = math.tmin;
 
@@ -14508,18 +14507,18 @@ inline fn intCast(comptime Result: type, comptime Type: type, rhs: Type, comptim
                 rhs & imax(Type) >> bits,
             .vector => rhs | @select(
                 Scalar(Type),
-                ct_rhs < splat(Type, 0),
-                splat(Type, imin(Scalar(Type)) >> bits),
-                splat(Type, 0),
+                ct_rhs < @as(Type, @splat(0)),
+                @as(Type, @splat(imin(Scalar(Type)) >> bits)),
+                @as(Type, @splat(0)),
             ) & ~@select(
                 Scalar(Type),
-                ct_rhs >= splat(Type, 0),
-                splat(Type, imin(Scalar(Type)) >> bits),
-                splat(Type, 0),
+                ct_rhs >= @as(Type, @splat(0)),
+                @as(Type, @splat(imin(Scalar(Type)) >> bits)),
+                @as(Type, @splat(0)),
             ),
             else => comptime unreachable,
         },
-        .mask => |bits| if (bits == rhs_info.bits) rhs else rhs & splat(Type, (1 << bits) - 1),
+        .mask => |bits| if (bits == rhs_info.bits) rhs else rhs & @as(Type, @splat((1 << bits) - 1)),
     });
 }
 test intCast {

test/behavior/x86_64/math.zig

@@ -72,19 +72,13 @@ pub fn Log2Int(comptime Type: type) type {
 pub fn Log2IntCeil(comptime Type: type) type {
     return ChangeScalar(Type, math.Log2IntCeil(Scalar(Type)));
 }
-pub fn splat(comptime Type: type, scalar: Scalar(Type)) Type {
-    return switch (@typeInfo(Type)) {
-        else => scalar,
-        .vector => @splat(scalar),
-    };
-}
 pub fn sign(rhs: anytype) ChangeScalar(@TypeOf(rhs), bool) {
     const Int = ChangeScalar(@TypeOf(rhs), switch (@typeInfo(Scalar(@TypeOf(rhs)))) {
         .int, .comptime_int => Scalar(@TypeOf(rhs)),
         .float => |float| @Int(.signed, float.bits),
         else => @compileError(@typeName(@TypeOf(rhs))),
     });
-    return @as(Int, @bitCast(rhs)) < splat(Int, 0);
+    return @as(Int, @bitCast(rhs)) < @as(Int, @splat(0));
 }
 pub fn select(cond: anytype, lhs: anytype, rhs: @TypeOf(lhs)) @TypeOf(lhs) {
     return switch (@typeInfo(@TypeOf(cond))) {
@@ -117,19 +111,19 @@ pub noinline fn checkExpected(expected: anytype, actual: @TypeOf(expected), comp
                     .approx_or_overflow => @exp2(@log2(epsilon) * 0.4),
                 };
                 const approx_unequal = @abs(expected - actual) > @max(
-                    @abs(expected) * splat(Expected, tolerance),
-                    splat(Expected, switch (compare) {
+                    @abs(expected) * @as(Expected, @splat(tolerance)),
+                    @as(Expected, @splat(switch (compare) {
                         .strict, .relaxed => comptime unreachable,
                         .approx, .approx_or_overflow => tolerance,
                         .approx_int => 1,
-                    }),
+                    })),
                 );
                 break :unexpected switch (compare) {
                     .strict, .relaxed => comptime unreachable,
                     .approx, .approx_int => approx_unequal,
                     .approx_or_overflow => approx_unequal &
-                        (((@abs(expected) != splat(Expected, inf(Expected))) &
-                            (@abs(actual) != splat(Expected, inf(Expected)))) |
+                        (((@abs(expected) != @as(Expected, @splat(inf(Expected)))) &
+                            (@abs(actual) != @as(Expected, @splat(inf(Expected))))) |
                             (sign(expected) != sign(actual))),
                 };
             },

test/cases/compile_errors/splat_bad_result_type.zig

@@ -1,7 +0,0 @@
-export fn f() void {
-    _ = @as(u32, @splat(5));
-}
-
-// error
-//
-// :2:18: error: expected array or vector type, found 'u32'