tiny.simd.ShardMulPlan

1pub const ShardMulPlan = struct {2    best_seed: u64,3    key_counts: [bucket_count]usize,4    extra_counts: [bucket_count]usize,5    offsets: [bucket_count + 1]usize,6    slot_count: usize,7    scratch_len: usize,89    const Self = @This();1011    pub fn inspect(keys: []const u64, extra_outputs: []const u32) ShardMulBuildError!Self {12        const total = std.math.add(usize, keys.len, extra_outputs.len) catch13            return error.CapacityExceeded;14        if (total > std.math.maxInt(u32)) return error.CapacityExceeded;15        const engine = random.AesCtrEngine.initDeterministic();16        var extra_counts: [bucket_count]usize = @splat(0);17        for (extra_outputs) |value| extra_counts[value >> 28] += 1;1819        var best_seed: u64 = feistel_candidate_count;20        var best_ratio: f32 = std.math.inf(f32);21        for (0..feistel_candidate_count) |candidate| {22            const candidate_seed: u64 = @intCast(candidate);23            const shard = ShardMul.init(.{ .keys = makeFeistelKeys(&engine, candidate_seed) });24            var counts = extra_counts;25            for (keys) |key| counts[ShardMul.bucketIndex(shard.feistel(key).left)] += 1;26            var minimum = counts[0];27            var maximum = counts[0];28            for (counts[1..]) |count| {29                minimum = @min(minimum, count);30                maximum = @max(maximum, count);31            }32            if (minimum == 0) continue;33            const ratio = @as(f32, @floatFromInt(maximum)) /34                @as(f32, @floatFromInt(minimum));35            if (ratio < best_ratio) {36                best_ratio = ratio;37                best_seed = candidate_seed;38            }39        }4041        const shard = ShardMul.init(.{ .keys = makeFeistelKeys(&engine, best_seed) });42        var key_counts: [bucket_count]usize = @splat(0);43        for (keys) |key| key_counts[ShardMul.bucketIndex(shard.feistel(key).left)] += 1;44        var offsets: [bucket_count + 1]usize = @splat(0);45        var maximum_bucket_size: usize = 0;46        for (0..bucket_count) |bucket| {47            const bucket_size = std.math.add(48                usize,49                key_counts[bucket],50                extra_counts[bucket],51            ) catch return error.CapacityExceeded;52            if (bucket_size > @as(usize, 1) << 28) return error.CapacityExceeded;53            maximum_bucket_size = @max(maximum_bucket_size, bucket_size);54            offsets[bucket + 1] = std.math.add(usize, offsets[bucket], bucket_size) catch55                return error.CapacityExceeded;56        }57        std.debug.assert(offsets[bucket_count] == total);58        const slot_count = try slotsFor(maximum_bucket_size);59        const slot_u64 = std.math.divCeil(usize, slot_count, 2) catch unreachable;60        const scratch_len = std.math.add(usize, total, slot_u64) catch61            return error.CapacityExceeded;62        return .{63            .best_seed = best_seed,64            .key_counts = key_counts,65            .extra_counts = extra_counts,66            .offsets = offsets,67            .slot_count = slot_count,68            .scratch_len = scratch_len,69        };70    }7172    pub fn build(73        self: Self,74        scratch: []u64,75        keys: []const u64,76        extra_outputs: []const u32,77    ) ShardMulBuildError!ShardMulData {78        if (scratch.len < self.scratch_len) return error.ScratchTooSmall;79        const total = std.math.add(usize, keys.len, extra_outputs.len) catch80            return error.CapacityExceeded;81        if (total != self.offsets[bucket_count]) return error.PlanMismatch;82        const engine = random.AesCtrEngine.initDeterministic();83        const feistel_keys = makeFeistelKeys(&engine, self.best_seed);84        const shard = ShardMul.init(.{ .keys = feistel_keys });85        const records = scratch[0..total];86        const slot_bytes = std.mem.sliceAsBytes(scratch[total..self.scratch_len]);87        const all_slots = std.mem.bytesAsSlice(u32, slot_bytes)[0..self.slot_count];88        var key_cursor = self.offsets[0..bucket_count].*;89        var extra_cursor: [bucket_count]usize = undefined;90        for (0..bucket_count) |bucket| {91            extra_cursor[bucket] = self.offsets[bucket] + self.key_counts[bucket];92        }93        for (keys) |key| {94            const pair = shard.feistel(key);95            const bucket = ShardMul.bucketIndex(pair.left);96            if (key_cursor[bucket] >= self.offsets[bucket] + self.key_counts[bucket]) {97                return error.PlanMismatch;98            }99            records[key_cursor[bucket]] = encodePair(pair);100            key_cursor[bucket] += 1;101        }102        for (extra_outputs) |value| {103            const bucket = value >> 28;104            if (extra_cursor[bucket] >= self.offsets[bucket + 1]) return error.PlanMismatch;105            records[extra_cursor[bucket]] = value;106            extra_cursor[bucket] += 1;107        }108        for (0..bucket_count) |bucket| {109            if (key_cursor[bucket] != self.offsets[bucket] + self.key_counts[bucket] or110                extra_cursor[bucket] != self.offsets[bucket + 1])111            {112                return error.PlanMismatch;113            }114        }115116        var data = ShardMulData{ .keys = feistel_keys };117        for (0..bucket_count) |bucket| {118            const bucket_size = self.key_counts[bucket] + self.extra_counts[bucket];119            const slots = all_slots[0..try slotsFor(bucket_size)];120            const first = self.offsets[bucket];121            const key_records = records[first .. first + self.key_counts[bucket]];122            const extras = records[first + self.key_counts[bucket] .. self.offsets[bucket + 1]];123            for (0..max_attempts_per_bucket) |attempt| {124                @memset(slots, 0);125                var stream = random.RngStream.init(126                    &engine,127                    multiplierSeed(self.best_seed, bucket, attempt),128                );129                const muls = makeMultiplierPair(&stream);130                var collision = false;131                for (key_records) |encoded| {132                    const pair = decodePair(encoded);133                    const output = mulAndXorScalar(pair.left, pair.right, muls);134                    if (!insertUnique(slots, output)) {135                        collision = true;136                        break;137                    }138                }139                if (!collision) {140                    for (extras) |encoded| {141                        if (!insertUnique(slots, @truncate(encoded))) {142                            collision = true;143                            break;144                        }145                    }146                }147                if (!collision) {148                    data.table[bucket] = muls;149                    data.attempts[bucket] = @intCast(attempt + 1);150                    break;151                }152            }153            if (data.table[bucket] == 0) return .{};154        }155        return data;156    }157};

1pub const ShardMulPlan = shardmul.ShardMulPlan;