Poly

const Poly = struct

File

Code

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

    const encoded_length = N / 2 * 3;
    const zero: Poly = .{ .cs = @splat(0) };

    // Add two polynomials (coefficients not normalized)
    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 (coefficients not normalized)
    fn sub(a: Poly, b: Poly) Poly {
        var ret: Poly = undefined;
        for (0..N) |i| {
            ret.cs[i] = a.cs[i] - b.cs[i];
        }
        return ret;
    }

    // Executes a forward "NTT" on p.
    //
    // Assumes the coefficients are in absolute value ≤q.  The resulting
    // coefficients are in absolute value ≤7q.  If the input is in Montgomery
    // form, then the result is in Montgomery form and so (by linearity of the NTT)
    // if the input is in regular form, then the result is also in regular form.
    fn ntt(a: Poly) Poly {
        // Note that ℤ_q does not have a primitive 512ᵗʰ root of unity (as 512
        // does not divide into q-1) and so we cannot do a regular NTT.  ℤ_q
        // does have a primitive 256ᵗʰ root of unity, the smallest of which
        // is ζ := 17.
        //
        // Recall that our base ring R := ℤ_q[x] / (x²⁵⁶ + 1).  The polynomial
        // x²⁵⁶+1 will not split completely (as its roots would be 512ᵗʰ roots
        // of unity.)  However, it does split almost (using ζ¹²⁸ = -1):
        //
        // x²⁵⁶ + 1 = (x²)¹²⁸ - ζ¹²⁸
        //          = ((x²)⁶⁴ - ζ⁶⁴)((x²)⁶⁴ + ζ⁶⁴)
        //          = ((x²)³² - ζ³²)((x²)³² + ζ³²)((x²)³² - ζ⁹⁶)((x²)³² + ζ⁹⁶)
        //          ⋮
        //          = (x² - ζ)(x² + ζ)(x² - ζ⁶⁵)(x² + ζ⁶⁵) … (x² + ζ¹²⁷)
        //
        // Note that the powers of ζ that appear (from the second line down) are
        // in binary
        //
        // 0100000 1100000
        // 0010000 1010000 0110000 1110000
        // 0001000 1001000 0101000 1101000 0011000 1011000 0111000 1111000
        //         …
        //
        // That is: brv(2), brv(3), brv(4), …, where brv(x) denotes the 7-bit
        // bitreversal of x.  These powers of ζ are given by the Zetas array.
        //
        // The polynomials x² ± ζⁱ are irreducible and coprime, hence by
        // the Chinese Remainder Theorem we know
        //
        //  ℤ_q[x]/(x²⁵⁶+1) → ℤ_q[x]/(x²-ζ) x … x  ℤ_q[x]/(x²+ζ¹²⁷)
        //
        // given by a ↦ ( a mod x²-ζ, …, a mod x²+ζ¹²⁷ )
        // is an isomorphism, which is the "NTT".  It can be efficiently computed by
        //
        //
        //  a ↦ ( a mod (x²)⁶⁴ - ζ⁶⁴, a mod (x²)⁶⁴ + ζ⁶⁴ )
        //    ↦ ( a mod (x²)³² - ζ³², a mod (x²)³² + ζ³²,
        //        a mod (x²)⁹⁶ - ζ⁹⁶, a mod (x²)⁹⁶ + ζ⁹⁶ )
        //
        //      et cetera
        // If N was 8 then this can be pictured in the following diagram:
        //
        //  https://cnx.org/resources/17ee4dfe517a6adda05377b25a00bf6e6c93c334/File0026.png
        //
        // Each cross is a Cooley-Tukey butterfly: it's the map
        //
        //  (a, b) ↦ (a + ζb, a - ζb)
        //
        // for the appropriate power ζ for that column and row group.
        var p = a;
        var k: usize = 0; // index into zetas

        var l = N >> 1;
        while (l > 1) : (l >>= 1) {
            // On the nᵗʰ iteration of the l-loop, the absolute value of the
            // coefficients are bounded by nq.

            // offset effectively loops over the row groups in this column; it is
            // the first row in the row group.
            var offset: usize = 0;
            while (offset < N - l) : (offset += 2 * l) {
                k += 1;
                const z = @as(i32, zetas[k]);

                // j loops over each butterfly in the row group.
                for (offset..offset + l) |j| {
                    const t = montReduce(z * @as(i32, p.cs[j + l]));
                    p.cs[j + l] = p.cs[j] - t;
                    p.cs[j] += t;
                }
            }
        }

        return p;
    }

    // Executes an inverse "NTT" on p and multiply by the Montgomery factor R.
    //
    // Assumes the coefficients are in absolute value ≤q.  The resulting
    // coefficients are in absolute value ≤q.  If the input is in Montgomery
    // form, then the result is in Montgomery form and so (by linearity)
    // if the input is in regular form, then the result is also in regular form.
    fn invNTT(a: Poly) Poly {
        var k: usize = 127; // index into zetas
        var r: usize = 0; // index into invNTTReductions
        var p = a;

        // We basically do the oppposite of NTT, but postpone dividing by 2 in the
        // inverse of the Cooley-Tukey butterfly and accumulate that into a big
        // division by 2⁷ at the end.  See the comments in the ntt() function.

        var l: usize = 2;
        while (l < N) : (l <<= 1) {
            var offset: usize = 0;
            while (offset < N - l) : (offset += 2 * l) {
                // As we're inverting, we need powers of ζ⁻¹ (instead of ζ).
                // To be precise, we need ζᵇʳᵛ⁽ᵏ⁾⁻¹²⁸. However, as ζ⁻¹²⁸ = -1,
                // we can use the existing zetas table instead of
                // keeping a separate invZetas table as in Dilithium.

                const minZeta = @as(i32, zetas[k]);
                k -= 1;

                for (offset..offset + l) |j| {
                    // Gentleman-Sande butterfly: (a, b) ↦ (a + b, ζ(a-b))
                    const t = p.cs[j + l] - p.cs[j];
                    p.cs[j] += p.cs[j + l];
                    p.cs[j + l] = montReduce(minZeta * @as(i32, t));

                    // Note that if we had |a| < αq and |b| < βq before the
                    // butterfly, then now we have |a| < (α+β)q and |b| < q.
                }
            }

            // We let the invNTTReductions instruct us which coefficients to
            // Barrett reduce.
            while (true) {
                const i = inv_ntt_reductions[r];
                r += 1;
                if (i < 0) {
                    break;
                }
                p.cs[@as(usize, @intCast(i))] = feBarrettReduce(p.cs[@as(usize, @intCast(i))]);
            }
        }

        for (0..N) |j| {
            // Note 1441 = (128)⁻¹ R².  The coefficients are bounded by 9q, so
            // as 1441 * 9 ≈ 2¹⁴ < 2¹⁵, we're within the required bounds
            // for montReduce().
            p.cs[j] = montReduce(r2_over_128 * @as(i32, p.cs[j]));
        }

        return p;
    }

    // Normalizes coefficients.
    //
    // Ensures each coefficient is in {0, …, q-1}.
    fn normalize(a: Poly) Poly {
        var ret: Poly = undefined;
        for (0..N) |i| {
            ret.cs[i] = csubq(feBarrettReduce(a.cs[i]));
        }
        return ret;
    }

    // Put p in Montgomery form.
    fn toMont(a: Poly) Poly {
        var ret: Poly = undefined;
        for (0..N) |i| {
            ret.cs[i] = feToMont(a.cs[i]);
        }
        return ret;
    }

    // Barret reduce coefficients.
    //
    // Beware, this does not fully normalize coefficients.
    fn barrettReduce(a: Poly) Poly {
        var ret: Poly = undefined;
        for (0..N) |i| {
            ret.cs[i] = feBarrettReduce(a.cs[i]);
        }
        return ret;
    }

    fn compressedSize(comptime d: u8) usize {
        return @divTrunc(N * d, 8);
    }

    // Returns packed Compress_q(p, d).
    //
    // Assumes p is normalized.
    fn compress(p: Poly, comptime d: u8) [compressedSize(d)]u8 {
        @setEvalBranchQuota(10000);
        const q_over_2: u32 = comptime @divTrunc(Q, 2); // (q-1)/2
        const two_d_min_1: u32 = comptime (1 << d) - 1; // 2ᵈ-1
        var in_off: usize = 0;
        var out_off: usize = 0;

        const batch_size: usize = comptime math.lcm(d, 8);
        const in_batch_size: usize = comptime batch_size / d;
        const out_batch_size: usize = comptime batch_size / 8;

        const out_length: usize = comptime @divTrunc(N * d, 8);
        comptime assert(out_length * 8 == d * N);
        var out: [out_length]u8 = @splat(0);

        while (in_off < N) {
            // First we compress into in.
            var in: [in_batch_size]u16 = undefined;
            inline for (0..in_batch_size) |i| {
                // Compress_q(x, d) = ⌈(2ᵈ/q)x⌋ mod⁺ 2ᵈ
                //                  = ⌊(2ᵈ/q)x+½⌋ mod⁺ 2ᵈ
                //                  = ⌊((x << d) + q/2) / q⌋ mod⁺ 2ᵈ
                //                  = DIV((x << d) + q/2, q) & ((1<<d) - 1)
                const t = @as(u24, @intCast(p.cs[in_off + i])) << d;
                // Division by invariant multiplication, equivalent to DIV(t + q/2, q).
                // A division may not be a constant-time operation, even with a constant denominator.
                // Here, side channels would leak information about the shared secret, see https://kyberslash.cr.yp.to
                // Multiplication, on the other hand, is a constant-time operation on the CPUs we currently support.
                comptime assert(d <= 11);
                comptime assert(((20642679 * @as(u64, Q)) >> 36) == 1);
                const u: u32 = @intCast((@as(u64, t + q_over_2) * 20642679) >> 36);
                in[i] = @intCast(u & two_d_min_1);
            }

            // Now we pack the d-bit integers from `in' into out as bytes.
            comptime var in_shift: usize = 0;
            comptime var j: usize = 0;
            comptime var i: usize = 0;
            inline while (i < in_batch_size) : (j += 1) {
                comptime var todo: usize = 8;
                inline while (todo > 0) {
                    const out_shift = comptime 8 - todo;
                    out[out_off + j] |= @as(u8, @truncate((in[i] >> in_shift) << out_shift));

                    const done = comptime @min(@min(d, todo), d - in_shift);
                    todo -= done;
                    in_shift += done;

                    if (in_shift == d) {
                        in_shift = 0;
                        i += 1;
                    }
                }
            }

            in_off += in_batch_size;
            out_off += out_batch_size;
        }

        return out;
    }

    // Set p to Decompress_q(m, d).
    fn decompress(comptime d: u8, in: *const [compressedSize(d)]u8) Poly {
        @setEvalBranchQuota(10000);
        const in_len = comptime @divTrunc(N * d, 8);
        comptime assert(in_len * 8 == d * N);
        var ret: Poly = undefined;
        var in_off: usize = 0;
        var out_off: usize = 0;

        const batch_size: usize = comptime math.lcm(d, 8);
        const in_batch_size: usize = comptime batch_size / 8;
        const out_batch_size: usize = comptime batch_size / d;

        while (out_off < N) {
            comptime var in_shift: usize = 0;
            comptime var j: usize = 0;
            comptime var i: usize = 0;
            inline while (i < out_batch_size) : (i += 1) {
                // First, unpack next coefficient.
                comptime var todo = d;
                var out: u16 = 0;

                inline while (todo > 0) {
                    const out_shift = comptime d - todo;
                    const m = comptime (1 << d) - 1;
                    out |= (@as(u16, in[in_off + j] >> in_shift) << out_shift) & m;

                    const done = comptime @min(@min(8, todo), 8 - in_shift);
                    todo -= done;
                    in_shift += done;

                    if (in_shift == 8) {
                        in_shift = 0;
                        j += 1;
                    }
                }

                // Decompress_q(x, d) = ⌈(q/2ᵈ)x⌋
                //                    = ⌊(q/2ᵈ)x+½⌋
                //                    = ⌊(qx + 2ᵈ⁻¹)/2ᵈ⌋
                //                    = (qx + (1<<(d-1))) >> d
                const qx = @as(u32, out) * @as(u32, Q);
                ret.cs[out_off + i] = @as(i16, @intCast((qx + (1 << (d - 1))) >> d));
            }

            in_off += in_batch_size;
            out_off += out_batch_size;
        }

        return ret;
    }

    // Returns the "pointwise" multiplication a o b.
    //
    // That is: invNTT(a o b) = invNTT(a) * invNTT(b).  Assumes a and b are in
    // Montgomery form.  Products between coefficients of a and b must be strictly
    // bounded in absolute value by 2¹⁵q.  a o b will be in Montgomery form and
    // bounded in absolute value by 2q.
    fn mulHat(a: Poly, b: Poly) Poly {
        // Recall from the discussion in ntt(), that a transformed polynomial is
        // an element of ℤ_q[x]/(x²-ζ) x … x  ℤ_q[x]/(x²+ζ¹²⁷);
        // that is: 128 degree-one polynomials instead of simply 256 elements
        // from ℤ_q as in the regular NTT.  So instead of pointwise multiplication,
        // we multiply the 128 pairs of degree-one polynomials modulo the
        // right equation:
        //
        //  (a₁ + a₂x)(b₁ + b₂x) = a₁b₁ + a₂b₂ζ' + (a₁b₂ + a₂b₁)x,
        //
        // where ζ' is the appropriate power of ζ.

        var p: Poly = undefined;
        var k: usize = 64;
        var i: usize = 0;
        while (i < N) : (i += 4) {
            const z = @as(i32, zetas[k]);
            k += 1;

            const a1b1 = montReduce(@as(i32, a.cs[i + 1]) * @as(i32, b.cs[i + 1]));
            const a0b0 = montReduce(@as(i32, a.cs[i]) * @as(i32, b.cs[i]));
            const a1b0 = montReduce(@as(i32, a.cs[i + 1]) * @as(i32, b.cs[i]));
            const a0b1 = montReduce(@as(i32, a.cs[i]) * @as(i32, b.cs[i + 1]));

            p.cs[i] = montReduce(a1b1 * z) + a0b0;
            p.cs[i + 1] = a0b1 + a1b0;

            const a3b3 = montReduce(@as(i32, a.cs[i + 3]) * @as(i32, b.cs[i + 3]));
            const a2b2 = montReduce(@as(i32, a.cs[i + 2]) * @as(i32, b.cs[i + 2]));
            const a3b2 = montReduce(@as(i32, a.cs[i + 3]) * @as(i32, b.cs[i + 2]));
            const a2b3 = montReduce(@as(i32, a.cs[i + 2]) * @as(i32, b.cs[i + 3]));

            p.cs[i + 2] = a2b2 - montReduce(a3b3 * z);
            p.cs[i + 3] = a2b3 + a3b2;
        }

        return p;
    }

    // Sample p from a centered binomial distribution with n=2η and p=½ - viz:
    // coefficients are in {-η, …, η} with probabilities
    //
    //  {ncr(0, 2η)/2^2η, ncr(1, 2η)/2^2η, …, ncr(2η,2η)/2^2η}
    fn noise(comptime eta: u8, nonce: u8, seed: *const [32]u8) Poly {
        var h = sha3.Shake256.init(.{});
        const suffix: [1]u8 = .{nonce};
        h.update(seed);
        h.update(&suffix);

        // The distribution at hand is exactly the same as that
        // of (a₁ + a₂ + … + a_η) - (b₁ + … + b_η) where a_i,b_i~U(1).
        // Thus we need 2η bits per coefficient.
        const buf_len = comptime 2 * eta * N / 8;
        var buf: [buf_len]u8 = undefined;
        h.squeeze(&buf);

        // buf is interpreted as a₁…a_ηb₁…b_ηa₁…a_ηb₁…b_η…. We process
        // multiple coefficients in one batch.

        const T = switch (builtin.target.cpu.arch) {
            .x86_64, .x86 => u32, // Generates better code on Intel CPUs
            else => u64, // u128 might be faster on some other CPUs.
        };

        comptime var batch_count: usize = undefined;
        comptime var batch_bytes: usize = undefined;
        comptime var mask: T = 0;
        comptime {
            batch_count = @bitSizeOf(T) / @as(usize, 2 * eta);
            while (@rem(N, batch_count) != 0 and batch_count > 0) : (batch_count -= 1) {}
            assert(batch_count > 0);
            assert(@rem(2 * eta * batch_count, 8) == 0);
            batch_bytes = 2 * eta * batch_count / 8;

            for (0..2 * eta * batch_count) |_| {
                mask <<= eta;
                mask |= 1;
            }
        }

        var ret: Poly = undefined;
        for (0..comptime N / batch_count) |i| {
            // Read coefficients into t. In the case of η=3,
            // we have t = a₁ + 2a₂ + 4a₃ + 8b₁ + 16b₂ + …
            var t: T = 0;
            inline for (0..batch_bytes) |j| {
                t |= @as(T, buf[batch_bytes * i + j]) << (8 * j);
            }

            // Accumulate `a's and `b's together by masking them out, shifting
            // and adding. For η=3, we have  d = a₁ + a₂ + a₃ + 8(b₁ + b₂ + b₃) + …
            var d: T = 0;
            inline for (0..eta) |j| {
                d += (t >> j) & mask;
            }

            // Extract each a and b separately and set coefficient in polynomial.
            inline for (0..batch_count) |j| {
                const mask2 = comptime (1 << eta) - 1;
                const a = @as(i16, @intCast((d >> (comptime (2 * j * eta))) & mask2));
                const b = @as(i16, @intCast((d >> (comptime ((2 * j + 1) * eta))) & mask2));
                ret.cs[batch_count * i + j] = a - b;
            }
        }

        return ret;
    }

    fn uniform(seed: [32]u8, x: u8, y: u8) Poly {
        const domain_sep: [2]u8 = .{ x, y };
        return sampleUniformRejection(
            Poly,
            Q,
            12,
            N,
            &seed,
            &domain_sep,
        );
    }

    // Packs p.
    //
    // Assumes p is normalized (and not just Barrett reduced).
    fn toBytes(p: Poly) [encoded_length]u8 {
        var ret: [encoded_length]u8 = undefined;
        for (0..comptime N / 2) |i| {
            const t0 = @as(u16, @intCast(p.cs[2 * i]));
            const t1 = @as(u16, @intCast(p.cs[2 * i + 1]));
            ret[3 * i] = @as(u8, @truncate(t0));
            ret[3 * i + 1] = @as(u8, @truncate((t0 >> 8) | (t1 << 4)));
            ret[3 * i + 2] = @as(u8, @truncate(t1 >> 4));
        }
        return ret;
    }

    // Unpacks a Poly from buf.
    //
    // p will not be normalized; instead 0 ≤ p[i] < 4096.
    fn fromBytes(buf: *const [encoded_length]u8) Poly {
        var ret: Poly = undefined;
        for (0..comptime N / 2) |i| {
            const b0 = @as(i16, buf[3 * i]);
            const b1 = @as(i16, buf[3 * i + 1]);
            const b2 = @as(i16, buf[3 * i + 2]);
            ret.cs[2 * i] = b0 | ((b1 & 0xf) << 8);
            ret.cs[2 * i + 1] = (b1 >> 4) | b2 << 4;
        }
        return ret;
    }
}