Specification.cry

/**
 * ML-KEM key encapsulation mechanism for establishing a shared secret key
 * between two parties, with a fast NTT implementation.
 *
 * This implements the final version of ML-KEM as standardized in [FIPS-203].
 * The executable spec differs from the spec in several important ways:
 * - In the spec, two ML-KEM functions (`KeyGen` and `Encaps`) require the
 *   generation of randomness. In the exectable spec, this
 *   randomness is passed as input, because Cryptol cannot generate randomness.
 *   Implementors must manually verify that their functions generate randomness
 *   using an approved RBG with a suitable security strength for their chosen
 *   parameter set, and that the sampling of random values is performed by the
 *   cryptographic module.
 * - The spec requires that intermediate values must be destroyed prior to the
 *   termination of each algorithm, with two exceptions (see Section 3.3
 *   "Destruction of intermediate values.") Cryptol cannot express this
 *   requirement. Implementors must manually verify that intermediate data is
 *   destroyed as soon as it is no longer needed.
 *
 * The spec defines other requirements that cannot be verified by Cryptol!
 * For full details, see [FIPS-203] Section 3.3.
 * - K-PKE cannot be used as a stand-alone cryptographic scheme. This
 *   executable spec makes the K-PKE module private; implementors must manually
 *   verify that their own implementations do not make the K-PKE API public.
 * - The internal functions `KeyGen_internal`, `Encaps_internal`, and
 *   `Decaps_internal` must not be made available to other applications except
 *   for testing purposes. This executable spec makes the internal functions
 *   private; implementors must manually verify that the same is true in their
 *   implementations.
 * - The shared secret key generated by `Encaps` and `Decaps` can be used
 *   directly as a key for symmetric cryptography. If further key derivation is
 *   needed, it must be done using an approved technique. Cryptol cannot verify
 *   that the shared secret key is used in an appropriate way.
 * - As mentioned, randomness must be generated using an approved RBG, and a
 *   fresh string of random bytes must be generated for each invocation.
 *   Cryptol cannot verify either of these requirements.
 * - The input to `Encaps` and `Decaps` must be checked. Cryptol cannot verify
 *   that the input has been checked.
 * - As mentioned, intermediate values must be destroyeda fter use. Cryptol
 *   cannot verify that this has been done.
 *
 * @copyright Galois, Inc
 * @author Marios Georgiou <marios@galois.com>
 * @editor Marcella Hastings <marcella@galois.com>
 *
 * @copyright Amazon.com or its affiliates.
 * @author Rod Chapman <rodchap@amazon.com>
 *
 * Sources:
 * [FIPS-203]: National Institute of Standards and Technology. Module-Lattice-
 *     Based Key-Encapsulation Mechanism Standard. (Department of Commerce,
 *     Washington, D.C.), Federal Information Processing Standards Publication
 *     (FIPS) NIST FIPS 203. August 2024.
 *     @see https://doi.org/10.6028/NIST.FIPS.203
 */
module Primitive::Asymmetric::KEM::ML_KEM::Specification where

import Primitive::Keyless::Hash::SHAKE::SHAKE256 as SHAKE256
import Primitive::Keyless::Hash::SHAKE::SHAKE128 as SHAKE128
import Primitive::Keyless::Hash::SHA3::SHA3_256 as SHA3_256
import Primitive::Keyless::Hash::SHA3::SHA3_512 as SHA3_512

private
    /*
     * [FIPS-203] Section 2.3.
     */
    type n = 256
    /*
     * [FIPS-203] Section 2.3.
     */
    type q = 3329
    /**
     * A primitive n-th root of unity modulo `q`.
     * [FIPS-203] Section 2.3.
     */
    zeta = 17 : Z q

/**
 * Defines the set {0, 1, ..., 255} of unsigned 8-bit integers.
 * [FIPS-203] Section 2.3.
 */
type Byte = [8]

