Mat

fn Mat(comptime k: u8, comptime l: u8) type

File

Code

fn Mat(comptime k: u8, comptime l: u8) type {
    return struct {
        rows: [k]PolyVec(l),

        const Self = @This();
        const VecL = PolyVec(l);
        const VecK = PolyVec(k);

        /// Expand matrix A from seed rho using SHAKE-128
        /// This is the ExpandA function from FIPS 204
        fn derive(rho: *const [32]u8) Self {
            var m: Self = undefined;
            for (0..k) |i| {
                if (i + 1 < k) {
                    @prefetch(&m.rows[i + 1], .{ .rw = .write, .locality = 2 });
                }
                for (0..l) |j| {
                    // Nonce is i*256 + j
                    const nonce: u16 = (@as(u16, @intCast(i)) << 8) | @as(u16, @intCast(j));
                    m.rows[i].ps[j] = polyDeriveUniform(rho, nonce);
                }
            }
            return m;
        }

        /// Multiply matrix by vector in NTT domain and return result in regular domain.
        /// Takes a vector in NTT form and returns the product in regular form.
        fn mulVec(self: Self, v_hat: VecL) VecK {
            var result = VecK.zero;
            for (0..k) |i| {
                result.ps[i] = dotHat(l, self.rows[i], v_hat);
                result.ps[i] = result.ps[i].reduceLe2Q();
                result.ps[i] = result.ps[i].invNTT();
            }
            return result;
        }

        /// Multiply matrix by vector in NTT domain and return result in NTT domain.
        /// Takes a vector in NTT form and returns the product in NTT form.
        fn mulVecHat(self: Self, v_hat: VecL) VecK {
            var result: VecK = undefined;
            for (0..k) |i| {
                result.ps[i] = dotHat(l, self.rows[i], v_hat);
            }
            return result;
        }
    };
}