PolyVec

fn PolyVec(comptime len: u8) type

File

Code

fn PolyVec(comptime len: u8) type {
    return struct {
        ps: [len]Poly,

        const Self = @This();
        const zero: Self = .{ .ps = @splat(.zero) };

        /// Apply a unary operation to each polynomial in the vector
        fn map(v: Self, comptime op: fn (Poly) Poly) Self {
            var ret: Self = undefined;
            inline for (0..len) |i| {
                ret.ps[i] = op(v.ps[i]);
            }
            return ret;
        }

        /// Apply a binary operation pairwise to two vectors
        fn mapBinary(a: Self, b: Self, comptime op: fn (Poly, Poly) Poly) Self {
            var ret: Self = undefined;
            inline for (0..len) |i| {
                ret.ps[i] = op(a.ps[i], b.ps[i]);
            }
            return ret;
        }

        /// Apply a binary operation between a vector and a scalar polynomial
        fn mapBinaryPoly(v: Self, scalar: Poly, comptime op: fn (Poly, Poly) Poly) Self {
            var ret: Self = undefined;
            inline for (0..len) |i| {
                ret.ps[i] = op(v.ps[i], scalar);
            }
            return ret;
        }

        fn add(a: Self, b: Self) Self {
            return mapBinary(a, b, Poly.add);
        }

        fn sub(a: Self, b: Self) Self {
            return mapBinary(a, b, Poly.sub);
        }

        fn ntt(v: Self) Self {
            return map(v, Poly.ntt);
        }

        fn invNTT(v: Self) Self {
            return map(v, Poly.invNTT);
        }

        fn normalize(v: Self) Self {
            return map(v, Poly.normalize);
        }

        fn reduceLe2Q(v: Self) Self {
            return map(v, Poly.reduceLe2Q);
        }

        fn normalizeAssumingLe2Q(v: Self) Self {
            return map(v, Poly.normalizeAssumingLe2Q);
        }

        // Check if any polynomial in the vector exceeds the bound
        fn exceeds(v: Self, bound: u32) bool {
            var result = false;
            for (0..len) |i| {
                result = result or v.ps[i].exceeds(bound);
            }
            return result;
        }

        /// Apply Power2Round to each polynomial in the vector
        /// Returns both t0 and t1 vectors
        fn power2Round(v: Self, t0_out: *Self) Self {
            var t1: Self = undefined;
            for (0..len) |i| {
                const result = v.ps[i].power2RoundPoly();
                t0_out.ps[i] = result.t0;
                t1.ps[i] = result.t1;
            }
            return t1;
        }

        /// Generic packing function for vectors
        fn packWith(
            v: Self,
            buf: []u8,
            comptime poly_size: usize,
            comptime pack_fn: fn (Poly, []u8) void,
        ) void {
            inline for (0..len) |i| {
                const offset = i * poly_size;
                pack_fn(v.ps[i], buf[offset..][0..poly_size]);
            }
        }

        /// Generic unpacking function for vectors
        fn unpackWith(
            comptime poly_size: usize,
            comptime unpack_fn: fn ([]const u8) Poly,
            buf: []const u8,
        ) Self {
            var result: Self = undefined;
            inline for (0..len) |i| {
                const offset = i * poly_size;
                result.ps[i] = unpack_fn(buf[offset..][0..poly_size]);
            }
            return result;
        }

        /// Pack T1 vector to bytes
        fn packT1(v: Self, buf: []u8) void {
            const poly_size = (N * (Q_BITS - D)) / 8;
            packWith(v, buf, poly_size, polyPackT1);
        }

        /// Unpack T1 vector from bytes
        fn unpackT1(bytes: []const u8) Self {
            const poly_size = (N * (Q_BITS - D)) / 8;
            return unpackWith(poly_size, polyUnpackT1, bytes);
        }

        /// Pack T0 vector to bytes
        fn packT0(v: Self, buf: []u8) void {
            const poly_size = (N * D) / 8;
            packWith(v, buf, poly_size, polyPackT0);
        }

        /// Unpack T0 vector from bytes
        fn unpackT0(buf: []const u8) Self {
            const poly_size = (N * D) / 8;
            return unpackWith(poly_size, polyUnpackT0, buf);
        }

        /// Pack vector with coefficients in [-eta, eta]
        fn packLeqEta(v: Self, comptime eta: u8, buf: []u8) void {
            const poly_size = if (eta == 2) 96 else 128;
            const pack_fn = struct {
                fn pack(p: Poly, b: []u8) void {
                    polyPackLeqEta(p, eta, b);
                }
            }.pack;
            packWith(v, buf, poly_size, pack_fn);
        }

        /// Unpack vector with coefficients in [-eta, eta]
        fn unpackLeqEta(comptime eta: u8, buf: []const u8) Self {
            const poly_size = if (eta == 2) 96 else 128;
            const unpack_fn = struct {
                fn unpack(b: []const u8) Poly {
                    return polyUnpackLeqEta(eta, b);
                }
            }.unpack;
            return unpackWith(poly_size, unpack_fn, buf);
        }

        /// Pack vector of polynomials with coefficients < gamma1
        fn packLeGamma1(v: Self, comptime gamma1_bits: u8, buf: []u8) void {
            const poly_size = ((gamma1_bits + 1) * N) / 8;
            const pack_fn = struct {
                fn pack(p: Poly, b: []u8) void {
                    polyPackLeGamma1(p, gamma1_bits, b);
                }
            }.pack;
            packWith(v, buf, poly_size, pack_fn);
        }

        /// Unpack vector of polynomials with coefficients < gamma1
        fn unpackLeGamma1(comptime gamma1_bits: u8, buf: []const u8) Self {
            const poly_size = ((gamma1_bits + 1) * N) / 8;
            const unpack_fn = struct {
                fn unpack(b: []const u8) Poly {
                    return polyUnpackLeGamma1(gamma1_bits, b);
                }
            }.unpack;
            return unpackWith(poly_size, unpack_fn, buf);
        }

        /// Pack high bits w1 for signature verification
        fn packW1(v: Self, comptime gamma1_bits: u8, buf: []u8) void {
            const poly_size = (N * (Q_BITS - gamma1_bits)) / 8;
            const pack_fn = struct {
                fn pack(p: Poly, b: []u8) void {
                    polyPackW1(p, gamma1_bits, b);
                }
            }.pack;
            packWith(v, buf, poly_size, pack_fn);
        }

        /// Decompose each polynomial in the vector into high and low bits
        fn decomposeVec(v: Self, comptime gamma2: u32, w0_out: *Self) Self {
            var w1: Self = undefined;
            for (0..len) |i| {
                for (0..N) |j| {
                    const r = decompose(v.ps[i].cs[j], gamma2);
                    w0_out.ps[i].cs[j] = r.a0_plus_q;
                    w1.ps[i].cs[j] = r.a1;
                }
            }
            return w1;
        }

        /// Create hints for vector, returns hint population count
        fn makeHintVec(w0mcs2pct0: Self, w1: Self, comptime gamma2: u32) struct { hint: Self, pop: u32 } {
            var hint: Self = undefined;
            var pop: u32 = 0;
            for (0..len) |i| {
                const result = polyMakeHint(w0mcs2pct0.ps[i], w1.ps[i], gamma2);
                hint.ps[i] = result.hint;
                pop += result.count;
            }
            return .{ .hint = hint, .pop = pop };
        }

        /// Apply hints to recover high bits
        fn useHint(v: Self, hint: Self, comptime gamma2: u32) Self {
            var result: Self = undefined;
            for (0..len) |i| {
                result.ps[i] = polyUseHint(v.ps[i], hint.ps[i], gamma2);
            }
            return result;
        }

        /// Multiply vector by 2^D (left shift)
        fn mulBy2toD(v: Self) Self {
            var result: Self = undefined;
            for (0..len) |i| {
                for (0..N) |j| {
                    result.ps[i].cs[j] = v.ps[i].cs[j] << D;
                }
            }
            return result;
        }

        /// Sample vector with coefficients uniformly in (-gamma1, gamma1]
        /// Wraps expandMask (FIPS 204: ExpandMask)
        fn deriveUniformLeGamma1(comptime gamma1_bits: u8, seed: *const [64]u8, nonce: u16) Self {
            var result: Self = undefined;
            for (0..len) |i| {
                result.ps[i] = expandMask(gamma1_bits, seed, nonce + @as(u16, @intCast(i)));
            }
            return result;
        }

        /// Pack hints into bytes
        /// Format: for each polynomial, find positions where hint[i]=1, encode those positions
        fn packHint(v: Self, comptime omega: u16, buf: []u8) bool {
            var idx: usize = 0;
            var count: u32 = 0;

            for (0..len) |i| {
                for (0..N) |j| {
                    if (v.ps[i].cs[j] != 0) {
                        count += 1;
                    }
                }
            }

            if (count > omega) {
                return false;
            }

            // Hint encoding format per FIPS 204:
            // First omega bytes: positions of set bits across all polynomials
            // Last len bytes: boundary indices showing where each polynomial's hints end
            for (0..len) |i| {
                for (0..N) |j| {
                    if (v.ps[i].cs[j] != 0) {
                        buf[idx] = @intCast(j);
                        idx += 1;
                    }
                }
                buf[omega + i] = @intCast(idx);
            }

            while (idx < omega) : (idx += 1) {
                buf[idx] = 0;
            }

            return true;
        }

        /// Unpack hints from bytes
        fn unpackHint(comptime omega: u16, buf: []const u8) ?Self {
            var result: Self = .{ .ps = @splat(.zero) };
            var prev_sop: u8 = 0; // previous switch-over-point

            for (0..len) |i| {
                const sop = buf[omega + i]; // switch-over-point
                if (sop < prev_sop or sop > omega) {
                    return null; // ensures switch-over-points are increasing
                }

                var j = prev_sop;
                while (j < sop) : (j += 1) {
                    // Validation: indices must be strictly increasing within each polynomial
                    if (j > prev_sop and buf[j] <= buf[j - 1]) {
                        return null;
                    }
                    const pos = buf[j];
                    if (pos >= N) {
                        return null;
                    }
                    result.ps[i].cs[pos] = 1;
                }
                prev_sop = sop;
            }

            var j = prev_sop;
            while (j < omega) : (j += 1) {
                if (buf[j] != 0) {
                    return null;
                }
            }

            return result;
        }
    };
}