ML_KEM.cs

using System;
using Waher.Security.SHA3;
 
namespace Waher.Security.PQC
{
    public class ML_KEM : ML_Common
    {
        public static readonly ML_KEM ML_KEM_512 = new ML_KEM(2, 3, 2, 10, 4);
 
        public static readonly ML_KEM ML_KEM_768 = new ML_KEM(3, 2, 2, 10, 4);
 
        public static readonly ML_KEM ML_KEM_1024 = new ML_KEM(4, 2, 2, 11, 5);
 
        public static ML_KEM GetModel(string Name)
        {
            switch (Name.ToUpper())
            {
                case "ML-KEM-512":
                    return ML_KEM_512;
 
                case "ML-KEM-768":
                    return ML_KEM_768;
 
                case "ML-KEM-1024":
                    return ML_KEM_1024;
 
                default:
                    throw new ArgumentException("Unknown model name: " + Name, nameof(Name));
            }
        }
 
        private static readonly ushort[] nttTransformZeta =
        {
            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
        };
 
        private static readonly ushort[] nttTransformZeta2 =
        {
            17, q - 17, 2761, q - 2761, 583, q - 583, 2649, q - 2649,
            1637, q - 1637, 723, q - 723, 2288, q - 2288, 1100, q - 1100,
            1409, q - 1409, 2662, q - 2662, 3281, q - 3281, 233, q - 233,
            756, q - 756, 2156, q - 2156, 3015, q - 3015, 3050, q - 3050,
            1703, q - 1703, 1651, q - 1651, 2789, q - 2789, 1789, q - 1789,
            1847, q - 1847, 952, q - 952, 1461, q - 1461, 2687, q - 2687,
            939, q - 939, 2308, q - 2308, 2437, q - 2437, 2388, q - 2388,
            733, q - 733, 2337, q - 2337, 268, q - 268, 641, q - 641,
            1584, q - 1584, 2298, q - 2298, 2037, q - 2037, 3220, q - 3220,
            375, q - 375, 2549, q - 2549, 2090, q - 2090, 1645, q - 1645,
            1063, q - 1063, 319, q - 319, 2773, q - 2773, 757, q - 757,
            2099, q - 2099, 561, q - 561, 2466, q - 2466, 2594, q - 2594,
            2804, q - 2804, 1092, q - 1092, 403, q - 403, 1026, q - 1026,
            1143, q - 1143, 2150, q - 2150, 2775, q - 2775, 886, q - 886,
            1722, q - 1722, 1212, q - 1212, 1874, q - 1874, 1029, q - 1029,
            2110, q - 2110, 2935, q - 2935, 885, q - 885, 2154, q - 2154
        };
 
        private const int n = 256;
        private const ushort q = 3329;
        private const ushort half_q_rounded = (q + 1) >> 1;
 
        private readonly int k384;
        private readonly int cipherTextLength;
        private readonly byte k;
        private readonly int η1;
        private readonly int η2;
        private readonly int dᵤ;
        private readonly int dᵥ;
 
        public ML_KEM(byte k, int η1, int η2, int dᵤ, int dᵥ)
        {
            this.k = k;
            this.η1 = η1;
            this.η2 = η2;
            this.dᵤ = dᵤ;
            this.dᵥ = dᵥ;
 
            this.k384 = this.k * 384;
            this.cipherTextLength = 32 * (this.dᵤ * this.k + this.dᵥ);
        }
 
        public override int PublicKeyLength => 384 * this.k + 32;
 
        public override int PrivateKeyLength => 768 * this.k + 96;
 
        public static byte[] PRF(byte[] Seed, byte Data, int η)
        {
            SHAKE256 HashFunction = new SHAKE256(8 * 64 * η);
            int c = Seed.Length;
            byte[] Bin = new byte[c + 1];
            Buffer.BlockCopy(Seed, 0, Bin, 0, c);
            Bin[c] = Data;
            byte[] Result = HashFunction.ComputeVariable(Bin);
            Clear(Bin);
            return Result;
        }
 
        public static byte[] H(byte[] Data)
        {
            return new SHA3_256().ComputeVariable(Data);
        }
 
        public static byte[] J(byte[] Data)
        {
            return new SHAKE256(n).ComputeVariable(Data);
        }
 
