Accepting review comments.

src/java.base/share/classes/sun/security/provider/SHA3.java

@@ -67,9 +67,21 @@ public abstract class SHA3 extends DigestBase {
         0x8000000000008080L, 0x80000001L, 0x8000000080008008L,
     };
 
+    // the starting 3 or 5 bits of the domain separator and 10*1 padding
+    // (11111... for the SHAKE functions and 011... for the SHA3-nnn ones)
     private final byte suffix;
+
+    // the stae matrix flattened into an array
     private long[] state = new long[DM*DM];
-    private int squeezeOffset = -1;
+
+    // The byte offset in the state where the next sqeeze() will start.
+    // -1 indicates that either we are in the absorbing phase (only
+    // update() calls were made so far)
+    // The first squeeze() call (after a possibly empty sequence of update()
+    // calls) will set it to 0 at its start.
+    // When a squeeze() call uses up all available bytes from this state
+    // and so a new keccak() call is made, squeezeOffset is reset to 0.
+    protected int squeezeOffset = -1;
 
     static final VarHandle asLittleEndian
             = MethodHandles.byteArrayViewVarHandle(long[].class,
@@ -120,6 +132,8 @@ void finishAbsorb() {
      * DigestBase calls implReset() when necessary.
      */
     void implDigest(byte[] out, int ofs) {
+        // Moving this allocation to the block where t is used causes a little
+        // performance drop, that is why it is here.
         byte[] byteState = new byte[8];
         if (engineGetDigestLength() == 0) {
             // This is an XOF, so the digest() call is illegal.
@@ -139,22 +153,21 @@ void implDigest(byte[] out, int ofs) {
             numBytes -= availableBytes;
             keccak();
         }
-        int numLongs = (numBytes + 7) / 8;
+        int numLongs = numBytes / 8;
 
-        for (int i = 0; i < numLongs - 1; i++) {
+        for (int i = 0; i < numLongs; i++) {
             asLittleEndian.set(out, ofs, state[i]);
             ofs += 8;
         }
-        if (numBytes == numLongs * 8) {
-            asLittleEndian.set(out, ofs, state[numLongs - 1]);
-        } else {
-            asLittleEndian.set(byteState, 0, state[numLongs - 1]);
-            System.arraycopy(byteState, 0,
-                    out, ofs, numBytes - (numLongs - 1) * 8);
+        if (numBytes % 8 != 0) {
+            asLittleEndian.set(byteState, 0, state[numLongs]);
+            System.arraycopy(byteState, 0, out, ofs, numBytes % 8);
         }
     }
 
     void implSqueeze(byte[]output, int offset, int numBytes) {
+        // Moving this allocation to the block where t is used causes a little
+        // performance drop, that is why it is here.
         byte[] byteState = new byte[8];
         if (engineGetDigestLength() != 0) {
             // This is not an XOF, so the squeeze() call is illegal.
@@ -168,12 +181,6 @@ void implSqueeze(byte[]output, int offset, int numBytes) {
 
         int availableBytes = blockSize - squeezeOffset;
 
-        if (availableBytes == 0) {
-            keccak();
-            squeezeOffset = 0;
-            availableBytes = blockSize;
-        }
-
         while (numBytes > availableBytes) {
             int longOffset = squeezeOffset / 8;
             int bytesToCopy = 0;
@@ -207,6 +214,8 @@ void implSqueeze(byte[]output, int offset, int numBytes) {
             numBytes -= bytesToCopy;
             offset += bytesToCopy;
             squeezeOffset += bytesToCopy;
+
+            if (numBytes == 0) return;
         }
 
         int longsToConvert =
@@ -428,34 +437,27 @@ public SHA512() {
         }
     }
 
-
-    /*
-     * The SHAKE128 extendable output function.
-     */
-    public static final class SHAKE128 extends SHA3 {
-        // d is the required number of output bytes.
-        // If this constructor is used with d > 0, the squeezing methods
-        // will throw a ProviderException.
-        public SHAKE128(int d) {
-            super("SHAKE128", d, (byte) 0x1F, 32);
-        }
-
-        // If this constructor is used to get an instance of the class, then,
-        // after the last update, one can get the generated bytes using the
-        // squeezing methods.
-        // Calling digest method will throw a ProviderException.
-        public SHAKE128() {
-            super("SHAKE128", 0, (byte) 0x1F, 32);
+    public static abstract class SHA3XOF extends SHA3 {
+        public SHA3XOF(String name, int digestLength, byte offset, int c) {
+            super(name, digestLength, offset, c);
         }
-
         public void update(byte in) {
+            if (squeezeOffset  != -1) {
+                throw new ProviderException("update() after squeeze() is not allowed.");
+            }
             engineUpdate(in);
         }
         public void update(byte[] in, int off, int len) {
+            if (squeezeOffset  != -1) {
+                throw new ProviderException("update() after squeeze() is not allowed.");
+            }
             engineUpdate(in, off, len);
         }
 
         public void update(byte[] in) {
+            if (squeezeOffset  != -1) {
+                throw new ProviderException("update() after squeeze() is not allowed.");
+            }
             engineUpdate(in, 0, in.length);
         }
 
@@ -474,10 +476,31 @@ public void reset() {
             engineReset();
         }
     }
+
+    /*
+     * The SHAKE128 extendable output function.
+     */
+    public static final class SHAKE128 extends SHA3XOF {
+        // d is the required number of output bytes.
+        // If this constructor is used with d > 0, the squeezing methods
+        // will throw a ProviderException.
+        public SHAKE128(int d) {
+            super("SHAKE128", d, (byte) 0x1F, 32);
+        }
+
+        // If this constructor is used to get an instance of the class, then,
+        // after the last update, one can get the generated bytes using the
+        // squeezing methods.
+        // Calling digest method will throw a ProviderException.
+        public SHAKE128() {
+            super("SHAKE128", 0, (byte) 0x1F, 32);
+        }
+    }
+
     /*
      * The SHAKE256 extendable output function.
      */
-    public static final class SHAKE256 extends SHA3 {
+    public static final class SHAKE256 extends SHA3XOF {
         // d is the required number of output bytes.
         // If this constructor is used with d > 0, the squeezing methods will
         // throw a ProviderException.
@@ -492,32 +515,6 @@ public SHAKE256(int d) {
         public SHAKE256() {
             super("SHAKE256", 0, (byte) 0x1F, 64);
         }
-
-        public void update(byte in) {
-            engineUpdate(in);
-        }
-        public void update(byte[] in, int off, int len) {
-            engineUpdate(in, off, len);
-        }
-
-        public void update(byte[] in) {
-            engineUpdate(in, 0, in.length);
-        }
-
-        public byte[] digest() {
-            return engineDigest();
-        }
-
-        public void squeeze(byte[] output, int offset, int numBytes) {
-            implSqueeze(output, offset, numBytes);
-        }
-        public byte[] squeeze(int numBytes) {
-            return implSqueeze(numBytes);
-        }
-
-        public void reset() {
-            engineReset();
-        }
     }
 
 }