Decode small Flate streams through asyncGetBytes (bug 2086217)

src/core/flate_stream.js

@@ -122,14 +122,16 @@ const fixedDistCodeTab = [
 ];
 
 class FlateStream extends DecodeStream {
-  // Only mark the FlateStream as async if its compressed length is above
-  // this threshold. Using the browser;s built-in `DecompressionStream` is
-  // faster than the JS decoder per byte, but it has a higher fixed cost
-  // that is not worth paying on small streams.
-  static #MIN_ASYNC_COMPRESSED_LENGTH = 1024;
+  // Use JS for streams with a small compressed size.
+  static #MIN_NATIVE_COMPRESSED_LENGTH = 1024;
+
+  // Prefer native decoding for larger expected outputs.
+  static #MAX_JS_DECODED_LENGTH = 16384;
 
   #isAsync = true;
 
+  #decodeInJS = false;
+
   constructor(str, maybeLength) {
     super(maybeLength);
 
@@ -138,10 +140,10 @@ class FlateStream extends DecodeStream {
 
     if (
       str instanceof Stream &&
-      str.length < FlateStream.#MIN_ASYNC_COMPRESSED_LENGTH
+      str.length < FlateStream.#MIN_NATIVE_COMPRESSED_LENGTH
     ) {
       assert(!str.isAsync, "Stream sub-classes are never asynchronous");
-      this.#isAsync = false;
+      this.#decodeInJS = true;
     }
 
     const cmf = str.getByte();
@@ -166,17 +168,32 @@ class FlateStream extends DecodeStream {
   }
 
   async getImageData(length, _decoderOptions) {
-    const data = await this.asyncGetBytes();
+    const data = await this.asyncGetBytes(length);
     if (!data) {
       return this.getBytes(length);
     }
     return data.length <= length ? data : data.subarray(0, length);
   }
 
-  async asyncGetBytes() {
+  /**
+   * @param {number} [decodedLength] - Expected output length, or 0 if unknown.
+   */
+  async asyncGetBytes(decodedLength = 0) {
     if (!this.#isAsync) {
       return null;
     }
+    if (
+      this.#decodeInJS &&
+      (decodedLength <= FlateStream.#MAX_JS_DECODED_LENGTH ||
+        this.bufferLength > 0 ||
+        this.eof)
+    ) {
+      // Returning bytes lets callers wrap them in a plain `Stream`.
+      this.reset();
+      return this.getBytes(decodedLength);
+    }
+    // After native decoding succeeds, the JS decoder cannot reuse the input.
+    this.#decodeInJS = false;
     const { decompressed, compressed } =
       await this.asyncGetBytesFromDecompressionStream("deflate");
     if (decompressed) {

test/unit/stream_spec.js

@@ -135,6 +135,46 @@ describe("stream", function () {
         expect(await flate.asyncGetBytes()).toEqual(new Uint8Array(0));
         expect(flate.isAsync).toEqual(true);
       });
+
+      it("should decode a small stream in JS and return its bytes", async function () {
+        const small = data.subarray(0, 300);
+        const flate = new FlateStream(new Stream(await deflate(small)));
+        const spy = spyOn(
+          flate,
+          "asyncGetBytesFromDecompressionStream"
+        ).and.callThrough();
+
+        expect(flate.isAsync).toEqual(true);
+        expect(await flate.asyncGetBytes()).toEqual(small);
+        expect(await flate.asyncGetBytes()).toEqual(small);
+        expect(spy).not.toHaveBeenCalled();
+      });
+
+      it("should choose a decoder for a small image stream and reuse it", async function () {
+        const zeros = new Uint8Array(100000);
+        const compressedZeros = await deflate(zeros);
+        expect(compressedZeros.length).toBeLessThan(1024);
+
+        for (const lengths of [
+          [zeros.length, 1000],
+          [1000, zeros.length],
+        ]) {
+          const flate = new FlateStream(new Stream(compressedZeros));
+          const spy = spyOn(
+            flate,
+            "asyncGetBytesFromDecompressionStream"
+          ).and.callThrough();
+          for (const length of lengths) {
+            flate.reset();
+            expect(await flate.getImageData(length)).toEqual(
+              zeros.subarray(0, length)
+            );
+          }
+          expect(spy).toHaveBeenCalledTimes(
+            lengths[0] === zeros.length ? lengths.length : 0
+          );
+        }
+      });
     });
   });