        public static byte[] G(byte[] Data)
        {
            return new SHA3_512().ComputeVariable(Data);
        }
 
        public static byte[] XOF(byte[] Input, int OutputLength)
        {
            return new SHAKE128(OutputLength << 3).ComputeVariable(Input);
        }
 
        public static ushort Compress(ushort Value, int d)
        {
            if (d < 1 || d > 12)
                throw new ArgumentOutOfRangeException(nameof(d), "d must be between 1 and 12.");
 
            uint i = Value;
            i <<= d + 1;
            i /= q;
            i++;
            i >>= 1;
            return (ushort)(i & ushortBitMask[d]);
        }
 
        public static void Compress(ushort[] Value, int d)
        {
            if (d < 1 || d > 12)
                throw new ArgumentOutOfRangeException(nameof(d), "d must be between 1 and 12.");
 
            int j, c = Value.Length;
            uint i;
 
            d++;
 
            for (j = 0; j < c; j++)
            {
                i = Value[j];
                i <<= d;
                i /= q;
                i++;
                i >>= 1;
                Value[j] = (ushort)i;
            }
        }
 
        public static void Compress(ushort[][] Value, int d)
        {
            int i, c = Value.Length;
 
            for (i = 0; i < c; i++)
                Compress(Value[i], d);
        }
 
        public static ushort Decompress(ushort Value, int d)
        {
            if (d < 1 || d > 12)
                throw new ArgumentOutOfRangeException(nameof(d), "d must be between 1 and 12.");
 
            uint i = Value;
            i *= q;
            i >>= d - 1;
            i++;
            i >>= 1;
            return (ushort)i;
        }
 
        public static void Decompress(ushort[] Value, int d)
        {
            if (d < 1 || d > 12)
                throw new ArgumentOutOfRangeException(nameof(d), "d must be between 1 and 12.");
 
            int j, c = Value.Length;
            uint i;
 
            for (j = 0; j < c; j++)
            {
                i = Value[j];
                i *= q;
                i >>= d - 1;
                i++;
                i >>= 1;
                Value[j] = (ushort)i;
            }
        }
 
        public static byte[] ByteEncode(ushort[] Values, int d)
        {
            int c = Values.Length;
            int NrBits = c * d;
            int NrBytes = (NrBits + 7) >> 3;
            byte[] Result = new byte[NrBytes];
 
            ByteEncode(Values, d, Result, 0);
 
            return Result;
        }
 
        public static int ByteEncode(ushort[] Values, int d, byte[] Output, int Index)
        {
            if (d < 1 || d > 12)
                throw new ArgumentOutOfRangeException(nameof(d), "d must be between 1 and 12.");
 
            int i, c = Values.Length;
            int BitOffset = 0;
            int Index0 = Index;
 
            for (i = 0; i < c; i++)
            {
                ushort Value = Values[i];
 
                Value &= ushortBitMask[d];
                Output[Index] |= (byte)(Value << BitOffset);
                BitOffset += d;
                while (BitOffset >= 8)
                {
                    Index++;
                    BitOffset -= 8;
 
                    if (BitOffset > 0)
                        Output[Index] = (byte)(Value >> (d - BitOffset));
                }
            }
 
            return Index - Index0;
        }
 
        public static int ByteEncode(ushort[][] Values, int d, byte[] Output, int Index)
        {
            int Pos = Index;
            int i, c = Values.Length;
 
            for (i = 0; i < c; i++)
            {
                int NrBytes = ByteEncode(Values[i], d, Output, Pos);
                Pos += NrBytes;
            }
 
            return Pos - Index;
        }
 
        public static ushort[] ByteDecode(byte[] Data, int d)
        {
            return ByteDecode(Data, 0, Data.Length, d);
        }
 
