Zig radix sort

radix-sort.zig

const std = @import("std");
const Log2Int = std.math.Log2Int;
const IntType = std.meta.IntType;

fn radix(comptime T: type, comptime bits: Log2Int(T), shift: Log2Int(T), value: T) IntType(false, bits) {
    const RadixType = IntType(false, bits);
    return @intCast(RadixType, (value >> shift) & std.math.maxInt(RadixType));
}

fn partialRadixSort(comptime T: type, comptime bits: Log2Int(T), shift: Log2Int(T), src: []const T, dst: []T) void {
    std.debug.assert(src.len == dst.len);
    std.debug.assert(bits > 0);
    const buckets = 1 << bits;

    // Calculate the initial number of elements in each bucket
    var prefix = [_]usize{0} ** buckets;
    for (src) |elem| {
        prefix[radix(T, bits, shift, elem)] += 1;
    }

    // Calculate the offset in `dst` for each bucket
    var accum: usize = 0;
    for (prefix) |*elem| {
        const tmp = accum + elem.*;
        elem.* = accum;
        accum = tmp;
    }

    // Sort the items into the buckets
    for (src) |elem| {
        const bucket = radix(T, bits, shift, elem);
        dst[prefix[bucket]] = elem;
        prefix[bucket] += 1;
    }
}

fn radixSort(arr: []u32, tmp: []u32) void {
    partialRadixSort(u32, 8, 0, arr, tmp);
    partialRadixSort(u32, 8, 8, tmp, arr);
    partialRadixSort(u32, 8, 16, arr, tmp);
    partialRadixSort(u32, 8, 24, tmp, arr);
}

pub fn main() !void {
    const amount = 100000000;

    var rng = std.rand.DefaultPrng.init(0);
    var arr = try std.heap.page_allocator.alloc(u32, amount);
    var tmp = try std.heap.page_allocator.alloc(u32, amount);

    for (arr) |*i| i.* = rng.random.int(u32);

    radixSort(arr, tmp);
    // std.sort.sort(u32, arr, std.sort.asc(u32));
}