ktMultiThreaded

fn ktMultiThreaded(
    comptime Variant: type,
    allocator: Allocator,
    io: Io,
    view: *const MultiSliceView,
    total_len: usize,
    output: []u8,
) !void

File

Code

fn ktMultiThreaded(
    comptime Variant: type,
    allocator: Allocator,
    io: Io,
    view: *const MultiSliceView,
    total_len: usize,
    output: []u8,
) !void {
    comptime assert(bytes_per_batch % (optimal_vector_len * chunk_size) == 0);

    const cv_size = Variant.cv_size;
    const StateType = Variant.StateType;
    const leaves_per_batch = bytes_per_batch / chunk_size;
    const remaining_bytes = total_len - chunk_size;
    const total_leaves = std.math.divCeil(usize, remaining_bytes, chunk_size) catch unreachable;

    var final_state = StateType.init(.{});

    var first_chunk_buffer: [chunk_size]u8 = undefined;
    if (view.tryGetSlice(0, chunk_size)) |first_chunk| {
        final_state.update(first_chunk);
    } else {
        view.copyRange(0, chunk_size, &first_chunk_buffer);
        final_state.update(&first_chunk_buffer);
    }

    const padding = [_]u8{ 0x03, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00 };
    final_state.update(&padding);

    const full_leaves = remaining_bytes / chunk_size;
    const has_partial_leaf = (remaining_bytes % chunk_size) != 0;
    const partial_leaf_size = if (has_partial_leaf) remaining_bytes % chunk_size else 0;

    if (full_leaves > 0) {
        const total_batches = std.math.divCeil(usize, full_leaves, leaves_per_batch) catch unreachable;
        const max_concurrent: usize = @min(256, total_batches);

        const Result = BatchResult(Variant);
        const SelectResult = union(enum) { batch: Result };
        const Select = Io.Select(SelectResult);

        const select_buf = try allocator.alloc(SelectResult, max_concurrent);
        defer allocator.free(select_buf);

        // Buffer for out-of-order results (select_buf slots get reused)
        const pending_cv_buf = try allocator.alloc([leaves_per_batch * cv_size]u8, max_concurrent);
        defer allocator.free(pending_cv_buf);
        var pending_cv_lens: [256]usize = @splat(0);

        var select_outstanding: usize = 0;
        var select: Select = .init(io, select_buf);
        defer select.cancel();
        var batches_spawned: usize = 0;
        var next_to_process: usize = 0;

        while (next_to_process < total_batches) {
            while (batches_spawned < total_batches and batches_spawned - next_to_process < max_concurrent) {
                const batch_start_leaf = batches_spawned * leaves_per_batch;
                const batch_leaves = @min(leaves_per_batch, full_leaves - batch_start_leaf);
                const start_offset = chunk_size + batch_start_leaf * chunk_size;

                select_outstanding += 1;
                select.async(.batch, SelectLeafContext(Variant).process, .{SelectLeafContext(Variant){
                    .view = view,
                    .batch_idx = batches_spawned,
                    .start_offset = start_offset,
                    .num_leaves = batch_leaves,
                }});
                batches_spawned += 1;
            }

            select_outstanding -= 1;
            const result = try select.await();
            const batch = result.batch;
            const slot = batch.batch_idx % max_concurrent;

            if (batch.batch_idx == next_to_process) {
                final_state.update(batch.cvs[0..batch.cv_len]);
                next_to_process += 1;

                // Drain pending batches that are now ready
                while (next_to_process < total_batches) {
                    const pending_slot = next_to_process % max_concurrent;
                    const pending_len = pending_cv_lens[pending_slot];
                    if (pending_len == 0) break;

                    final_state.update(pending_cv_buf[pending_slot][0..pending_len]);
                    pending_cv_lens[pending_slot] = 0;
                    next_to_process += 1;
                }
            } else {
                @memcpy(pending_cv_buf[slot][0..batch.cv_len], batch.cvs[0..batch.cv_len]);
                pending_cv_lens[slot] = batch.cv_len;
            }
        }

        assert(select_outstanding == 0);
    }

    if (has_partial_leaf) {
        var cv_buffer: [64]u8 = undefined;
        var leaf_buffer: [chunk_size]u8 = undefined;

        const start_offset = chunk_size + full_leaves * chunk_size;
        if (view.tryGetSlice(start_offset, start_offset + partial_leaf_size)) |leaf_data| {
            const cv_slice = MultiSliceView.init(leaf_data, &[_]u8{}, &[_]u8{});
            Variant.turboShakeToBuffer(&cv_slice, 0x0B, cv_buffer[0..cv_size]);
        } else {
            view.copyRange(start_offset, start_offset + partial_leaf_size, leaf_buffer[0..partial_leaf_size]);
            const cv_slice = MultiSliceView.init(leaf_buffer[0..partial_leaf_size], &[_]u8{}, &[_]u8{});
            Variant.turboShakeToBuffer(&cv_slice, 0x0B, cv_buffer[0..cv_size]);
        }
        final_state.update(cv_buffer[0..cv_size]);
    }

    const n_enc = rightEncode(total_leaves);
    final_state.update(n_enc.slice());
    const terminator = [_]u8{ 0xFF, 0xFF };
    final_state.update(&terminator);

    final_state.final(output);
}