        public static ushort[] ByteDecode(byte[] Data, int Offset, int Length, int d)
        {
            if (d < 1 || d > 12)
                throw new ArgumentOutOfRangeException(nameof(d), "d must be between 1 and 12.");
 
            int NrBits = Length << 3;
            int NrIntegers = (NrBits + d - 1) / d;
            ushort[] Result = new ushort[NrIntegers];
            int Pos = 0;
            int BitOffset = 0;
            int BitsLeft;
            int BitsToUse;
 
            while (Length-- > 0)
            {
                byte b = Data[Offset++];
                BitsLeft = 8;
 
                while (BitsLeft > 0)
                {
                    BitsToUse = d - BitOffset;
                    if (BitsToUse > BitsLeft)
                        BitsToUse = BitsLeft;
 
                    Result[Pos] |= (ushort)((b & ushortBitMask[BitsToUse]) << BitOffset);
                    BitOffset += BitsToUse;
                    BitsLeft -= BitsToUse;
                    b >>= BitsToUse;
 
                    if (BitOffset >= d)
                    {
                        Pos++;
                        BitOffset = 0;
                    }
                }
            }
 
            return Result;
        }
 
        private static ushort[] SampleNTT(byte[] Seed)
        {
            SHAKE128 HashFunction = new SHAKE128(0);
            Keccak1600.Context Context = HashFunction.Absorb(Seed);
            ushort[] Result = new ushort[n];
            int Pos = 0;
 
            while (Pos < n)
            {
                byte[] C = Context.Squeeze(3);
                ushort d1 = (ushort)(C[0] + ((C[1] & 15) << 8));
                ushort d2 = (ushort)((C[1] >> 4) + (C[2] << 4));
 
                if (d1 < q)
                    Result[Pos++] = d1;
 
                if (d2 < q && Pos < n)
                    Result[Pos++] = d2;
            }
 
            return Result;
        }
 
        public static ushort[] SamplePolyCBD(int η)
        {
            if (η < 2 || η > 3)
                throw new ArgumentOutOfRangeException(nameof(η), "η must be either 2 or 3.");
 
            return SamplePolyCBD(CreateSeed(η << 6));
        }
 
        public static ushort[] SamplePolyCBD(byte[] Seed)
        {
            int c = Seed.Length;
 
            if (c != 128 && c != 192)
                throw new ArgumentOutOfRangeException(nameof(Seed), "Seed must be either 128 or 192 bytes.");
 
            int η = c >> 6;
            int i, j, k;
            ushort x, y;
            ushort dx, dy;
            ushort[] Result = new ushort[n];
 
            for (i = 0; i < n; i++)
            {
                for (j = x = y = 0; j < η; j++)
                {
                    k = 2 * i * η + j;
                    dx = (Seed[k >> 3] & (1 << (k & 7))) != 0 ? (ushort)1 : (ushort)0;
                    x += dx;    // To avoid different CPU instructions to execute based on if bit is 0 or 1.
 
                    k += η;
                    dy = (Seed[k >> 3] & (1 << (k & 7))) != 0 ? (ushort)1 : (ushort)0;
                    y += dy;    // To avoid different CPU instructions to execute based on if bit is 0 or 1.
                }
 
                if (x < y)
                    Result[i] = (ushort)(q + x - y);
                else
                    Result[i] = (ushort)(x - y);
            }
 
            return Result;
        }
 
        public static void NTT(ushort[] f)
        {
            if (f.Length != n)
                throw new ArgumentException("Polynomial must have " + n + " coefficients.", nameof(f));
 
            int i = 1;
            int j;
            int Len;
            int Start;
            ushort ζ;
            ushort t;
 
            for (Len = n >> 1; Len >= 2; Len >>= 1)
            {
                for (Start = 0; Start < n; Start += Len << 1)
                {
                    ζ = nttTransformZeta[i++];
 
                    for (j = Start; j < Start + Len; j++)
                    {
                        t = (ushort)(ζ * f[j + Len] % q);
                        f[j + Len] = (ushort)((f[j] + q - t) % q);
                        f[j] = (ushort)((f[j] + t) % q);
                    }
                }
            }
        }
 
        public static void NTT(ushort[][] f)
        {
            int i, c = f.Length;
 
            for (i = 0; i < c; i++)
                NTT(f[i]);
        }
 
