tiny.simd.ShardMul

1pub const ShardMul = struct {2    table: [bucket_count]u32,3    rounds: [4]hash_mod.WeakTwoMul,45    const Self = @This();67    pub const FeistelPair = struct {8        left: u32,9        right: u32,10    };1112    pub fn init(data: ShardMulData) Self {13        return .{14            .table = data.table,15            .rounds = .{16                hash_mod.WeakTwoMul.initKey(data.keys[0]),17                hash_mod.WeakTwoMul.initKey(data.keys[1]),18                hash_mod.WeakTwoMul.initKey(data.keys[2]),19                hash_mod.WeakTwoMul.initKey(data.keys[3]),20            },21        };22    }2324    pub fn isEmpty(self: Self) bool {25        if (self.table[0] == 0) return true;26        for (self.table) |entry| {27            if (entry == 0) @panic("ShardMul table contains a zero multiplier");28        }29        return false;30    }3132    pub fn feistel(self: Self, input: u64) FeistelPair {33        var left: u32 = @truncate(input);34        var right: u32 = @truncate(input >> 32);35        left ^= self.rounds[0].hash(right);36        right ^= self.rounds[1].hash(left);37        left ^= self.rounds[2].hash(right);38        right ^= self.rounds[3].hash(left);39        return .{ .left = left, .right = right };40    }4142    pub fn bucketIndex(left: u32) u32 {43        return left >> 28;44    }4546    pub fn lookupMul(self: Self, bucket: u32) u32 {47        std.debug.assert(bucket < bucket_count);48        return self.table[bucket];49    }5051    pub fn hash(self: Self, input: u64) u32 {52        std.debug.assert(!self.isEmpty());53        const pair = self.feistel(input);54        return mulAndXorScalar(pair.left, pair.right, self.lookupMul(bucketIndex(pair.left)));55    }5657    pub fn oneVec(self: Self, comptime D64: type, input: D64.Vector) D64.rebind(u32).Vector {58        requireTag(D64, u64);59        std.debug.assert(!self.isEmpty());60        const D32 = D64.rebind(u32);61        var left: D32.Vector = undefined;62        var right: D32.Vector = undefined;63        inline for (0..D64.lane_count) |lane| {64            left[lane] = @truncate(input[lane]);65            right[lane] = @truncate(input[lane] >> 32);66        }67        return self.resultFromFeistel(D32, self.feistelVector(D32, left, right));68    }6970    pub fn twoVec(71        self: Self,72        comptime D32: type,73        first: D32.repartition(u64).Vector,74        second: D32.repartition(u64).Vector,75    ) D32.Vector {76        requireTag(D32, u32);77        if (D32.lane_count < 2) @compileError("ShardMul twoVec requires at least two u32 lanes");78        std.debug.assert(!self.isEmpty());79        const D64 = D32.repartition(u64);80        var left: D32.Vector = undefined;81        var right: D32.Vector = undefined;82        inline for (0..D64.lane_count) |lane| {83            left[lane] = @truncate(first[lane]);84            right[lane] = @truncate(first[lane] >> 32);85            left[D64.lane_count + lane] = @truncate(second[lane]);86            right[D64.lane_count + lane] = @truncate(second[lane] >> 32);87        }88        return self.resultFromFeistel(D32, self.feistelVector(D32, left, right));89    }9091    pub fn mulAndXor(92        comptime D32: type,93        left: D32.Vector,94        right: D32.Vector,95        muls: D32.Vector,96    ) D32.Vector {97        requireTag(D32, u32);98        const D16 = D32.repartition(u16);99        const products = multiply.mulHigh(100            D16,101            @as(D16.Vector, @bitCast(right)),102            @as(D16.Vector, @bitCast(muls)),103        );104        const mixed: D32.Vector = @bitCast(products);105        return left ^ (mixed & @as(D32.Vector, @splat(0x0fff_ffff)));106    }107108    fn feistelVector(109        self: Self,110        comptime D32: type,111        initial_left: D32.Vector,112        initial_right: D32.Vector,113    ) VectorPair(D32) {114        var left = initial_left;115        var right = initial_right;116        left ^= self.rounds[0].oneVec(D32, right);117        right ^= self.rounds[1].oneVec(D32, left);118        left ^= self.rounds[2].oneVec(D32, right);119        right ^= self.rounds[3].oneVec(D32, left);120        return .{ .left = left, .right = right };121    }122123    fn resultFromFeistel(self: Self, comptime D32: type, pair: VectorPair(D32)) D32.Vector {124        const buckets = shift.shiftRight(D32, 28, pair.left);125        const muls = table.lookup16(D32, &self.table, buckets);126        return mulAndXor(D32, pair.left, pair.right, muls);127    }128};

1pub const ShardMul = shardmul.ShardMul;