@noble/post-quantum 0.5.4 → 0.6.1

package/src/_crystals.ts

@@ -6,54 +6,118 @@
 import { FFTCore, reverseBits } from '@noble/curves/abstract/fft.js';
 import { shake128, shake256 } from '@noble/hashes/sha3.js';
 import type { TypedArray } from '@noble/hashes/utils.js';
-import { type BytesCoderLen, cleanBytes, type Coder, getMask } from './utils.ts';
+import {
+  type BytesCoderLen,
+  cleanBytes,
+  type Coder,
+  getMask,
+  type TArg,
+  type TRet,
+} from './utils.ts';
 
+/** Extendable-output reader used by the CRYSTALS implementations. */
 export type XOF = (
   seed: Uint8Array,
   blockLen?: number
 ) => {
+  /**
+   * Read diagnostic counters for the current XOF session.
+   * @returns Current call and XOF block counters.
+   */
   stats: () => { calls: number; xofs: number };
+  /**
+   * Select one `(x, y)` coordinate pair and get a block reader for it.
+   * Only one coordinate stream is live at a time: a later `get(...)` call rebinds the shared
+   * SHAKE state and invalidates older readers.
+   * Each squeeze aliases one mutable internal output buffer, so callers must copy blocks they
+   * want to retain before the next read.
+   * @param x - First matrix coordinate.
+   * @param y - Second matrix coordinate.
+   * @returns Lazy block reader for that coordinate pair.
+   */
   get: (x: number, y: number) => () => Uint8Array; // return block aligned to blockLen and 3
+  /** Wipe any buffered state once the reader is no longer needed. */
   clean: () => void;
 };
 
 /** CRYSTALS (ml-kem, ml-dsa) options */
+/** Shared polynomial and NTT parameters for CRYSTALS algorithms. */
 export type CrystalOpts<T extends TypedArray> = {
+  /**
+   * Allocate one zeroed polynomial/vector container.
+   * @param n - Number of coefficients to allocate.
+   * @returns Fresh typed container.
+   */
   newPoly: TypedCons<T>;
-  N: number; // poly size, 256
-  Q: number; // modulo
-  F: number; // 256**−1 mod q for dilithium, 128**−1 mod q for kyber
+  /** Polynomial size, typically `256`. */
+  N: number;
+  /** Prime modulus used for all coefficient arithmetic. */
+  Q: number;
+  /** Inverse transform normalization factor:
+   * `256**-1 mod q` for Dilithium, `128**-1 mod q` for Kyber.
+   */
+  F: number;
+  /** Principal root of unity for the transform domain. */
   ROOT_OF_UNITY: number;
-  brvBits: number; // bits for bitReversal
+  /** Number of bits used for bit-reversal ordering. */
+  brvBits: number;
+  /** `true` for Kyber/ML-KEM mode, `false` for Dilithium/ML-DSA mode. */
   isKyber: boolean;
 };
 
+/** Constructor function for typed polynomial containers. */
 export type TypedCons<T extends TypedArray> = (n: number) => T;
 