        public static void InverseNTT(ushort[] f)
        {
            if (f.Length != n)
                throw new ArgumentException("Polynomial must have " + n + " coefficients.", nameof(f));
 
            int i = 127;
            int j;
            int Len;
            int Len2;
            int StartLen;
            int Start;
            ushort ζ;
            ushort t;
 
            for (Len = 2; Len <= 128; Len <<= 1)
            {
                Len2 = Len << 1;
 
                for (Start = 0; Start < n; Start += Len2)
                {
                    ζ = nttTransformZeta[i--];
 
                    StartLen = Start + Len;
                    for (j = Start; j < StartLen; j++)
                    {
                        t = f[j];
                        f[j] = (ushort)((t + f[j + Len]) % q);
                        f[j + Len] = (ushort)(ζ * (f[j + Len] + q - t) % q);
                    }
                }
            }
 
            for (i = 0; i < n; i++)
                f[i] = (ushort)(3303 * f[i] % q);   // 3303 = 128^-1 mod q
        }
 
        public static void InverseNTT(ushort[][] f)
        {
            int i, c = f.Length;
 
            for (i = 0; i < c; i++)
                InverseNTT(f[i]);
        }
 
        public static ushort[] MultiplyNTTs(ushort[] f, ushort[] g)
        {
            ushort[] Result = new ushort[n];
            MultiplyNTTsAndAdd(f, g, Result);
            return Result;
        }
 
        public static void MultiplyNTTsAndAdd(ushort[] f, ushort[] g, ushort[] Result)
        {
            if (f.Length != n || g.Length != n)
                throw new ArgumentException("Polynomials must have " + n + " coefficients.", nameof(f));
 
            int i;
            ushort a0, a1, b0, b1, γ;
 
            for (i = 0; i < n; i += 2)
            {
                a0 = f[i];
                a1 = f[i + 1];
                b0 = g[i];
                b1 = g[i + 1];
                γ = nttTransformZeta2[i >> 1];
 
                Result[i] = (ushort)((Result[i] + a0 * b0 + γ * a1 % q * b1) % q);  // Three multiplications without modulus can generate integer overflow.
                Result[i + 1] = (ushort)((Result[i + 1] + a0 * b1 + a1 * b0) % q);
            }
        }
 
        public static void AddTo(ushort[] f, ushort[] g)
        {
            int i;
 
            for (i = 0; i < n; i++)
                f[i] = (ushort)((f[i] + g[i]) % q);
        }
 
        public static void AddTo(ushort[][] f, ushort[][] g)
        {
            int i, c = f.Length;
            if (g.Length != c)
                throw new ArgumentException("Vectors must have the same number of polynomials.", nameof(g));
 
            for (i = 0; i < c; i++)
                AddTo(f[i], g[i]);
        }
 
        public static void Negate(ushort[] f)
        {
            int i, c = f.Length;
 
            for (i = 0; i < c; i++)
                f[i] = (ushort)((q - f[i]) % q);   // To avoid different CPU instructions to execute based on if bit is 0 or 1.
        }
 
        public static ushort[] DotProductNTT(ushort[][] v1, ushort[][] v2)
        {
            int i, c = v1.Length;
 
            if (v2.Length != c)
                throw new ArgumentException("Vectors must have the same number of polynomials.", nameof(v2));
 
            ushort[] Result = new ushort[n];
 
            for (i = 0; i < c; i++)
                MultiplyNTTsAndAdd(v1[i], v2[i], Result);
 
            return Result;
        }
 
        public K_PKE_Keys K_PKE_KeyGen()
        {
            byte[] Seed = CreateSeed();
 
            K_PKE_Keys Result = this.K_PKE_KeyGen(Seed);
            Clear(Seed);
 
            return Result;
        }
 
