package org.bouncycastle.pqc.crypto.frodo;

import org.bouncycastle.crypto.BlockCipher;
import org.bouncycastle.crypto.Xof;
import org.bouncycastle.crypto.digests.SHAKEDigest;
import org.bouncycastle.crypto.engines.AESEngine;
import org.bouncycastle.crypto.params.KeyParameter;
import org.bouncycastle.util.Arrays;
import org.bouncycastle.util.Pack;

abstract class FrodoMatrixGenerator
{
    int n;
    int q;

    public FrodoMatrixGenerator(int n, int q)
    {
        this.n = n;
        this.q = q;
    }

    abstract short[] genMatrix(byte[] seedA);

    static class Shake128MatrixGenerator
            extends FrodoMatrixGenerator
    {
        public Shake128MatrixGenerator(int n, int q)
        {
            super(n, q);
        }

        short[] genMatrix(byte[] seedA)
        {
            short[] A = new short[n*n];
            short i, j;
            byte[] b, tmp = new byte[(16 * n) / 8];
            for (i = 0; i < n; i++)
            {
                // 1. b = i || seedA in {0,1}^{16 + len_seedA}, where i is encoded as a 16-bit integer in little-endian byte order
                b = Arrays.concatenate(Pack.shortToLittleEndian(i), seedA);

                // 2. c_{i,0} || c_{i,1} || ... || c_{i,n-1} = SHAKE128(b, 16n) (length in bits) where each c_{i,j} is parsed as a 16-bit integer in little-endian byte order format
                Xof digest = new SHAKEDigest(128);
                digest.update(b, 0, b.length);
                digest.doFinal(tmp, 0, tmp.length);
                for (j = 0; j < n; j++)
                {
                    A[i*n+j] = (short) (Pack.littleEndianToShort(tmp, 2 * j) % q);
                }
            }
            return A;
        }

    }
    static class Aes128MatrixGenerator
            extends FrodoMatrixGenerator
    {
        private final BlockCipher cipher;
        public Aes128MatrixGenerator(int n, int q)
        {
            super(n, q);
            cipher = new AESEngine();
        }

        short[] genMatrix(byte[] seedA)
        {
            //        """Generate matrix A using AES-128 (FrodoKEM specification, Algorithm 7)"""
            //        A = [[None for j in range(self.n)] for i in range(self.n)]
            short[] A = new short[n*n];
            byte[] b = new byte[16];
            byte[] c = new byte[16];

            KeyParameter kp = new KeyParameter(seedA);
            cipher.init(true, kp);

            // 1. for i = 0; i < n; i += 1
            for (int i = 0; i < n; i++)
            {
                // 2. for j = 0; j < n; j += 8
                for (int j = 0; j < n; j+=8)
                {

                    // 3. b = i || j || 0 || ... || 0 in {0,1}^128, where i and j are encoded as 16-bit integers in little-endian byte order
                    System.arraycopy(Pack.shortToLittleEndian((short) (i&0xffff)), 0, b, 0, 2);
                    System.arraycopy(Pack.shortToLittleEndian((short) (j&0xffff)), 0, b, 2, 2);
                    //                b = bytearray(16)
                    //                struct.pack_into('<H', b, 0, i)
                    //                struct.pack_into('<H', b, 2, j)
                    // 4. c = AES128(seedA, b)
                    cipher.processBlock(b, 0, c, 0);
                    // 5. for k = 0; k < 8; k += 1
                    for (int k = 0; k < 8; k++)
                    {
                        // 6. A[i][j+k] = c[k] where c is treated as a sequence of 8 16-bit integers each in little-endian byte order
                        A[i*n+ j + k] = (short) (Pack.littleEndianToShort(c, 2 * k) % q);
                    }
                }
            }
            return A;
        }
    }
}

@@ -275,45 +275,8 @@
 er;%0A
-import org.bouncycastle.util.Arrays;%0A
 impo
@@ -306,16 +306,16 @@
 l.Pack;%0A
+
 %0Aabstrac
@@ -833,19 +833,16 @@
   byte%5B%5D
- b,
  tmp = n
@@ -864,16 +864,176 @@
 ) / 8%5D;%0A
+            byte%5B%5D b = new byte%5B2 + seedA.length%5D;%0A            System.arraycopy(seedA, 0, b, 2, seedA.length);%0A%0A            Xof digest = new SHAKEDigest(128);%0A%0A
         
@@ -1225,31 +1225,8 @@
     
-b = Arrays.concatenate(
 Pack
@@ -1251,16 +1251,14 @@
 an(i
-)
 , 
-seedA
+b, 0
 );%0A%0A
@@ -1442,59 +1442,8 @@
 mat%0A
-                Xof digest = new SHAKEDigest(128);%0A
     
@@ -1672,19 +1672,25 @@
  2 * j) 
-%25 q
+& (q - 1)
 );%0A     
@@ -1748,23 +1748,23 @@
       %7D%0A
-%0A
     %7D%0A
+%0A
     stat
@@ -1845,50 +1845,8 @@
   %7B%0A
-        private final BlockCipher cipher;%0A
     
@@ -1931,46 +1931,8 @@
 q);%0A
-            cipher = new AESEngine();%0A
     
@@ -2298,48 +2298,43 @@
     
-KeyParameter kp = new KeyParameter(seedA
+BlockCipher cipher = new AESEngine(
 );%0A 
@@ -2366,10 +2366,31 @@
 ue, 
-kp
+new KeyParameter(seedA)
 );%0A%0A
@@ -2596,17 +2596,16 @@
       %7B%0A
-%0A
         
@@ -2757,33 +2757,16 @@
         
-System.arraycopy(
 Pack.sho
@@ -2793,32 +2793,15 @@
 ort)
- (i&0xffff)), 0
+i
 , b, 0
-, 2
 );%0A 
@@ -2819,33 +2819,16 @@
         
-System.arraycopy(
 Pack.sho
@@ -2855,230 +2855,17 @@
 ort)
- (j&0xffff)), 0, b, 2, 2);%0A                    //                b = bytearray(16)%0A                    //                struct.pack_into('%3CH', b, 0, i)%0A                    //                struct.pack_into('%3CH', b, 2, j)
+j, b, 2);
 %0A   
@@ -3304,11 +3304,17 @@
  k) 
-%25 q
+& (q - 1)
 );%0A 

Refactor FrodoMatrixGenerator (performance)