Ones, Zeros in Zig

ones_zeros.zig

const expect = @import("std").testing.expect;

/// The additive identity for this scalar type.
inline fn zero(comptime T: type) T {
    return 0;
}

/// The multiplicative identity for this scalar type.
inline fn one(comptime T: type) T {
    return 1;
}

/// The additive identity for this array/vector type.
inline fn zeros(comptime T: type) T {
    comptime var len: comptime_int = undefined;
    comptime var child: type = undefined;

    switch (@typeInfo(T)) {
        .Array => |artype| ablk: {
            len = artype.len;
            child = artype.child;
            break :ablk;
        },
        .Vector => |vectype| vblk: {
            len = vectype.len;
            child = vectype.child;
            break :vblk;
        },
        else => @compileError("zeros(T) must be an Array or Vector type."),
    }

    var result: T = undefined;
    comptime var i = 0;
    inline while (i < len) : (i += 1) {
        result[i] = zero(child);
    }
    return result;
}

/// The multiplicative identity for this array/vector type.
inline fn ones(comptime T: type) T {
    comptime var len: comptime_int = undefined;
    comptime var child: type = undefined;

    switch (@typeInfo(T)) {
        .Array => |artype| ablk: {
            len = artype.len;
            child = artype.child;
            break :ablk;
        },
        .Vector => |vectype| vblk: {
            len = vectype.len;
            child = vectype.child;
            break :vblk;
        },
        else => @compileError("ones(T) must be an Array or Vector type."),
    }

    var result: T = undefined;
    comptime var i = 0;
    inline while (i < len) : (i += 1) {
        result[i] = one(child);
    }
    return result;
}

test "zero" {
    const f32_zero: f32 = 0;
    const f64_zero: f64 = 0;
    const i32_zero: i32 = 0;
    const u32_zero: u32 = 0;

    expect(zero(f32) == f32_zero);
    expect(zero(f64) == f64_zero);
    expect(zero(i32) == i32_zero);
    expect(zero(u32) == u32_zero);
}

test "one" {
    const f32_one: f32 = 1;
    const f64_one: f64 = 1;
    const i32_one: i32 = 1;
    const u32_one: u32 = 1;

    expect(one(f32) == f32_one);
    expect(one(f64) == f64_one);
    expect(one(i32) == i32_one);
    expect(one(u32) == u32_one);
}

test "zeros" {
    const a3_f32_zeros: [3]f32 = .{ 0.0, 0.0, 0.0 };
    const v3_f32_zeros: @Vector(3, f32) = a3_f32_zeros;

    const a3_f32_zeros_res = zeros([3]f32);
    const v3_f32_zeros_res = zeros(@Vector(3, f32));

    expect(a3_f32_zeros_res[0] == 0);
    expect(a3_f32_zeros_res[1] == 0);
    expect(a3_f32_zeros_res[2] == 0);

    expect(v3_f32_zeros_res[0] == 0);
    expect(v3_f32_zeros_res[1] == 0);
    expect(v3_f32_zeros_res[2] == 0);
}

test "ones" {
    const a3_f32_ones: [3]f32 = .{ 1.0, 1.0, 1.0 };
    const v3_f32_ones: @Vector(3, f32) = a3_f32_ones;

    const a3_f32_ones_res = ones([3]f32);
    const v3_f32_ones_res = ones(@Vector(3, f32));

    expect(a3_f32_ones_res[0] == 1);
    expect(a3_f32_ones_res[1] == 1);
    expect(a3_f32_ones_res[2] == 1);

    expect(v3_f32_ones_res[0] == 1);
    expect(v3_f32_ones_res[1] == 1);
    expect(v3_f32_ones_res[2] == 1);
}

ones and zeros methods are super useful in the julia programming language for writing generic scientific algorithms, I wanted to see how hard it would be to implement them in zig. Turns out: not too hard!