        public K_PKE_Keys K_PKE_KeyGen(byte[] d)
        {
            if (d.Length != 32)
                throw new ArgumentException("Seed must be 32 bytes long.", nameof(d));
 
            byte[] Temp = new byte[33];
            Buffer.BlockCopy(d, 0, Temp, 0, 32);
            Temp[32] = this.k;
            byte[] Bin = G(Temp);
            Clear(Temp);
 
            byte[] ρ = new byte[32];
            Buffer.BlockCopy(Bin, 0, ρ, 0, 32);
 
            byte[] σ = new byte[32];
            Buffer.BlockCopy(Bin, 32, σ, 0, 32);
            Clear(Bin);
 
            byte N = 0;
            int i, j;
            ushort[,][] Â = new ushort[this.k, this.k][];
            ushort[][] s = new ushort[this.k][];
            ushort[][] e = new ushort[this.k][];
 
            byte[] B = new byte[34];
            Buffer.BlockCopy(ρ, 0, B, 0, 32);
 
            for (i = 0; i < this.k; i++)
            {
                B[33] = (byte)i;
 
                for (j = 0; j < this.k; j++)
                {
                    B[32] = (byte)j;
                    Â[i, j] = SampleNTT(B);
                }
            }
 
            Clear(B);
 
            for (i = 0; i < this.k; i++)
                s[i] = SamplePolyCBD(PRF(σ, N++, this.η1));
 
            for (i = 0; i < this.k; i++)
                e[i] = SamplePolyCBD(PRF(σ, N++, this.η1));
 
            NTT(s);
            NTT(e);
 
            ushort[][] t = new ushort[this.k][];
            ushort[] f;
 
            for (i = 0; i < this.k; i++)
            {
                t[i] = f = (ushort[])e[i].Clone();
 
                for (j = 0; j < this.k; j++)
                    MultiplyNTTsAndAdd(Â[i, j], s[j], f);
            }
 
            byte[] EncryptionKey = new byte[32 + this.k384];
            byte[] DecryptionKey = new byte[this.k384];
            int Pos = 0;
 
            for (i = 0; i < this.k; i++)
            {
                ByteEncode(t[i], 12, EncryptionKey, Pos);
                ByteEncode(s[i], 12, DecryptionKey, Pos);
                Pos += 384;
            }
 
            Buffer.BlockCopy(ρ, 0, EncryptionKey, Pos, 32);
 
            Clear(ρ);
            Clear(σ);
            Clear(s);
            Clear(t);
            Clear(e);
 
            return new K_PKE_Keys(Â, EncryptionKey, DecryptionKey);
        }
 
        public byte[] K_PKE_Encrypt(byte[] EncryptionKey, byte[] Message)
        {
            return this.K_PKE_Encrypt(ML_KEM_Keys.FromEncapsulationKey(EncryptionKey), Message);
        }
 
        public byte[] K_PKE_Encrypt(ML_KEM_Keys Keys, byte[] Message)
        {
            byte[] Seed = CreateSeed();
 
            byte[] Result = this.K_PKE_Encrypt(Keys, Message, Seed);
            Clear(Seed);
 
            return Result;
        }
 
        public byte[] K_PKE_Encrypt(byte[] EncryptionKey, byte[] Message, byte[] Seed)
        {
            return this.K_PKE_Encrypt(ML_KEM_Keys.FromEncapsulationKey(EncryptionKey), Message, Seed);
        }
 
