Poly

const Poly = struct

File

Code

const Poly = struct {
    cs: [N]u32,

    const zero: Poly = .{ .cs = @splat(0) };

    // Add two polynomials (no normalization)
    fn add(a: Poly, b: Poly) Poly {
        var ret: Poly = undefined;
        for (0..N) |i| {
            ret.cs[i] = a.cs[i] + b.cs[i];
        }
        return ret;
    }

    // Subtract two polynomials (assumes b coefficients < 2q)
    fn sub(a: Poly, b: Poly) Poly {
        var ret: Poly = undefined;
        for (0..N) |i| {
            ret.cs[i] = a.cs[i] +% (@as(u32, 2 * Q) -% b.cs[i]);
        }
        return ret;
    }

    // Reduce each coefficient to < 2q
    fn reduceLe2Q(p: Poly) Poly {
        var ret = p;
        for (0..N) |i| {
            ret.cs[i] = le2Q(ret.cs[i]);
        }
        return ret;
    }

    // Normalize coefficients to [0, q)
    fn normalize(p: Poly) Poly {
        var ret = p;
        for (0..N) |i| {
            ret.cs[i] = modQ(ret.cs[i]);
        }
        return ret;
    }

    // Normalize assuming coefficients already < 2q
    fn normalizeAssumingLe2Q(p: Poly) Poly {
        var ret = p;
        for (0..N) |i| {
            ret.cs[i] = le2qModQ(ret.cs[i]);
        }
        return ret;
    }

    // Pointwise multiplication in NTT domain (Montgomery form)
    fn mulHat(a: Poly, b: Poly) Poly {
        var ret: Poly = undefined;
        for (0..N) |i| {
            ret.cs[i] = montReduceLe2Q(@as(u64, a.cs[i]) * @as(u64, b.cs[i]));
        }
        return ret;
    }

    // Forward NTT
    fn ntt(p: Poly) Poly {
        var ret = p;
        ret.nttInPlace();
        return ret;
    }

    // In-place forward NTT
    fn nttInPlace(p: *Poly) void {
        var k: usize = 0;
        var l: usize = N / 2;

        while (l > 0) : (l >>= 1) {
            var offset: usize = 0;
            while (offset < N - l) : (offset += 2 * l) {
                k += 1;
                const zeta: u64 = zetas[k];

                for (offset..offset + l) |j| {
                    const t = montReduceLe2Q(zeta * @as(u64, p.cs[j + l]));
                    p.cs[j + l] = p.cs[j] +% (2 * Q -% t);
                    p.cs[j] +%= t;
                }
            }
        }
    }

    // Inverse NTT
    fn invNTT(p: Poly) Poly {
        var ret = p;
        ret.invNTTInPlace();
        return ret;
    }

    // In-place inverse NTT
    fn invNTTInPlace(p: *Poly) void {
        var k: usize = 0;
        var l: usize = 1;

        while (l < N) : (l <<= 1) {
            var offset: usize = 0;
            while (offset < N - l) : (offset += 2 * l) {
                const zeta: u64 = inv_zetas[k];
                k += 1;

                for (offset..offset + l) |j| {
                    const t = p.cs[j];
                    p.cs[j] = t +% p.cs[j + l];
                    p.cs[j + l] = montReduceLe2Q(zeta * @as(u64, t +% 256 * Q -% p.cs[j + l]));
                }
            }
        }

        for (0..N) |j| {
            p.cs[j] = montReduceLe2Q(@as(u64, R_OVER_256) * @as(u64, p.cs[j]));
        }
    }

    /// Apply Power2Round to all coefficients
    /// Returns both t0 and t1 polynomials
    fn power2RoundPoly(p: Poly) struct { t0: Poly, t1: Poly } {
        var t0 = Poly.zero;
        var t1 = Poly.zero;
        for (0..N) |i| {
            const result = power2Round(p.cs[i]);
            t0.cs[i] = result.a0_plus_q;
            t1.cs[i] = result.a1;
        }
        return .{ .t0 = t0, .t1 = t1 };
    }

    // Check if infinity norm exceeds bound
    fn exceeds(p: Poly, bound: u32) bool {
        var result: u32 = 0;
        for (0..N) |i| {
            const x = @as(i32, @intCast((Q - 1) / 2)) - @as(i32, @intCast(p.cs[i]));
            const abs_x = x ^ (x >> 31);
            const norm = @as(i32, @intCast((Q - 1) / 2)) - abs_x;
            const exceeds_bit = @intFromBool(@as(u32, @intCast(norm)) >= bound);
            result |= exceeds_bit;
        }
        return result != 0;
    }
}