private
    /**
     * An element in the ring `R_q`.
     *
     * An element in the ring is a polynomial of degree at most 255 (e.g. with
     * 256 terms). The `i`th element in this array represents the coefficient
     * of the degree-`i` term.
     *
     * [FIPS-203] Section 2.3 (definition of the ring).
     * [FIPS-203] Section 2.4.4, Equation 2.5 (definition of the representation
     *     of elements in the ring).
     */
    type Rq = [n](Z q)

    /**
     * An element in the ring `T_q`.
     *
     * An element in this ring (sometimes called the "NTT representation") is a
     * tuple of 128 polynomials, each of degree at most one (e.g. with two
     * terms). The `2i` and `2i+1`th terms in this array represent the degree-0
     * and degree-1 coefficients of the `i`th polynomial, respectively.
     *
     * [FIPS-203] Section 2.3 (definition of the `T_q`).
     * [FIPS-203] Section 2.4.4 Equation 2.7 (definition of the representation
     *     of an element in `T_q`).
     */
    type Tq = [n](Z q)

    /**
     * Pseudorandom function (PRF).
     * [FIPS-203] Section 4.1, Equations 4.2 and 4.3.
     */
    PRF : {eta} (2 <= eta, eta <= 3) => [32]Byte -> Byte -> [64 * eta]Byte
    PRF s b = SHAKE256::xofBytes`{8 * 64 * eta} (s # [b])

    /**
     * One of the hash functions used in the protocol.
     * [FIPS-203] Section 4.1, Equation 4.4.
     */
    H : {hinl} (fin hinl) => [hinl]Byte -> [32]Byte
    H s = SHA3_256::hashBytes s

    /**
     * One of the hash functions used in the protocol.
     * [FIPS-203] Section 4.1, Equation 4.4.
     */
    J : {hinl} (fin hinl) => [hinl]Byte -> [32]Byte
    J s = SHAKE256::xofBytes`{8 * 32} s

    /**
     * One of the hash functions used in the protocol.
     * [FIPS-203] Section 4.1, Equation 4.5.
     */
    G : {ginl} (fin ginl) => [ginl]Byte -> ([32]Byte, [32]Byte)
    G c = (a, b) where
        [a, b] = split (SHA3_512::hashBytes c)

    /**
     * eXtendable-Output Function (XOF) wrapper.
     * [FIPS-203] Section 4.1, Equation 4.6.
     */
    XOF : ([34]Byte) -> [inf]Byte
    XOF = SHAKE128::xofBytes

    /**
     * Conversion from little-endian bit arrays to byte arrays.
     * [FIPS-203] Section 4.2.1, Algorithm 3.
     */
    BitsToBytes : {ell} (fin ell) => [8 * ell]Bit -> [ell]Byte
    BitsToBytes b
        | ell == 0 => zero
        | ell > 0 => B where
            // Group the bits into the B[⌊i / 8⌋] sets; pad them to support
            // subsequent operations, and correlate each bit with its index `i`.
            b' = groupBy`{8} [(zext [bi], i)
                | bi <- b
                | i <- [0..8 * ell - 1]]

            // Steps 2-4.
            B = [sum [bi * (2 ^^ (i % 8))
                    | (bi, i) <- bi8]
                | bi8 <- b']

    /**
     * Conversion from byte arrays to bit arrays.
     * [FIPS-203] Section 4.2.1, Algorithm 4.
     */
    BytesToBits : {ell} (fin ell) => [ell]Byte -> [ell*8]Bit
    BytesToBits C
        | ell == 0 => []
        | ell > 0 => join [[ b8ij where
                // Step 4. Taking the last bit is the same as modding by 2. (See
                // `mod2IsFinalBit`).
                b8ij = Ci' ! 0
                // Step 5. Shifting right is the same as the iterative
                // division (see `div2IsShiftR`). This accounts for all the
                // divisions "up to this point" (e.g. none when `j = 0`), which
                // is why we use `Ci'` to evaluate `b8ij` above.
                Ci' = Ci >> j
            // Step 3.
            | j <- [0..7]]
            // Step 2. We iterate over `C` directly instead of indexing into it.
            | Ci <- C ]

    /**
     * The iterative division by 2 in `BytesToBits` is the same as shifting
     * right.
     * ```repl
     * :prove div2IsShiftR
     * ```
     */
    div2IsShiftR : Byte -> Bit
    div2IsShiftR C = take (d2 C) == shl where
        // Note: division here is floor'd by default.
        d2 c = [c] # d2 (c / 2)
        shl = [C >> j | j <- [0..7]]

    /**
     * The conversions between bits and bytes are each others' inverses.
     * [FIPS-203] Section 4.2.1 (see description on Algorithm 4).
     * The sample `ell` values here are a subset of the possible values in the
     * spec.
     * ```repl
     * :prove B2B2BInverts`{32}
     * :prove B2B2BInverts`{192}
     * :prove B2B2BInverts`{384}
     * ```
     */
    B2B2BInverts : {ell} (fin ell, ell > 0) => [ell * 8] -> Bit
    B2B2BInverts bits = bitsWorks && bytesWorks where
        bitsWorks = BytesToBits (BitsToBytes bits) == bits
        bytesWorks = BitsToBytes (BytesToBits (split bits)) == split bits

    /**
     * Check the example given in the spec for converting between bits and
     * bytes.
     * [FIPS-203] Section 4.2.1 "Converting between bits and bytes."
     * ```repl
     * :prove B2BExampleWorks
     * ```
     */
    B2BExampleWorks = BitsToBytes 0b11010001 == [139]

    /**
     * In Cryptol, rounding is computed via the built-in function `roundAway`.
     * [FIPS-203] Section 2.3.
     */
    property roundingWorks y = y >= 0 ==> roundUpWorks && roundDownWorks where
        y' = fromInteger y
        roundUpWorks = roundAway (y' + 0.5) == (y + 1)
        roundDownWorks = roundAway (y' + 0.4) == y

    /**
     * Compress an integer mod `q` into an integer mod `2^d`.
     * [FIPS-203] Section 4.2.1, Equation 4.7.
     */
    Compress : {d} (d < width q) => Z q -> [d]
    Compress x = y where
        // Convert from an integer mod `q` to a rational number.
        x' = fromInteger (fromZ x) : Rational
        // Compress. Note that `/.` denotes division of rationals.
        y' = roundAway (((2^^`d) /. `q) * x')
        // mod 2^^d (by converting from an integer to a d-bit vector).
        y = (fromInteger y') : [d]

    /**
     * Decompress an integer mod `2^d` into an integer mod `q`.
     * [FIPS-203] Section 4.2.1, Equation 4.8.
     */
    Decompress : {d} (d < width q) => [d] -> Z q
    Decompress y = x where
        // Convert from a d-length bit vector to a rational number.
        y' = fromInteger (toInteger y) : Rational
        // Decompress! As before, `/.` is division of rationals.
        x' = roundAway((`q /. (2^^`d)) * y')
        // Convert from an integer to an integer mod `q`.
        x = (fromInteger x') : Z q

    /**
     * Compression inverts decompression for all inputs and bit lengths.
     * We'll prove it for the bit lengths found in the
     * ```repl
     * :prove CompressInvertsDecompress`{1}
     * :exhaust CompressInvertsDecompress`{d_u}
     * :exhaust CompressInvertsDecompress`{d_v}
     * ```
     */
    CompressInvertsDecompress : {d} (d < width q) => [d] -> Bit
    property CompressInvertsDecompress y = Compress (Decompress y) == y

    /**
     * When `d` is large, compression followed by decompression must not
     * significantly alter the value.
     * This sets `d = d_u`, which is the largest value for `d` used in the
     * spec.
     * [FIPS-203] Section 4.2.1, "Compression and Decompression".
     * ```repl
     * :exhaust DecompressMostlyInvertsCompress
     * ```
     */
    DecompressMostlyInvertsCompress : Z q -> Bit
    property DecompressMostlyInvertsCompress x = errIsSmallEnough where
        x' = Decompress`{d_u} (Compress`{d_u} x)
        err = abs (modpm (x' - x))
        errIsSmallEnough = err <= B_q`{d_u}

        // The spec doesn't describe formally what "not significantly altered"
        // means; we use this equation.
        B_q : {d} (d < lg2 q) => Integer
        B_q = roundAway((`q/.(2^^(`d+1))))

        // Convert an integer mod `q` to a representation centered around 0
        // (and represented as an `Integer`).
        modpm : Z q -> Integer
        modpm r = if r' > (`q / 2) then r' - `q else r'
            where r' = fromZ r

    /**
     * Compression applied to a vector is equivalent to applying compression to
     * each individual element.
     * [FIPS-203] Section 2.4.8, Equation 2.15.
     */
    Compress_Vec : {d} (d < lg2 q) => [n](Z q) -> [n][d]
    Compress_Vec x = map Compress`{d} x

    /**
     * Decompression applied to a vector is equivalent to applying
     * decompression to each individual element.
     * [FIPS-203] Section 2.4.8.
     */
    Decompress_Vec : {d} (d < lg2 q) => [n][d] -> [n](Z q)
    Decompress_Vec x = map Decompress`{d} x

    /**
     * Compression applied to a matrix is equivalent to applying compression
     * to each individual element.
     * [FIPS-203] Section 2.4.8.
     */
    Compress_Mat : {d, k1} (d < lg2 q, fin k1) => [k1][n](Z q) -> [k1][n][d]
    Compress_Mat x = map Compress_Vec`{d} x

    /**
     * Decompression applied to a matrix is equivalent to applying
     * decompression to each individual element.
     * [FIPS-203] Section 2.4.8.
     */
    Decompress_Mat : {d, k1} (d < lg2 q, fin k1) => [k1][n][d] -> [k1][n](Z q)
    Decompress_Mat x = map Decompress_Vec`{d} x

    /**
     * Encode an array of `d`-bit integers into a byte array, for `d < 12`.
     * [FIPS-203] Section 4.2.1 Algorithm 5.
     *
     * Note: In this implementation, we treat the `d < 12` case separately from
     * `d == 12` because it allows us to better express the different integer
     * types. For `d < 12`, we use the bit vector type.
     * See `ByteEncode12` for the `d = 12` case.
     */
    ByteEncode : {d} (1 <= d, d < 12) => [256][d] -> [32 * d]Byte
    ByteEncode F = B where
        // Step 1-3, 5.
        // We iterate over `F` directly, instead of indexing into it with `i`.
        // In a bit vector, iteratively subtracting the last bit and dividing
        // by 2 is the same as bit-shifting to the right (see
        // `subAndDivIsShift`).
        as = [[ a >> j | j <- [0..d-1]] | a <- F]
        // Step 4. In a bit vector, taking the value `% 2` is the same as taking
        // the final bit. See `mod2IsFinalBit`.
        b = join [ [ a ! 0 | a <- ajs] | ajs <- as]
        // Step 8.
        B = BitsToBytes b

    /**
     * Encode `k` vectors of `d`-bit integers into a byte array.
     * [FIPS-203] Section 2.4.8.
     */
    ByteEncode_Vec : {d} (fin d, 1 <= d, d < 12) => [k][256][d] -> [k * 32 * d]Byte
    ByteEncode_Vec F_vec = join (map ByteEncode F_vec)

    /**
     * Encode an array of integers mod `q` into a byte array.
     * [FIPS-203] Section 4.2.1 Algorithm 5.
     *
     * Note: In this implementation, we treat the `d < 12` case separately from
     * `d == 12` because it allows us to better express the different integer
     * types. For `d == 12`, we use the integers-mod (`Z`) type.
     * See `ByteEncode` for the `d < 12` case.
     */
    ByteEncode12 : [256](Z q) -> [32 * 12]Byte
    ByteEncode12 F = B where
        type d = 12

        // We need to explicitly convert from integers mod `q` to 12-bit
        // vectors.
        toBitVec : Z q -> [d]
        toBitVec f = fromInteger (fromZ f)
        F' = map toBitVec F

        // The following is the same as in `ByteEncode`. See notes above.
        as = [[ a >> j | j <- [0..d-1]] | a <- F' ]
        b = join [ [ a ! 0 | a <- ajs] | ajs <- as]
        B = BitsToBytes b

    /**
     * Encode a set of `k` vectors of integers mod `q` into a byte array.
     * [FIPS-203] Section 2.4.8.
     */
    ByteEncode12_Vec : [k][256](Z q) -> [k * 32 * 12]Byte
    ByteEncode12_Vec F_vec = join (map ByteEncode12 F_vec)

    /**
     * The subtract-and-divide algorithm applied to `a` in `ByteEncode` is the
     * same as shifting right.
     *
     * This is in reference to [FIPS-203] Section 4.2.1 Algorithm 5, Step 4-5.
     * That algorithm converts an integer repesentation `a` into a bitwise
     * representation `b`, and needs to iteratively update `a` to discard each
     * bit that it has already retrieved for `b`.
     *
     * Note: The type constraint for `d` does not allow `1` because `% 2` is
     * not a legal operation on 1-bit vectors (since 2 cannot be represented
     * as a 1-bit vector). It remains true, though (with a result of 0).
     *
     * These tests use the values of `d` that are actually used in the
     * algorithm. `d_u` and `d_v` are functor parameters for this module.
     * ```repl
     * :prove subAndDivIsShiftR`{d_u}
     * :prove subAndDivIsShiftR`{d_v}
     * :prove subAndDivIsShiftR`{12}
     * ```
     */
    subAndDivIsShiftR : {d} (fin d, d > 1) => [d] -> Bit
    property subAndDivIsShiftR a = take (sad a) == shift a where
        // Recursive definition of Step 4-5 (step 4 is inline here)
        sad x = [x] # sad ((x - (x % 2)) / 2)
        shift x = [x >> i | i <- [0..d-1]]

    /**
     * Computing a bit vector mod 2 is the same as taking its final (least
     * significant) bit.
     *
     * This is in reference to [FIPS-203] Section 4.2.1 Algorithm 5, Step 4.
     * That algorithm converts the integer representation of `a` into a bitwise
     * representation `b` and uses `mod` to iteratively get bits out of `a`.
     *
     * Note: The type constraint for `d` does not allow `1` because `% 2` is
     * not a legal operation on 1-bit vectors (since 2 cannot be represented
     * as a 1-bit vector.) It's trivially true, though.
     *
     * These tests use the values of `d` that are actually used in the
     * algorithm. `d_u` and `d_v` are functor parameters for this module.
     * ```repl
     * :prove mod2IsFinalBit`{d_u}
     * :prove mod2IsFinalBit`{d_v}
     * :prove mod2IsFinalBit`{12}
     * ```
     */
    mod2IsFinalBit : {d} (fin d, d > 1) => [d] -> Bit
    property mod2IsFinalBit a = a % 2 == zext [a ! 0]

    /**
     * Decode a byte array into an array of `d`-bit integers, for `d < 12`.
     * [FIPS-203] Section 4.2.1 Algorithm 6.
     *
     * Note: In this implementation, we treat the `d < 12` case separately from
     * `d == 12` because it allows us to better express the different integer
     * types. For `d < 12`, we use the bit vector type.
     * See `ByteDecode12` for the `d = 12` case.
     */
    ByteDecode : {d} (1 <= d, d < 12) => [32 * d]Byte -> [256][d]
    ByteDecode B = F where
        // Step 1.
        b = BytesToBits B
        // Steps 2-4.
        // - In FIPS 203, these operations are taken `mod m`, where `m = 2^d`.
        //   The `mod` is implicit here because computations on the `[d]` type
        //   are `mod 2^d` by definition.
        // - In a bit vector, multiplying by `2^j` is the same as left shift
        //   (see mul2jIsShiftL`).
        // - `bidj` is a single Bit; we use `zext` to convert to a `d`-length
        //   bit vector (with zeros in the higher-order bits) to support
        //   subsequentoperations.
        F = [ sum [ (zext [bidj]) << j
                | bidj <- bid
                | j <- [0..d-1]]
            | bid <- split`{256} b]

    /**
     * Decode `k` arrays of `d`-bit integers into a byte array.
     * [FIPS-203] Section 2.4.8.
     */
    ByteDecode_Vec : {d} (1 <= d, d < 12) => [k * 32 * d]Byte -> [k][256][d]
    ByteDecode_Vec B_vec = map ByteDecode (split B_vec)

    /**
     * Decode a byte array into an array of integers mod `q`.
     * [FIPS-203] Section 4.2.1 Algorithm 6.
     *
     * Note: In this implementation, we treat the `d < 12` case separately from
     * `d == 12` because it allows us to better express the different integer
     * types. For `d = 12`, we use the integers-mod (`Z`) type.
     * See `ByteDecode` for the `d < 12` case.
     */
    ByteDecode12 : [32 * 12]Byte -> [256](Z q)
    ByteDecode12 B = F' where
        type d = 12
        // These steps are the same as in `ByteDecode`. See that function for
        // notes.
        b = BytesToBits B
        F = [ sum [ (zext`{d} [bidj]) << j
                | bidj <- bid
                | j <- [0..d-1]]
            | bid <- split`{256} b]

        // In this case, since `m = q` (and not `2^12`), we need to explicitly
        // convert each integer from a 12-bit array to an integer mod `q`.
        toZq f = fromInteger (toInteger f)
        F' = map toZq F

    /**
     * Decode `k` joined byte arrays into `k` arrays of integers mod `q`.
     */
    ByteDecode12_Vec : [k * 32 * 12]Byte -> [k][256](Z q)
    ByteDecode12_Vec B = map ByteDecode12 (split B)

    /**
     * Multiplying a value by `2^j` is the same as bit-shifting it left by `j`
     * bits.
     * Note: The type constraint does not allow `d == 1` because the hard-coded
     * `2` does not fit in a single bit, but the property is trivial:
     * `j` can only be 0. On the left, `b * 2^0 == b * 1 == b`. On the right,
     * `b << 0 == b`.
     *
     * These tests use the values of `d` that are actually used in the
     * algorithm. `d_u` and `d_v` are functor parameters for this module.
     * ```repl
     * :prove mul2jIsShiftL`{d_u}
     * :prove mul2jIsShiftL`{d_v}
     * :prove mul2jIsShiftL`{12}
     * ```
     */
    mul2jIsShiftL : {d} (fin d, 1 < d) => [d] -> Bit
    property mul2jIsShiftL b = and [( b * (2^^j)) == (b << j) | j <- [0..d-1]]

    /**
     * When `d < 12`, the byte encoding and decoding functions should be
     * one-to-one inverses of each other.
     * [FIPS-203] Section 4.2.1 "Encoding and decoding", first point.
     *
     * These tests use the values of `d` that are actually used in the
     * algorithm. `d_u` and `d_v` are functor parameters for this module.
     * ```repl
     * :prove ByteEncodeInvertsByteDecode`{1}
     * :prove ByteEncodeInvertsByteDecode`{d_u}
     * :prove ByteEncodeInvertsByteDecode`{d_v}
     * ```
     */
    ByteEncodeInvertsByteDecode : {d} (fin d, 1 <= d, d < 12) => [256][d] -> Bit
    property ByteEncodeInvertsByteDecode bits =
        decode_encode_works && encode_decode_works where
        decode_encode_works = ByteDecode (ByteEncode bits) == bits
        // Rearrange random input to be valid for the decode function
        bytes = split (join bits)
        encode_decode_works = ByteEncode (ByteDecode bytes) == bytes

    /**
     * Byte decoding is the inverse of byte decoding for `d = 12`.
     * [FIPS-203] Section 4.2.1 "Encoding and decoding", second point.
     *
     * Note that the reverse property (decoding, then encoding) is not true!
     * ```repl
     * :check ByteDecode12InvertsByteEncode12
     * ```
     */
    ByteDecode12InvertsByteEncode12 : [256](Z q) -> Bit
    ByteDecode12InvertsByteEncode12 bits =
        ByteDecode12 (ByteEncode12 bits) == bits

    /**
     * Uniformly sample NTT representations.
     *
     * This uses a seed `B` to generate a pseudorandom stream, which is parsed into
     * a polynomial in `T_q` drawn from a distribution indistinguishable from the
     * uniform distribution.
     *
     * [FIPS-203] Section 4.2.2, Algorithm 7.
     */
    SampleNTT : [34]Byte -> Tq
    SampleNTT B = a_hat' where
        // Steps 1-2, 5.
        // We (lazily) take an infinite stream from the XOF and remove only as
        // many bytes as are needed to compute the function. See [FIPS-203]
        // Section 4.1 Equation 4.6 for a discussion of the equivalence of this
        // form to the one in Algorithm 7.
        ctx0 = XOF B

        // Step 3. Since Cryptol is not imperative, we implement this loop using
        // recursion. The `j` counter is not made explicit; instead we lazily
        // generate an infinite stream of coefficients in `T_q` and `take` the
        // correct length in the next line.

        // Step 4-16. `take` fulfills the `j < 256` condition in Steps 4 and 12.
        a_hat = take`{256} (filter ctx0)

        // `filter` parses an infinite stream from the XOF, computing
        // potential elements `d1` and `d2` from the first 3 bytes in the stream
        // and adding them to the output if they are valid elements in `Z q`.
        filter: [inf]Byte -> [inf][12]
        filter XOFSqueeze = a_hat_j where
            // Step 5.
            (C # ctx) = XOFSqueeze

            // The conversion from 8-bit to 12-bit vectors (with the same
            // value!) is implicit in the spec -- see notes on Step 5 and 6/7.
            // In Cryptol, we need to convert manually to do the subsequent
            // computations.
            [C0, C1, C2] = map zext`{12} C

            // Step 6.
            d1 = C0 + 256 * (C1 % 16)

            // Step 7. Cryptol uses integer division; it always takes the floor
            // of the result.
            d2 = (C1 / 16) + 16 * C2

            // Steps 4, 8 - 15.
            // Add `d1` and/or `d2` to the sampled vector `a_hat` if they are
            // valid elements in `Z q`.
            // The `while` loop in Step 4 is equivalent to the recursive call to
            // `filter` in each condition.
            a_hat_j = if (d1 < `q) && (d2 < `q) then
                    [d1, d2] # filter ctx
                else if d1 < `q then
                    [d1] # filter ctx
                else if d2 < `q then
                    [d2] # filter ctx
                else filter ctx

        // This conversion is implicit in the implementation -- see the notes on
        // Step 6/7 and 9.
        toZq : [12] -> Z q
        toZq x = fromInteger (toInteger x)

        a_hat' = map toZq a_hat

    /**
     * Sample a special, centered distribution of polynomials in `R_q` with small
     * coefficients.
     *
     * The input stream `B` must be uniformly random bytes!
     *
     * [FIPS-203] Section 4.2.2, Algorithm 8.
     */
    SamplePolyCBD: {eta} (2 <= eta, eta <= 3) => [64 * eta]Byte -> Rq
    SamplePolyCBD B = f where
        // Step 1.
        b = BytesToBits B
        // This conversion is implicit in the implementation. Convert each bit into
        // an element of Z q.
        BitToZ : Bit -> Z q
        BitToZ bit = if bit then 1 else 0
        b' = map BitToZ b

        // Step 3.
        x i = sum [b'@(2 * i * `eta + j) | j <- [0 .. (eta-1)]]
        // Step 4.
        y i = sum [b'@(2 * i * `eta + `eta + j) | j <- [0 .. (eta-1)]]

        // Steps 2, 5. The `mod q` is not explicit here because `x` and `y`
        // return elements of `Z q`.
        f = [(x i) - (y i) | i <- [0 .. 255]]

    /**
     * [FIPS-203] Section 4.3 "The mathematical structure of the NTT."
     * ```repl
     * :prove QisCorrectlyDefined
     * ```
     */
    QisCorrectlyDefined: Bit
    property QisCorrectlyDefined = `q == 2^^8 * 13 + 1

    /**
     * `zeta` is a primitive 256-th root of unity modulo `q`.
     * [FIPS-203] Section 4.3 "The mathematical structure of the NTT."
     *
     * ```repl
     * :prove zetaIsPrimitiveRoot
     * ```
     */
    property zetaIsPrimitiveRoot = zeta ^^ 128 == -1

    /**
     * Proves that `zeta` is correctly set to be the 256th root of `q`.
     * ```repl
     * :exhaust Is256thRootOfq
     * ```
     */
    Is256thRootOfq : [lg2 q] -> Bit
    property Is256thRootOfq p = (p == 0) || (p >= 256) || (zeta^^p != 1)

    /**
     * Reverse the unsigned 7-bit value corresponding to an input integer in
     * `[0, ..., 127]`.
     * [FIPS-203] Section 4.3 "The mathematical structure of the NTT."
     *
     * This diverges from the spec by operating over an 8-bit vector;
     * this is to ease prior and subsequent computations that would overflow a
     * 7-bit vector, like:
     * - `2 * (BitRev7 i) + 1`
     * - `2 * i + 1`
     *
     * A "pure" implementation of `BitRev7` in Cryptol is the `reverse` function
     * on 7-bit vectors. This mini-property shows equivalence:
     * ```repl
     * :prove \(x:[7]) -> ([0] # reverse x) == BitRev7 ([0] # x)
     * ```
     */
    BitRev7 : [8] -> [8]
    BitRev7 i = if i > 255 then error "BitRev7 called with invalid input"
        else (reverse i) >> 1

/**
 * This section specifies the number-theoretic transform (NTT).
 *
 * It includes the version from [FIPS-203] Section 4.3 as well
 * as a faster O(N log N) version, and a proof of their equivalence.
 *
 * This is explicitly allowed by the spec: "For every computational procedure
 * [...] a conforming implementation may replace the given set of steps with
 * any mathematically equivalent set of steps". The equivalence properties
 * prove mathematical equivalence.
 * [FIPS-203] Introduction, "7. Implementations".
 */
import submodule NTT
submodule NTT where
    private
        /**
         * Number theoretic transform: compute the "NTT representation" in
         * `T_q` of a polynomial in `R_q`.
         *
         * [FIPS-203] Section 4.3, Algorithm 9.
         */
        NaiveNTT : Rq -> Tq
        NaiveNTT f = join [[f2i i, f2iPlus1 i] | i <- [0 .. 127]] where
            f2i i = sum [f @(2*j) * zeta_term i j | j <- [0 .. 127]]
            f2iPlus1 i = sum [f @(2*j+1) * zeta_term i j | j <- [0 .. 127]]
            zeta_term i j = zeta ^^ ((2 * BitRev7 i + 1) * j)

        /**
         * Inverse of the number theoretic transform: converts from the "NTT
         * representation" in `T_q` to a polynomial in `R_q`.
         *
         * [FIPS-203] Section 4.3, Algorithm 10.
         */
        NaiveNTTInv : Tq -> Rq
        NaiveNTTInv f_hat = [f_i * 3303 | f_i <- f] where
            f = join [[f2i i, f2iPlus1 i] | i <- [0 .. 127]]
            f2i i = sum [f_hat @(2*j) * zeta_term i j | j <- [0 .. 127]]
            f2iPlus1 i = sum [f_hat @(2*j+1) * zeta_term i j | j <- [0 .. 127]]
            zeta_term i j = (recip zeta) ^^ ((2 * BitRev7 j + 1) * i)

        //////////////////////////////////////////////////////////////
        // This section specifies fast O(N log N) NTT and Inverse NTT
        //
        // A readable explanation of the derivation of this form of
        // the NTT is in "A Complete Beginner Guide to the Number
        // Theoretic Transform (NTT)" by Ardianto Satriawan,
        // Rella Mareta, and Hanho Lee. Available from:
        //    https://eprint.iacr.org/2024/585
        //
        // This section Copyright Amazon.com, Inc. or its affiliates.
        //////////////////////////////////////////////////////////////

        /**
         * Lookup table for `zeta` values computed in `NTT` and `NTTInv`.
         *
         * [FIPS-203] Appendix A.
         */
        zeta_expc  : [128](Z q)
        zeta_expc = [
            1, 1729, 2580, 3289, 2642, 630, 1897, 848,
            1062, 1919, 193, 797, 2786, 3260, 569, 1746,
            296, 2447, 1339, 1476, 3046, 56, 2240, 1333,
            1426, 2094, 535, 2882, 2393, 2879, 1974, 821,
            289, 331, 3253, 1756, 1197, 2304, 2277, 2055,
            650, 1977, 2513, 632, 2865, 33, 1320, 1915,
            2319, 1435, 807, 452, 1438, 2868, 1534, 2402,
            2647, 2617, 1481, 648, 2474, 3110, 1227, 910,
            17, 2761, 583, 2649, 1637, 723, 2288, 1100,
            1409, 2662, 3281, 233, 756, 2156, 3015, 3050,
            1703, 1651, 2789, 1789, 1847, 952, 1461, 2687,
            939, 2308, 2437, 2388, 733, 2337, 268, 641,
            1584, 2298, 2037, 3220, 375, 2549, 2090, 1645,
            1063, 319, 2773, 757, 2099, 561, 2466, 2594,
            2804, 1092, 403, 1026, 1143, 2150, 2775, 886,
            1722, 1212, 1874, 1029, 2110, 2935, 885, 2154
        ]


        // Top level entry point - start with lv=256, k=1
        fast_ntt : Rq -> Tq
        fast_ntt v = fast_nttl v 1

        // Recursive NTT function
        fast_nttl :
            {lv}  // Length of v is a member of {256,128,64,32,16,8,4}
            (lv >= 2, lv <= 8) =>
            [2^^lv](Z q) -> [8] -> [2^^lv](Z q)
        fast_nttl v k
            // Base case. lv==2 so just compute the butterfly and return
            | lv == 2 => ct_butterfly`{lv,lv-1} v (zeta_expc@k)

            // Recursive case. Butterfly what we have, then recurse on each half,
            // concatenate the results and return.
            | lv  > 2 => (fast_nttl`{lv-1} s0 (k * 2)) #
                    (fast_nttl`{lv-1} s1 (k * 2 + 1))
                        where
                        t = ct_butterfly`{lv,lv-1} v (zeta_expc@k)
                        // Split t into two halves s0 and s1
                        [s0, s1] = split t

        // Fast recursive CT-NTT
        ct_butterfly :
            {m, hm}
            (m >= 2, m <= 8, hm >= 1, hm <= 7, hm == m - 1) =>
            [2^^m](Z q) -> (Z q) -> [2^^m](Z q)
        ct_butterfly v z = new_v where
            halflen = 2^^`hm
            lower, upper : [2^^hm](Z q)
            lower@x = v@x + z * v@(x + halflen)
            upper@x = v@x - z * v@(x + halflen)
            new_v = lower # upper


        // Recursive inverse-NTT function
        fast_invnttl :
            {lv}  // Length of v is a member of {256,128,64,32,16,8,4}
            (lv >= 2, lv <= 8) =>
            [2^^lv](Z q) -> [8] -> [2^^lv](Z q)
        fast_invnttl v k
            // Base case. lv==2 so just compute the butterfly and return
            | lv == 2 => gs_butterfly`{lv,lv-1} v (zeta_expc@k)

            // Recursive case. Recurse on each half,
            // concatenate the results, butterfly that, and return.
            | lv  > 2 => gs_butterfly`{lv,lv-1} t (zeta_expc@k) where
                // Split t into two halves s0 and s1
                [s0, s1] = split v
                t = (fast_invnttl`{lv-1} s0 (k * 2 + 1)) #
                    (fast_invnttl`{lv-1} s1 (k * 2))

        // Fast recursive GS-Inverse-NTT
        gs_butterfly :
            {m, hm}
            (m >= 2, m <= 8, hm >= 1, hm <= 7, hm == m - 1) =>
            [2^^m](Z q) -> (Z q) -> [2^^m](Z q)
        gs_butterfly v z = new_v where
            halflen = 2^^`hm
            lower, upper : [2^^hm](Z q)
            lower@x = v@x  + v@(x + halflen)
            upper@x = z * (v@(x + halflen) - v@x)
            new_v = lower # upper

        // Top level entry point - start with lv=256, k=1
        fast_invntt : Tq -> Rq
        fast_invntt v = mul_recip128 (fast_invnttl v 1) where

            // Multiplicative inverse of 128, mod `q`.
            recip_128_modq = (recip 128) : (Z q)

            // Multiply all elements of v' by the reciprocal of 128 (modulo q)
            mul_recip128 : Tq -> Tq
            mul_recip128 v' = [ v'@x * recip_128_modq | x <- [0 .. <n] ]

        //////////////////////////////////////////////////////////////
        // Properties and proofs of Naive and Fast NTT
        //////////////////////////////////////////////////////////////

        /**
         * This property demonstrates that NaiveNTT is self-inverting.
         * ```repl
         * :prove NaiveNTT_Inverts
         * ```
         */
        NaiveNTT_Inverts : Rq -> Bit
        property NaiveNTT_Inverts f =  NaiveNTTInv (NaiveNTT f) == f

        /**
         * This property demonstrates that NaiveNTTInv is self-inverting.
         * ```repl
         * :prove NaiveNTTInv_Inverts
         * ```
         */
        NaiveNTTInv_Inverts : Tq -> Bit
        property NaiveNTTInv_Inverts f =  NaiveNTT (NaiveNTTInv f) == f

        /**
         * This property demonstrates that `fast_ntt` is the inverse of `fast_invntt`.
         * ```repl
         * :prove fast_ntt_inverts
         * ```
         */
        fast_ntt_inverts : Rq -> Bit
        property fast_ntt_inverts    f =  fast_invntt (fast_ntt f)    == f

        /**
         * This property demonstrates that `fast_invntt` is the inverse of `fast_ntt`.
         * ```repl
         * :prove fast_invntt_inverts
         * ```
         */
        fast_invntt_inverts : Tq -> Bit
        property fast_invntt_inverts f =  fast_ntt    (fast_invntt f) == f

        /**
         * This property demonstrates that `naive_ntt` is equivalent to `fast_ntt`.
         * ```repl
         * :prove naive_fast_ntt_equiv
         * ```
         */
        naive_fast_ntt_equiv : Rq -> Bit
        property naive_fast_ntt_equiv f =  NaiveNTT f == fast_ntt f

        /**
         * This property demonstrates that `naive_invntt` is equivalent to `fast_invntt`.
         * ```repl
         * :prove naive_fast_invntt_equiv
         * ```
         */
        naive_fast_invntt_equiv : Tq -> Bit
        property naive_fast_invntt_equiv f =  NaiveNTTInv f == fast_invntt f

    /**
     * The Number-Theoretic Transform (NTT) is used to improve the efficiency
     * of multiplication in the ring `R_q`. We choose to use the fast version
     * of NTT, which is equivalent to the version described in the spec.
     */
    NTT : Rq -> Tq
    NTT = fast_ntt

    /**
     * The inverse of the Number-Theoretic Transform (NTT) is used to improve
     * the efficiency of multiplication in the ring `R_q`. We choose to use
     * the fast version of inverse function, which is equivalent to the version
     * described in the spec.
     */
    NTTInv : Tq -> Rq
    NTTInv = fast_invntt

    /**
     * The notation `NTT` is overloaded to mean both a single application of `NTT`
     * to an element of `R_q` and also `k` applications of `NTT` to every element
     * of a `k`-length vector.
     * [FIPS-203] Section 2.4.6 Equation 2.9.
     */
    NTT_Vec v = map NTT v

    /**
     * The notation `NTTInv` is overloaded to mean both a single application of
     * `NTTInv` to an element of `R_q` and also `k` applications of `NTTInv` to
     * every element of a `k`-length vector.
     * [FIPS-203] Section 2.4.6.
     */
    NTTInv_Vec v = map NTTInv v


    /**
     * Compute the product of two degree-one polynomials with respect to a
     * quadratic modulus.
     * [FIPS-203] Section 4.3.1 Algorithm 12.
     */
    BaseCaseMultiply : (Z q) -> (Z q) -> (Z q) -> (Z q) -> (Z q) -> [2](Z q)
    BaseCaseMultiply a0 a1 b0 b1 γ = [c0, c1]
      where
        c0 = a0 * b0 + a1 * b1 * γ
        c1 = a0 * b1 + a1 * b0

    /**
     * Compute the product (in the ring `T_q`) of two NTT representations.
     * [FIPS-203] Section 4.3.1 Algorithm 11.
     */
    MultiplyNTTs : Tq -> Tq -> Tq
    MultiplyNTTs f_hat g_hat = join h_hat where
        h_hat = [ BaseCaseMultiply
            (f_hat @(2*i))
            (f_hat @(2*i+1))
            (g_hat @(2*i))
            (g_hat @(2*i+1))
            (zeta ^^(2 * BitRev7 i + 1))
        | i <- [0 .. 127] ]

    /**
     * Testing that (1+x)^2 = 1+2x+x^2.
     * ```repl
     * :prove TestMult
     * ```
     */
    TestMult : Bit
    property TestMult = prod f f == fsq where
      f = [1, 1] # [0 | i <- [3 .. 256]]
      fsq = [1,2,1] # [0 | i <- [4 .. 256]]

      prod : Rq -> Rq -> Rq
      prod a b = NTTInv (MultiplyNTTs (NTT a) (NTT b))

    /**
     * The cross product notation ×𝑇𝑞 is defined as the `MultiplyNTTs` function
     * (also referred to as `T_q` multiplication).
     * [FIPS-203] Section 2.4.5 Equation 2.8.
     */
    (×) = MultiplyNTTs

    /**
     * Overloaded `dot` function between two vectors is a standard dot-product
     * functionality with `T_q` multiplication as the base operation.
     * [FIPS-203] Section 2.4.7 Equation 2.14.
     */
    dotVecVec : {k1} (fin k1) => [k1]Tq -> [k1]Tq -> Tq
    dotVecVec v1 v2 = sum (zipWith (×) v1 v2)

    /**
     * Overloaded `dot` function between a matrix and a vector is standard
     * matrix-vector multiplication with `T_q` multiplication as the base
     * operation.
     * [FIPS-203] Section 2.4.7 Equation 2.12 and 2.13.
     */
    dotMatVec : {k1,k2} (fin k1, fin k2) => [k1][k2]Tq -> [k2]Tq -> [k1]Tq
    dotMatVec matrix vector = [dotVecVec v1 vector | v1 <- matrix]

    /**
     * Overloaded `dot` function between two matrices is standard matrix
     * multiplication with `T_q` multiplication as the base operation.
     * [FIPS-203] Section 2.4.7.
     */
    dotMatMat :{k1,k2,k3} (fin k1, fin k2, fin k3) =>
        [k1][k2]Tq -> [k2][k3]Tq -> [k1][k3]Tq
    dotMatMat matrix1 matrix2 = transpose [dotMatVec matrix1 vector | vector <- m']
        where m' = transpose matrix2

/**
 * The K-PKE component scheme.
 *
 * ⚠️ This scheme is not approved for stand-alone use! ⚠️
 * K-PKE is an encryption scheme consisting of three algorithms `(KeyGen,
 * Encrypt, Decrypt)`, which are used to instantiate the approved ML-KEM
 * scheme. It's not secure as a standalone scheme; it doesn't do any input
 * checking.
 * [FIPS-203] Section 5.
 */
private submodule K_PKE where
    // Encryption key for the K_PKE component scheme.
    // [FIPS-203] Section 5 Algorithm 13. See "Output".
    type EncryptionKey = [384 * k + 32]Byte

    // Decryption key for the K_PKE component scheme.
    // [FIPS-203] Section 5 Algorithm 13. See "Output".
    type DecryptionKey = [384 * k]Byte

    // Ciphertext generated by the K_PKE component scheme.
    // [FIPS-203] Section 5 Algorithm 14. See "Output".
    type Ciphertext = [32 * (d_u * k + d_v)]Byte

    /**
     * Key generation for the K-PKE component scheme.
     *
     * ⚠️ Warnings ⚠️
     * - This scheme is not approved for use in a stand-alone fashion! It does not
     *   do any input validation and should only be used as a subroutine of ML-KEM.
     * - The seed `d` passed as input and the decryption key `dkPKE` returned from
     *   this algorithm must be kept private!
     *
     * [FIPS-203] Section 5.1 Algorithm 13.
     */
    KeyGen: [32]Byte -> (EncryptionKey, DecryptionKey)
    KeyGen d = (ekPKE, dkPKE) where
        // Step 1.
        (ρ, σ) = G (d # [`(k)])
        // Steps 3-7.
        A_hat = [[ SampleNTT (ρ # [j] # [i])
            | j <- [0 .. k-1]]
            | i <- [0 .. k-1]]
        // Steps 2, 8-11.
        s = [SamplePolyCBD`{eta_1} (PRF σ N)
            | N <- [0 .. k-1]]
        // Steps 12 - 15.
        e = [SamplePolyCBD`{eta_1} (PRF σ N)
            | N <- [k .. 2 * k - 1]]
        // Step 16.
        s_hat = NTT_Vec s
        // Step 17.
        e_hat = NTT_Vec e
        // Step 18.
        t_hat = (dotMatVec A_hat s_hat) + e_hat
        // Step 19.
        ekPKE = (ByteEncode12_Vec t_hat) # ρ
        // Step 20.
        dkPKE = ByteEncode12_Vec s_hat

    /**
     * Encryption algorithm for the K-PKE component scheme.
     *
     * ⚠️ Warning ⚠️ This scheme is not approved for use in a stand-alone fashion!
     * It does not do any input validation and should only be used as a subroutine
     * of ML-KEM.
     *
     * [FIPS-203] Section 5.2 Algorithm 14.
     */
    Encrypt : EncryptionKey -> [32]Byte -> [32]Byte -> Ciphertext
    Encrypt ekPKE m r = c where
        // Step 2.
        t_hat = ByteDecode12_Vec (ekPKE @@[0 .. 384*k - 1])
        // Step 3.
        rho = ekPKE @@[384*k .. 384*k + 32 - 1]
        // Steps 4-8.
        A_hat = [[ SampleNTT (rho # [j] # [i])
            | j <- [0 .. k-1]]
            | i <- [0 .. k-1]]
        // Steps 1, 9-12.
        y = [SamplePolyCBD`{eta_1} (PRF r N)
            | N <- [0 .. k-1]]
        // Steps 13-16.
        e1 = [SamplePolyCBD`{eta_2} (PRF r N)
            | N <- [k .. 2 * k - 1]]
        // Step 17. In the spec, the second parameter is `N = 2k`. In this
        // implementation, `N` itself is out of scope, so we use the fixed
        // value instead.
        e2 = SamplePolyCBD`{eta_2} (PRF r (2 * `k))
        // Step 18.
        y_hat = NTT_Vec y
        // Step 19.
        u = NTTInv_Vec (dotMatVec (transpose A_hat) y_hat) + e1
        // Step 20.
        mu = Decompress_Vec`{1} (ByteDecode`{1} m)
        // Step 21.
        v = (NTTInv (dotVecVec t_hat y_hat)) + e2 + mu
        // Step 22.
        c1 = ByteEncode_Vec`{d_u} (Compress_Mat`{d_u} u)
        // Step 23.
        c2 = ByteEncode`{d_v} (Compress_Vec`{d_v} v)
        // Step 24.
        c = c1 # c2

    /**
     * Decryption algorithm for the K-PKE component scheme.
     *
     * ⚠️ Warning ⚠️ This scheme is not approved for use in a stand-alone fashion!
     * It does not do any input validation and should only be used as a subroutine
     * of ML-KEM.
     *
     * [FIPS-203] Section 5.3 Algorithm 15.
     */
    Decrypt : DecryptionKey -> Ciphertext -> [32]Byte
    Decrypt dkPKE c = m where
        // Step 1.
        c1 = c @@[0 .. 32 * d_u * k - 1]
        // Step 2.
        c2 = c @@[32 * d_u * k .. 32 * (d_u * k + d_v) - 1]
        // Step 3.j
        u' = Decompress_Mat`{d_u} (ByteDecode_Vec`{d_u} c1)
        // Step 4.
        v' = Decompress_Vec`{d_v} (ByteDecode`{d_v} c2)
        // Step 5.
        s_hat = ByteDecode12_Vec dkPKE
        // Step 6.
        w = v' - NTTInv (dotVecVec s_hat (NTT_Vec u'))
        // Step 7.
        m = ByteEncode`{1} (Compress_Vec`{1} w)

    /**
     * The K-PKE scheme must satisfy the basic properties of an encryption
     * scheme.
     * This must be `:check`ed because K-PKE is correct with probability
     * 1-delta and not 1. It is not provably correct because there is a
     * (very small!) fraction of seeds `d, r` that don't work.
     * ```repl
     * :set tests=3
     * :check CorrectnessPKE
     * ```
     */
    CorrectnessPKE : [32]Byte -> [32]Byte -> [32]Byte -> Bit
    property CorrectnessPKE d m r = (m' == m) where
        (pk, sk) = KeyGen d
        c = Encrypt pk m r
        m' = Decrypt sk c


private
    /**
     * Uses randomness to generate an encapsulation key and corresponding
     * decapsulation key.
     * [FIPS-203] Section 6.1 Algorithm 16.
     *
     * Warning: This function does not validate input and should not be used
     * externally!
     */
    KeyGen_internal : [32]Byte -> [32]Byte
        -> (EncapsulationKey, DecapsulationKey)
    KeyGen_internal d z = (ek, dk) where
      (ekPKE, dkPKE) = K_PKE::KeyGen d
      ek = ekPKE
      dk = dkPKE # ek # (H ek) # z

    /**
     * Uses the encapsulation key and randomness to generate a key and an
     * associated ciphertext.
     * [FIPS-203] Section 6.2 Algorithm 17.
     *
     * Warning: This function does not validate input and should not be used
     * externally!
     */
    Encaps_internal : EncapsulationKey -> [32]Byte
        -> (SharedSecretKey, Ciphertext)
    Encaps_internal ek m = (K, c) where
      (K, r) = G (m # (H ek))
      c = K_PKE::Encrypt ek m r

    /**
     * Uses the decapsulation key to produce a shared secret key from a
     * ciphertext.
     * [FIPS-203] Section 6.3 Algorithm 18.
     *
     * Warning: This function does not validate input and should not be used
     * externally!
     */
    Decaps_internal : DecapsulationKey -> Ciphertext -> SharedSecretKey
    Decaps_internal dk c = K where
        // Step 1. Extract (from KEM decaps key) the PKE decryption key.
        dkPKE = dk @@ [0 .. 384*k - 1]
        // Step 2. Extract PKE encryption key.
        ekPKE = dk @@ [384*k .. 768*k + 32 - 1]
        // Step 3. Extract hash of PKE encryption key.
        h = dk @@ [768*k + 32 .. 768*k + 64 - 1]
        // Step 4. Extract implicit rejection value.
        z = dk @@ [768*k + 64 .. 768*k + 96 - 1]
        // Step 5. Decrypt ciphertext.
        m' = K_PKE::Decrypt dkPKE c
        // Step 6.
        (K', r') = G (m' # h)
        // Step 7.
        Kbar = J (z # c)
        // Step 8. Re-encrypt using the derived randomness `r'`.
        c' = K_PKE::Encrypt ekPKE m' r'
        // Step 9 - 11. If ciphertextx do not match, "implicitly reject".
        K = if (c != c') then
                Kbar
            else
                K'

    /**
     * The ML-KEM scheme is correct with probability 1-delta and not 1. As a
     * result, running `:prove CorrectnessPKE` will not succeed since there is a
     * fraction delta of seeds `d`, `z`, `m` that do not work.
     * Cryptol does not currently support counting.
     * ```repl
     * :set tests=3
     * :check CorrectnessKEM
     * ```
     */
    CorrectnessKEM : [32]Byte -> [32]Byte -> [32]Byte -> Bit
    property CorrectnessKEM z d m = (K == K') where
        (ek, dk) = KeyGen_internal d z
        (K, c) = Encaps_internal ek m
        K' = Decaps_internal dk c

/*
 * Type of an encapsulation key used in the public ML-KEM API.
 * [FIPS-203], as in the output type in Algorithm 19.
 * See also the paragraph "Terminology for keys" in Section 3.2.
 */
type EncapsulationKey = [384 * k + 32]Byte

/*
 * Type of an decapsulation key used in the public ML-KEM API.
 * [FIPS-203], as in the output type in Algorithm 19.
 * See also the paragraph "Terminology for keys" in Section 3.2.
 */
type DecapsulationKey = [768 * k + 96]Byte

/*
 * Type of the shared secret key material generated by the public ML-KEM API.
 * [FIPS-203], as in the output type in Algorithm 20.
 * See also the paragraph "Terminology for keys" in Section 3.2.
 */
type SharedSecretKey = [32]Byte

/*
 * Type of the ciphertext generated by teh public ML-KEM API.
 * [FIPS-203], as in the output type in Algorithm 20.
 *
 * Usage: This ciphertext is only used in the context of ML-KEM, to derive a
 * shared secret key.
 */
type Ciphertext = [32 * (d_u * k + d_v)]Byte

/**
 * Generate an encapsulation key and a corresponding decapsulation key.
 * [FIPS-203] Section 7.1, Algorithm 19.
 *
 * This differs from the spec: it takes the randomness `d` and `z` as
 * parameters instead of generating them internally (because Cryptol
 * cannot generate randomness).
 *
 * The decapsulation key produced from this function MUST remain private. The
 * seed `(d, z)` produced (or in this case, passed as input) is sensitive data
 * and MUST be treated with the same safeguards as the decapsulation key.
 */
KeyGen : Option ([32]Byte) -> Option ([32]Byte)
    -> (EncapsulationKey, DecapsulationKey)
KeyGen d z = (ek, dk) where
    // Steps 1 and 2 are skipped because Cryptol cannot generate randomness.
    // Step 3 - 5.
    d' = case d of
        None -> error "Random bit generation failed; `d` not initialized."
        Some _d -> _d
    z' = case z of
        None -> error "Random bit generation failed; `z` not initialized."
        Some _z -> _z
    // Step 6.
    (ek, dk) = KeyGen_internal d' z'

/**
 * Perform checks on a key pair.
 *
 * The security of ML-KEM depends on the use of key pairs that were correctly
 * generated using `KeyGen`!
 * This checking process _does not_ guarantee that the input pair is a properly
 * produced output of `KeyGen`. While these checks can detect certain
 * corruptions, they do not guarantee that the key pair was properly generated.
 */
keyPairCheck : EncapsulationKey
    -> DecapsulationKey
    -> Option ([32]Byte, [32]Byte)
    -> [32]Byte
    -> Bit
keyPairCheck ek dk maybe_seed m =
    seedConsistency && encapsCheck && decapsCheck && pairwiseConsistency where
    seedConsistency = case maybe_seed of
        Some (d, z) -> KeyGen_internal d z == (ek, dk)
        None -> True

    encapsCheck = encapsulationKeyCheck ek
    decapsCheck = decapsulationKeyCheck dk

    pairwiseConsistency = K == K'
    (K, c) = Encaps_internal ek m
    K' = Decaps_internal dk c

/**
 * A properly produced output of the `KeyGen` function should always pass the
 * key pair checks.
 * ```repl
 * :set tests=3
 * :check KeyGenProducesValidKeys
 * ```
 */
KeyGenProducesValidKeys d z m = keyPairCheck ek dk (Some (d, z)) m where
    (ek, dk) = KeyGen (Some d) (Some z)

/**
 * Perform checks on a candidate encapsulation key.
 *
 * The type check is implicit; if this function compiles with a given `ek`,
 * then it must be of the correct type.
 *
 * This checking process _does not_ gurantee that `ek` is a properly produced
 * output of `KeyGen`.
 */
encapsulationKeyCheck : EncapsulationKey -> Bool
encapsulationKeyCheck ek = modulusCheck where
    test = ByteEncode12_Vec (ByteDecode12_Vec (ek @@ [0..384 * k - 1]))
    modulusCheck = test == (ek @@ [0..384 * k - 1])

/**
 * Use the encapsulation key to generate a shared secret an an associated
 * ciphertext.
 * [FIPS-203] Section 7.2, Algorithm 20.
 *
 * This differs from the spec: it takes the randomness `m` as a parameter
 * instead of generating it internally (because Cryptol cannot generate
 * randomness).
 *
 * ⚠️ Warning ⚠️
 * The encapsulation key passed as a parameter MUST be checked before use with
 * `encapsulationKeyCheck`! The caller must have assurance that the
 * encapsulation key is valid.
 * This executable spec (and Cryptol in general) cannot enforce this
 * requirement! Implementors must check it manually.
 */
Encaps : EncapsulationKey -> Option ([32]Byte) -> (SharedSecretKey, Ciphertext)
Encaps ek m = (K, c) where
    // Step 1 is skipped because Cryptol cannot generate randomness.
    // Step 2 - 4.
    m' = case m of
        None -> error "Random bit generation failed; `m` not initialized."
        Some _m -> _m
    // Step 5.
    (K, c) = Encaps_internal ek m'

/**
 * Perform three checks on candidate inputs to decapsulation.
 * [FIPS-203] Section 7.3 "Decapsulation input check" #1-3.
 *
 * The type checks are implicit; if this function compiles with a given `dk`
 * and `c`, then they must be of the correct type.
 *
 * This check _does not_ guarantee that `dk` is a propertly produced output
 * of `KeyGen`, or that `c` is a properly produced output of `Encaps`!
 */
decapsulationInputCheck : DecapsulationKey -> Ciphertext -> Bool
decapsulationInputCheck dk c = decapsulationKeyCheck dk

private
    /**
     * Perform a type check and a hash check on a candidate decapsulation key.
     * [FIPS-203] Section 7.3 "Decapsulation input check" #2-3.
     *
     * The type check is implicit; if this function compiles with a given `dk`,
     * then it must be the correct type.
     *
     * This check _does not_ guarantee that `dk` is a propertly produced output
     * of `KeyGen`!
     */
    decapsulationKeyCheck : DecapsulationKey -> Bool
    decapsulationKeyCheck dk = hash_check where
        test = H (dk @@ [384*k .. 768*k + 32 - 1])
        hash_check = test == (dk @@ [768*k + 32 .. 768*k + 64 - 1])


/**
 * Use the decapsulation key to produce a shared secret from a ciphertext.
 * [FIPS-203] Section 7.3, Algorithm 21.
 *
 * ⚠️ Warning ⚠️
 * The decapsulation key and the ciphertext MUST be checked before use with
 * `decapsulationInputCheck`! The caller must have assurance that the
 * decapsulation key and ciphertext are valid.
 * This executable spec (and Cryptol in general) cannot enforce this
 * requirement! Implementors must check it manually.
 */
Decaps : DecapsulationKey -> Ciphertext -> SharedSecretKey
Decaps = Decaps_internal

parameter
    /*
     * The parameter `k` determines the dimensions of the encryption key matrix
     * and the secret and noise vectors created and used in key generation and
     * encryption.
     * [FIPS-203] Section 8.
     *
     * The constraint on the width of `k` is drawn from `K-PKE-KeyGen`. In
     * that function, the variable `N` varies from 0 to `2k` and is passed as
     * the second parameter to the `PRF` function. `PRF` restricts the second
     * parameter to be exactly 1 byte. Therefore, `2k` must fit into a byte.
     * [FIPS-203] Section 5.1 Algorithm 13 (see lines 2 and 8-15) and Section
     * 4.1 Equation 4.2.
     */
    type k : #
    type constraint (width k > 0, width (2*k) <= 8)

    /*
     * The parameter `eta_1` specifies the distribution from which the secret
     * vectors are drawn in key generation and encryption.
     * [FIPS-203] Section 8.
     *
     * eta_1 must be in the set {2, 3} for use in a PRF and to parameterize
     * sampling from the centered binomial distribution.
     * [FIPS-203] Section 4.1 "Pseudorandom Function" and Section 4.2.2
     * "Sampling from the centered binomial distribution"
     */
    type eta_1 : #
    type constraint (fin eta_1, 2 <= eta_1, eta_1 <= 3)

    /*
     * The parameter eta_2 is required to specify the distribution from which
     * the noise vectors are drawn in encryption.
     * [FIPS-203] Section 8.
     *
     * eta_2 must be in the set {2, 3} for use in a PRF and to parameterize
     * sampling from the centered binomial distribution.
     * [FIPS-203] Section 4.1 "Pseudorandom Function" and Section 4.2.2
     * "Sampling from the centered binomial distribution"
     */
    type eta_2 : #
    type constraint (fin eta_2, 2 <= eta_2, eta_2 <= 3)

    /*
     * The parameter `d_u` is a parameter and input for the compression and
     * encoding functions.
     * [FIPS-203] Section 8.
     *
     * For compression, `d_u` must be smaller than the bit length of `q` (fixed
     * to 12).
     * [FIPS-203] Section 4.2.1 "Compression and decompression".
     * For encoding, the valid range of values for `d_u` is `1 ≤ d_u ≤ 12`.
     * [FIPS-203] Section 4.2.1 "Encoding and decoding".
     */
    type d_u : #
    type constraint (fin d_u, d_u < 12, d_u > 0)

    /*
     * The parameter `d_v` is a parameter and input for the compression and
     * encoding functions.
     * [FIPS-203] Section 8.
     *
     * For compression, `d_v` must be smaller than the bit length of `q` (fixed
     * to 12).
     * [FIPS-203] Section 4.2.1 "Compression and decompression".
     * For encoding, the valid range of values for `d_v` is `1 ≤ d_v ≤ 12`.
     * [FIPS-203] Section 4.2.1 "Encoding and decoding".
     */
    type d_v : #
    type constraint (fin d_v, d_v < 12, d_v > 0)