        public byte[] K_PKE_Encrypt(ML_KEM_Keys Keys, byte[] Message, byte[] Seed)
        {
            if (Keys is null)
                throw new ArgumentNullException(nameof(Keys), "Encryption key cannot be null.");
 
            if (Keys.EncapsulationKey.Length != this.k384 + 32)
                throw new ArgumentException("Encryption key must be 384k+32 bytes long.", nameof(Keys));
 
            if (Message.Length != 32)
                throw new ArgumentException("Message must be 32 bytes long.", nameof(Message));
 
            ushort[][] t = new ushort[this.k][];
            byte[] ρ = new byte[32];
            int Pos;
            int i, j;
            byte N = 0;
            byte k;
 
            for (i = Pos = 0; i < this.k; i++)
            {
                t[i] = ByteDecode(Keys.EncapsulationKey, Pos, 384, 12);
                Pos += 384;
            }
 
            Buffer.BlockCopy(Keys.EncapsulationKey, Pos, ρ, 0, 32);
 
            ushort[][] y = new ushort[this.k][];
            ushort[][] e1 = new ushort[this.k][];
            ushort[] e2;
            ushort[,][] Â = Keys.Â;
 
            if (Â is null)
            {
                Keys. =  = new ushort[this.k, this.k][];
 
                byte[] B = new byte[34];
                Buffer.BlockCopy(ρ, 0, B, 0, 32);
 
                for (i = 0; i < this.k; i++)
                {
                    B[33] = (byte)i;
 
                    for (j = 0; j < this.k; j++)
                    {
                        B[32] = (byte)j;
                        Â[i, j] = SampleNTT(B);
                    }
                }
 
                Clear(B);
            }
            else if (Â.GetLength(0) != this.k || Â.GetLength(1) != this.k)
                throw new ArgumentException("Matrix  must be " + this.k + "x" + this.k + ".", nameof(Â));
 
            for (i = 0; i < this.k; i++)
                y[i] = SamplePolyCBD(PRF(Seed, N++, this.η1));
 
            for (i = 0; i < this.k; i++)
                e1[i] = SamplePolyCBD(PRF(Seed, N++, this.η2));
 
            e2 = SamplePolyCBD(PRF(Seed, N++, this.η2));
 
            NTT(y);
 
            ushort[][] u = new ushort[this.k][];
            ushort[] f;
 
            for (i = 0; i < this.k; i++)
            {
                u[i] = f = new ushort[n];
 
                for (j = 0; j < this.k; j++)
                    MultiplyNTTsAndAdd(Â[j, i], y[j], f);
            }
 
            InverseNTT(u);
            AddTo(u, e1);
 
            ushort[] v = DotProductNTT(t, y);
 
            InverseNTT(v);
 
            ushort[] μ = new ushort[n];
 
            for (i = j = 0, k = 1; i < n; i++)
            {
                μ[i] = (Message[j] & k) != 0 ? half_q_rounded : (ushort)0;
 
                k <<= 1;
                if (k == 0)
                {
                    k = 1;
                    j++;
                }
            }
 
            AddTo(v, e2);
            AddTo(v, μ);
 
            int dᵤk32 = this.dᵤ * this.k << 5;
            int dᵥ32 = this.dᵥ << 5;
            byte[] CipherText = new byte[dᵤk32 + dᵥ32];
 
            Compress(u, this.dᵤ);
            ByteEncode(u, this.dᵤ, CipherText, 0);
 
            Compress(v, this.dᵥ);
            ByteEncode(v, this.dᵥ, CipherText, dᵤk32);
 
            return CipherText;
        }
 
        public byte[] K_PKE_Decrypt(byte[] DecryptionKey, byte[] CipherText)
        {
            if (DecryptionKey is null)
                throw new ArgumentNullException(nameof(DecryptionKey), "Decryption key cannot be null.");
 
            if (DecryptionKey.Length != this.k384)
                throw new ArgumentException("Decryption key must be 384k bytes long.", nameof(DecryptionKey));
 
            int dᵤ32 = this.dᵤ << 5;
            int dᵤk32 = dᵤ32 * this.k;
            int dᵥ32 = this.dᵥ << 5;
 
            if (CipherText.Length != dᵤk32 + dᵥ32)
                throw new ArgumentException("Message must be " + (dᵤk32 + dᵥ32) + " bytes long.", nameof(CipherText));
 
            ushort[][] u = new ushort[this.k][];
            ushort[] v;
            int Pos;
            int i;
 
            for (i = Pos = 0; i < this.k; i++)
            {
                u[i] = ByteDecode(CipherText, Pos, dᵤ32, this.dᵤ);
                Decompress(u[i], this.dᵤ);
                Pos += dᵤ32;
            }
 
            v = ByteDecode(CipherText, Pos, dᵥ32, this.dᵥ);
            Decompress(v, this.dᵥ);
 
            ushort[][] s = new ushort[this.k][];
 
            for (i = Pos = 0; i < this.k; i++)
            {
                s[i] = ByteDecode(DecryptionKey, Pos, 384, 12);
                Pos += 384;
            }
 
            NTT(u);
 
            ushort[] w = DotProductNTT(s, u);
 
            InverseNTT(w);
            Negate(w);
            AddTo(w, v);
 
            Compress(w, 1);
            return ByteEncode(w, 1);
        }
 
        public ML_KEM_Keys KeyGen()
        {
            return this.KeyGen(false);
        }
 
        public ML_KEM_Keys KeyGen(bool ReturnSeed)
        {
            byte[] d = CreateSeed();
            byte[] z = CreateSeed();
 
            ML_KEM_Keys Result = this.KeyGen_Internal(d, z, ReturnSeed);
 
            if (!ReturnSeed)
            {
                Clear(d);
                Clear(z);
            }
 
            return Result;
        }
 