-export const genCrystals = <T extends TypedArray>(
-  opts: CrystalOpts<T>
-): {
+type Crystals<T extends TypedArray> = {
   mod: (a: number, modulo?: number) => number;
   smod: (a: number, modulo?: number) => number;
   nttZetas: T;
   NTT: {
+    /** Forward transform in place. Mutates and returns `r`. */
     encode: (r: T) => T;
+    /** Inverse transform in place. Mutates and returns `r`. */
     decode: (r: T) => T;
   };
   bitsCoder: (d: number, c: Coder<number, number>) => BytesCoderLen<T>;
-} => {
+};
+
+/**
+ * Creates shared modular arithmetic, NTT, and packing helpers for CRYSTALS schemes.
+ * @param opts - Polynomial and transform parameters. See {@link CrystalOpts}.
+ * @returns CRYSTALS arithmetic and encoding helpers.
+ * @example
+ * Create shared modular arithmetic and NTT helpers for a CRYSTALS parameter set.
+ * ```ts
+ * const crystals = genCrystals({
+ *   newPoly: (n) => new Uint16Array(n),
+ *   N: 256,
+ *   Q: 3329,
+ *   F: 3303,
+ *   ROOT_OF_UNITY: 17,
+ *   brvBits: 7,
+ *   isKyber: true,
+ * });
+ * const reduced = crystals.mod(-1);
+ * ```
+ */
+export const genCrystals = <T extends TypedArray>(opts: CrystalOpts<T>): TRet<Crystals<T>> => {
   // isKyber: true means Kyber, false means Dilithium
   const { newPoly, N, Q, F, ROOT_OF_UNITY, brvBits, isKyber } = opts;
+  // Normalize JS `%` into the canonical Z_m representative `[0, modulo-1]` expected by
+  // FIPS 203 §2.3 / FIPS 204 §2.3 before downstream mod-q arithmetic.
   const mod = (a: number, modulo = Q): number => {
     const result = a % modulo | 0;
     return (result >= 0 ? result | 0 : (modulo + result) | 0) | 0;
   };
-  // -(Q-1)/2 < a <= (Q-1)/2
+  // FIPS 204 §7.4 uses the centered `mod ±` representative for low bits, keeping the
+  // positive midpoint when `modulo` is even.
+  // Center to `[-floor((modulo-1)/2), floor(modulo/2)]`.
   const smod = (a: number, modulo = Q): number => {
     const r = mod(a, modulo) | 0;
     return (r > modulo >> 1 ? (r - modulo) | 0 : r) | 0;
   };
-  // Generate zettas (different from roots of unity, negacyclic uses phi, where acyclic uses omega)
+  // Kyber uses the FIPS 203 Appendix A `BitRev_7` table here via the first 128 entries, while
+  // Dilithium uses the FIPS 204 §7.5 / Appendix B `BitRev_8` zetas table over all 256 entries.
   function getZettas() {
     const out = newPoly(N);
     for (let i = 0; i < N; i++) {
@@ -94,44 +158,57 @@ export const genCrystals = <T extends TypedArray>(
     },
     decode: (r: T): T => {
       dit(r as any);
+      // The inverse-NTT normalization factor is family-specific: FIPS 203 Algorithm 10 line 14
+      // uses `128^-1 mod q` for Kyber, while FIPS 204 Algorithm 42 lines 21-23 use `256^-1 mod q`.
       // kyber uses 128 here, because brv && stuff
       for (let i = 0; i < r.length; i++) r[i] = mod(F * r[i]);
       return r;
     },
   };
-  // Encode polynominal as bits
-  const bitsCoder = (d: number, c: Coder<number, number>): BytesCoderLen<T> => {
+  // Pack one little-endian `d`-bit word per coefficient, matching FIPS 203 ByteEncode /
+  // ByteDecode and the FIPS 204 BitsToBytes-based polynomial packing helpers.
+  const bitsCoder = (d: number, c: Coder<number, number>): TRet<BytesCoderLen<T>> => {
     const mask = getMask(d);
     const bytesLen = d * (N / 8);
     return {
       bytesLen,
-      encode: (poly: T): Uint8Array => {
+      encode: (poly_: TArg<T>): TRet<Uint8Array> => {
+        const poly = poly_ as T;
         const r = new Uint8Array(bytesLen);
         for (let i = 0, buf = 0, bufLen = 0, pos = 0; i < poly.length; i++) {
           buf |= (c.encode(poly[i]) & mask) << bufLen;
           bufLen += d;
           for (; bufLen >= 8; bufLen -= 8, buf >>= 8) r[pos++] = buf & getMask(bufLen);
         }
-        return r;
+        return r as TRet<Uint8Array>;
       },
-      decode: (bytes: Uint8Array): T => {
+      decode: (bytes: TArg<Uint8Array>): TRet<T> => {
         const r = newPoly(N);
         for (let i = 0, buf = 0, bufLen = 0, pos = 0; i < bytes.length; i++) {
           buf |= bytes[i] << bufLen;
           bufLen += 8;
           for (; bufLen >= d; bufLen -= d, buf >>= d) r[pos++] = c.decode(buf & mask);
         }
-        return r;
+        return r as TRet<T>;
       },
-    };
+    } as TRet<BytesCoderLen<T>>;
   };
 
-  return { mod, smod, nttZetas, NTT, bitsCoder };
+  return {
+    mod,
+    smod,
+    nttZetas: nttZetas as TRet<T>,
+    NTT: {
+      encode: (r: TArg<T>): TRet<T> => NTT.encode(r as T) as TRet<T>,
+      decode: (r: TArg<T>): TRet<T> => NTT.decode(r as T) as TRet<T>,
+    },
+    bitsCoder: bitsCoder as TRet<Crystals<T>>['bitsCoder'],
+  };
 };
 
 const createXofShake =
-  (shake: typeof shake128): XOF =>
-  (seed: Uint8Array, blockLen?: number) => {
+  (shake: typeof shake128): TRet<XOF> =>
+  (seed: TArg<Uint8Array>, blockLen?: number) => {
     if (!blockLen) blockLen = shake.blockLen;
     // Optimizations that won't mater:
     // - cached seed update (two .update(), on start and on the end)
@@ -148,6 +225,8 @@ const createXofShake =
     return {
       stats: () => ({ calls, xofs }),
       get: (x: number, y: number) => {
+        // Rebind to `seed || x || y` so callers can implement the spec's per-coordinate
+        // SHAKE inputs like `rho || j || i` and `rho || IntegerToBytes(counter, 2)`.
         _seed[seedLen + 0] = x;
         _seed[seedLen + 1] = y;
         h.destroy();
@@ -155,7 +234,7 @@ const createXofShake =
         calls++;
         return () => {
           xofs++;
-          return h.xofInto(buf);
+          return h.xofInto(buf) as TRet<Uint8Array>;
         };
       },
       clean: () => {
@@ -165,5 +244,37 @@ const createXofShake =
     };
   };
 
-export const XOF128: XOF = /* @__PURE__ */ createXofShake(shake128);
-export const XOF256: XOF = /* @__PURE__ */ createXofShake(shake256);
+/**
+ * SHAKE128-based extendable-output reader factory used by ML-KEM.
+ * `get(x, y)` selects one coordinate pair at a time; calling it again invalidates previously
+ * returned readers, and each squeeze reuses one mutable internal output buffer.
+ * @param seed - Seed bytes for the reader.
+ * @param blockLen - Optional output block length.
+ * @returns Stateful XOF reader.
+ * @example
+ * Build the ML-KEM SHAKE128 matrix expander and read one block.
+ * ```ts
+ * import { randomBytes } from '@noble/post-quantum/utils.js';
+ * import { XOF128 } from '@noble/post-quantum/_crystals.js';
+ * const reader = XOF128(randomBytes(32));
+ * const block = reader.get(0, 0)();
+ * ```
+ */
+export const XOF128: TRet<XOF> = /* @__PURE__ */ createXofShake(shake128);
+/**
+ * SHAKE256-based extendable-output reader factory used by ML-DSA.
+ * `get(x, y)` appends raw one-byte coordinates to the seed, invalidates previously returned
+ * readers, and reuses one mutable internal output buffer for each squeeze.
+ * @param seed - Seed bytes for the reader.
+ * @param blockLen - Optional output block length.
+ * @returns Stateful XOF reader.
+ * @example
+ * Build the ML-DSA SHAKE256 coefficient expander and read one block.
+ * ```ts
+ * import { randomBytes } from '@noble/post-quantum/utils.js';
+ * import { XOF256 } from '@noble/post-quantum/_crystals.js';
+ * const reader = XOF256(randomBytes(32));
+ * const block = reader.get(0, 0)();
+ * ```
+ */
+export const XOF256: TRet<XOF> = /* @__PURE__ */ createXofShake(shake256);