sampleUniformRejection

fn sampleUniformRejection(
    comptime PolyType: type,
    comptime q: comptime_int,
    comptime bits_per_coef: comptime_int,
    comptime n: comptime_int,
    seed: []const u8,
    domain_sep: []const u8,
) PolyType

File

Code

fn sampleUniformRejection(
    comptime PolyType: type,
    comptime q: comptime_int,
    comptime bits_per_coef: comptime_int,
    comptime n: comptime_int,
    seed: []const u8,
    domain_sep: []const u8,
) PolyType {
    var h = sha3.Shake128.init(.{});
    h.update(seed);
    h.update(domain_sep);

    const buf_len = sha3.Shake128.block_length; // 168 bytes
    var buf: [buf_len]u8 = undefined;

    var ret: PolyType = undefined;
    var coef_idx: usize = 0;

    if (bits_per_coef == 12) {
        // ML-KEM path: pack 2 coefficients per 3 bytes (12 bits each)
        outer: while (true) {
            h.squeeze(&buf);

            var j: usize = 0;
            while (j < buf_len) : (j += 3) {
                const b0 = @as(u16, buf[j]);
                const b1 = @as(u16, buf[j + 1]);
                const b2 = @as(u16, buf[j + 2]);

                const ts: [2]u16 = .{
                    b0 | ((b1 & 0xf) << 8),
                    (b1 >> 4) | (b2 << 4),
                };

                inline for (ts) |t| {
                    if (t < q) {
                        ret.cs[coef_idx] = @intCast(t);
                        coef_idx += 1;
                        if (coef_idx == n) break :outer;
                    }
                }
            }
        }
    } else if (bits_per_coef == 23) {
        // ML-DSA path: 1 coefficient per 3 bytes (23 bits)
        while (coef_idx < n) {
            h.squeeze(&buf);

            var j: usize = 0;
            while (j < buf_len and coef_idx < n) : (j += 3) {
                const t = (@as(u32, buf[j]) |
                    (@as(u32, buf[j + 1]) << 8) |
                    (@as(u32, buf[j + 2]) << 16)) & 0x7fffff;

                if (t < q) {
                    ret.cs[coef_idx] = @intCast(t);
                    coef_idx += 1;
                }
            }
        }
    } else {
        @compileError("bits_per_coef must be 12 or 23");
    }

    return ret;
}