        public ML_KEM_Keys KeyGen_Internal(byte[] d, byte[] z)
        {
            return this.KeyGen_Internal(d, z, false);
        }
 
        public ML_KEM_Keys KeyGen_Internal(byte[] d, byte[] z, bool ReturnSeed)
        {
            if (z.Length != 32)
                throw new ArgumentException("Seed must be 32 bytes long.", nameof(z));
 
            K_PKE_Keys Keys = this.K_PKE_KeyGen(d);
            byte[] DecryptionKey = new byte[768 * this.k + 96];
            int Pos;
 
            Buffer.BlockCopy(Keys.DecryptionKey, 0, DecryptionKey, 0, Pos = this.k384);
            Buffer.BlockCopy(Keys.EncryptionKey, 0, DecryptionKey, Pos, Pos + 32);
            Pos += Pos + 32;
 
            byte[] Bin = H(Keys.EncryptionKey);
            Buffer.BlockCopy(Bin, 0, DecryptionKey, Pos, 32);
            Pos += 32;
 
            Buffer.BlockCopy(z, 0, DecryptionKey, Pos, 32);
 
            if (!ReturnSeed)
                return new ML_KEM_Keys(Keys, DecryptionKey, null);
 
            byte[] Seed = new byte[64];
 
            Buffer.BlockCopy(d, 0, Seed, 0, 32);
            Buffer.BlockCopy(z, 0, Seed, 32, 32);
 
            return new ML_KEM_Keys(Keys, DecryptionKey, Seed);
        }
 
        public ML_KEM_Keys KeyGen_FromSeed(byte[] Seed, bool ReturnSeed)
        {
            if (Seed is null)
                throw new ArgumentNullException(nameof(Seed));
 
            if (Seed.Length != 64)
                throw new ArgumentException("Seed must be 64 bytes long.", nameof(Seed));
 
            byte[] d = new byte[32];
            byte[] z = new byte[32];
 
            Buffer.BlockCopy(Seed, 0, d, 0, 32);
            Buffer.BlockCopy(Seed, 32, z, 0, 32);
 
            return this.KeyGen_Internal(d, z, ReturnSeed);
        }
 
        public ML_KEM_Encapsulation Encapsulate(byte[] EncapsulationKey)
        {
            return this.Encapsulate(ML_KEM_Keys.FromEncapsulationKey(EncapsulationKey));
        }
 
        public ML_KEM_Encapsulation Encapsulate(ML_KEM_Keys Keys)
        {
            if (Keys.EncapsulationKey.Length < 384)
                throw new ArgumentException("Encapsulation key length mismatch.", nameof(Keys));
 
            ushort[] x = ByteDecode(Keys.EncapsulationKey, 0, 384, 12);
            int i, c = x.Length;
 
            for (i = 0; i < c; i++)
            {
                if (x[i] >= q)
                    throw new ArgumentException("Invalid encapsulation key.", nameof(Keys));
            }
 
            byte[] m = CreateSeed();
 
            ML_KEM_Encapsulation Result = this.Encapsulate_Internal(Keys, m);
            Clear(m);
 
            return Result;
        }
 
        public ML_KEM_Encapsulation Encapsulate_Internal(byte[] EncapsulationKey, byte[] m)
        {
            return this.Encapsulate_Internal(ML_KEM_Keys.FromEncapsulationKey(EncapsulationKey), m);
        }
 
        public ML_KEM_Encapsulation Encapsulate_Internal(ML_KEM_Keys Keys, byte[] m)
        {
            if (Keys.EncapsulationKey.Length != this.k384 + 32)
                throw new ArgumentException("Encapsulation key length mismatch.", nameof(Keys));
 
            if (m.Length != 32)
                throw new ArgumentException("Message length mismatch.", nameof(m));
 
            byte[] Bin = new byte[64];
 
            Buffer.BlockCopy(m, 0, Bin, 0, 32);
            Buffer.BlockCopy(H(Keys.EncapsulationKey), 0, Bin, 32, 32);
 
            byte[] Bin2 = G(Bin);
            Clear(Bin);
 
            byte[] K = new byte[32];
            Buffer.BlockCopy(Bin2, 0, K, 0, 32);
 
            byte[] r = new byte[32];
            Buffer.BlockCopy(Bin2, 32, r, 0, 32);
            Clear(Bin2);
 
            byte[] c = this.K_PKE_Encrypt(Keys, m, r);
            Clear(r);
 
            return new ML_KEM_Encapsulation(K, c);
        }
 
        public byte[] Decapsulate(byte[] DecapsulationKey, byte[] c)
        {
            if (DecapsulationKey.Length != 768 * this.k + 96)
                throw new ArgumentException("Decapsulation key length mismatch.", nameof(DecapsulationKey));
 
            byte[] Bin = new byte[this.k384 + 32];
            Buffer.BlockCopy(DecapsulationKey, this.k384, Bin, 0, this.k384 + 32);
 
            byte[] Test = H(Bin);
            Clear(Bin);
 
            for (int i = 0, j = (this.k384 << 1) + 32; i < 32; i++, j++)
            {
                if (Test[i] != DecapsulationKey[j])
                    throw new ArgumentException("Invalid decapsulation key.", nameof(DecapsulationKey));
            }
 
            return this.Decapsulate_Internal(DecapsulationKey, c);
        }
 
        public int CipherTextLength => this.cipherTextLength;
 
        public byte[] Decapsulate_Internal(byte[] DecapsulationKey, byte[] c)
        {
            if (DecapsulationKey.Length != 768 * this.k + 96)
                throw new ArgumentException("Decapsulation key length mismatch.", nameof(DecapsulationKey));
 
            if (c.Length != this.cipherTextLength)
                throw new ArgumentException("Cipher text length mismatch.", nameof(c));
 
            int Pos;
            byte[] DecryptionKey = new byte[Pos = this.k384];
            Buffer.BlockCopy(DecapsulationKey, 0, DecryptionKey, 0, Pos);
 
            int Len;
            byte[] EncryptionKey = new byte[Len = (768 - 384) * this.k + 32];
            Buffer.BlockCopy(DecapsulationKey, Pos, EncryptionKey, 0, Len);
            Pos += Len;
 
            byte[] h = new byte[32];
            Buffer.BlockCopy(DecapsulationKey, Pos, h, 0, 32);
 
            byte[] m = this.K_PKE_Decrypt(DecryptionKey, c);
 
            byte[] Bin = new byte[64];
            Buffer.BlockCopy(m, 0, Bin, 0, 32);
            Buffer.BlockCopy(h, 0, Bin, 32, 32);
 
            byte[] Bin2 = G(Bin);
 
            byte[] K = new byte[32];
            Buffer.BlockCopy(Bin2, 0, K, 0, 32);
 
            byte[] r = new byte[32];
            Buffer.BlockCopy(Bin2, 32, r, 0, 32);
 
            Clear(Bin2);
 
            // Fix of the decapsulation method to comply with FIPS-203 (Algorithm 18) and
            // to avoid a weakened IND-CCA2 attack, and a constant-time comparison check,
            // provided and republished, with permission, by Frank von der Loos.
 
            byte[] zc = new byte[32 + c.Length];                       // z ‖ full c (32 + cipherTextLength)
            Buffer.BlockCopy(DecapsulationKey, Pos + 32, zc, 0, 32);   // z
            Buffer.BlockCopy(c, 0, zc, 32, c.Length);                  // full c
            byte[] K2 = J(zc);
            Clear(zc);
 
            byte[] c2 = this.K_PKE_Encrypt(EncryptionKey, m, r);
            Clear(m);
            Clear(r);
 
            int diff = 0;
 
            for (int i = 0; i < c.Length; i++)   // full ciphertext, no early break, time constant
                diff |= c[i] ^ c2[i];
 
            Clear(c2);
 
            // diff == 0  -> ciphertext valid    -> K  (real secret)
            // diff != 0  -> implicit rejection   -> K2 (pseudo-random)
            byte mask = (byte)((diff | -diff) >> 31);   // 0x00 for equal, 0xFF is different
 
            byte[] Result = new byte[32];
            for (int i = 0; i < 32; i++)
                Result[i] = (byte)((K[i] & ~mask) | (K2[i] & mask));
 
            Clear(K);
            Clear(K2);
 
            return Result;
        }
    }
}