Decompress .z files in prefetch

eqrion · eqrion · commit 5055ed2b82b3 · 2025-08-29T10:49:55.000-05:00
diff --git a/JetStreamDriver.js b/JetStreamDriver.js
@@ -81,6 +81,10 @@ if (typeof(URLSearchParams) !== "undefined") {
         globalThis.prefetchResources = getBoolParam(urlParameters, "prefetchResources");
 }
 
+if (!isInBrowser) {
+    load("wasm-zlib.js");
+}
+
 if (!globalThis.prefetchResources)
     console.warn("Disabling resource prefetching!");
 
@@ -189,6 +193,47 @@ function uiFriendlyDuration(time) {
     return `${time.toFixed(3)} ms`;
 }
 
+// Files can be zlib compressed to reduce the download size. We don't use http
+// compression because we support running from the shell and don't want to
+// require a complicated server setup.
+//
+// zlib was chosen because we already have it in tree for the wasm-zlib test.
+function isCompressed(name) {
+    return name.endsWith(".z");
+}
+
+// Fallback for shell environments without TextDecoder. This only handles valid
+// UTF-8, invalid buffers will lead to unexpected results.
+function decodeUTF8(int8Array) {
+    let result = '';
+    let i = 0;
+    while (i < int8Array.length) {
+        let byte1 = int8Array[i++];
+        if (byte1 < 0x80) {
+            // 1-byte sequence (ASCII)
+            result += String.fromCharCode(byte1);
+        } else if ((byte1 & 0xE0) === 0xC0) {
+            // 2-byte sequence
+            let byte2 = int8Array[i++];
+            result += String.fromCharCode(((byte1 & 0x1F) << 6) | (byte2 & 0x3F));
+        } else if ((byte1 & 0xF0) === 0xE0) {
+            // 3-byte sequence
+            let byte2 = int8Array[i++];
+            let byte3 = int8Array[i++];
+            result += String.fromCharCode(((byte1 & 0x0F) << 12) | ((byte2 & 0x3F) << 6) | (byte3 & 0x3F));
+        } else if ((byte1 & 0xF8) === 0xF0) {
+            // 4-byte sequence (needs surrogate pairs)
+            let byte2 = int8Array[i++];
+            let byte3 = int8Array[i++];
+            let byte4 = int8Array[i++];
+            let codePoint = ((byte1 & 0x07) << 18) | ((byte2 & 0x3F) << 12) | ((byte3 & 0x3F) << 6) | (byte4 & 0x3F);
+            codePoint -= 0x10000;
+            result += String.fromCharCode(0xD800 + (codePoint >> 10), 0xDC00 + (codePoint & 0x3FF));
+        }
+    }
+    return result;
+}
+
 // TODO: Cleanup / remove / merge. This is only used for caching loads in the
 // non-browser setting. In the browser we use exclusively `loadCache`, 
 // `loadBlob`, `doLoadBlob`, `prefetchResourcesForBrowser` etc., see below.
@@ -201,14 +246,25 @@ class ShellFileLoader {
     // share common code.
     load(url) {
         console.assert(!isInBrowser);
-        if (!globalThis.prefetchResources)
+
+        // If we aren't supposed to prefetch this and don't need to decompress it,
+        // then return code snippet that will load the url on-demand.
+        let compressed = isCompressed(url);
+        if (!compressed && !globalThis.prefetchResources)
             return `load("${url}");`
 
         if (this.requests.has(url)) {
             return this.requests.get(url);
         }
 
-        const contents = readFile(url);
+        let contents;
+        if (isCompressed(url)) {
+            let bytes = new Int8Array(read(url, "binary"));
+            bytes = zlib.decompress(bytes);
+            contents = decodeUTF8(bytes);
+        } else {
+            contents = readFile(url);
+        }
         this.requests.set(url, contents);
         return contents;
     }
@@ -260,10 +316,14 @@ class Driver {
             performance.mark("update-ui");
             benchmark.updateUIAfterRun();
 
-            if (isInBrowser && globalThis.prefetchResources) {
+            if (isInBrowser) {
                 const cache = JetStream.blobDataCache;
                 for (const file of benchmark.files) {
                     const blobData = cache[file];
+                    // If we didn't prefetch this resource, then no need to free it
+                    if (!blobData.blob) {
+                        continue
+                    }
                     blobData.refCount--;
                     if (!blobData.refCount)
                         cache[file] = undefined;
@@ -415,6 +475,7 @@ class Driver {
 
     async prefetchResources() {
         if (!isInBrowser) {
+            await zlib.initialize();
             for (const benchmark of this.benchmarks)
                 benchmark.prefetchResourcesForShell();
             return;
@@ -629,6 +690,11 @@ class Scripts {
 }
 
 class ShellScripts extends Scripts {
+    constructor() {
+        super();
+        this.blobs = [];
+    }
+
     run() {
         let globalObject;
         let realm;
@@ -655,13 +721,23 @@ class ShellScripts extends Scripts {
             currentReject
         };
 
+        // Store shellBlobs on JetStreamBlobs so that getBinary can find them.
+        globalObject.JetStreamBlobs = {};
+        for (const [name, value] of this.blobs) {
+            globalObject.JetStreamBlobs[name] = value;
+        }
+
         globalObject.performance ??= performance;
         for (const script of this.scripts)
             globalObject.loadString(script);
 
         return isD8 ? realm : globalObject;
     }
 
+    addBlobs(blobs) {
+        this.blobs.push(...blobs);
+    }
+
     add(text) {
         this.scripts.push(text);
     }
@@ -694,7 +770,6 @@ class BrowserScripts extends Scripts {
         return magicFrame;
     }
 
-
     add(text) {
         this.scripts.push(`<script>${text}</script>`);
     }
@@ -714,6 +789,7 @@ class Benchmark {
         this.allowUtf16 = !!plan.allowUtf16;
         this.scripts = null;
         this.preloads = null;
+        this.shellBlobs = null;
         this.results = [];
         this._state = BenchmarkState.READY;
     }
@@ -820,10 +896,14 @@ class Benchmark {
         if (!!this.plan.exposeBrowserTest)
             scripts.addBrowserTest();
 
+        if (this.shellBlobs) {
+            scripts.addBlobs(this.shellBlobs);
+        }
         if (this.plan.preload) {
             let preloadCode = "";
-            for (let [ variableName, blobURLOrPath ] of this.preloads)
+            for (let [ variableName, blobURLOrPath ] of this.preloads) {
                 preloadCode += `JetStream.preload.${variableName} = "${blobURLOrPath}";\n`;
+            }
             scripts.add(preloadCode);
         }
 
@@ -838,7 +918,7 @@ class Benchmark {
         } else {
             const cache = JetStream.blobDataCache;
             for (const file of this.plan.files) {
-                scripts.addWithURL(globalThis.prefetchResources ? cache[file].blobURL : file);
+                scripts.addWithURL(cache[file].blobURL);
             }
         }
 
@@ -889,10 +969,15 @@ class Benchmark {
 
     async doLoadBlob(resource) {
         const blobData = JetStream.blobDataCache[resource];
-        if (!globalThis.prefetchResources) {
+
+        // If we aren't supposed to prefetch this and don't need to decompress it,
+        // then set the blobURL to just be the resource URL.
+        const compressed = isCompressed(resource);
+        if (!compressed && !globalThis.prefetchResources) {
             blobData.blobURL = resource;
             return blobData;
         }
+
         let response;
         let tries = 3;
         while (tries--) {
@@ -908,7 +993,15 @@ class Benchmark {
                 continue;
             throw new Error("Fetch failed");
         }
-        const blob = await response.blob();
+
+        // If we need to decompress this, then run it through a decompression
+        // stream.
+        if (compressed) {
+            const stream = response.body.pipeThrough(new DecompressionStream('deflate'))
+            response = new Response(stream);
+        }
+
+        let blob = await response.blob();
         blobData.blob = blob;
         blobData.blobURL = URL.createObjectURL(blob);
         return blobData;
@@ -1044,7 +1137,26 @@ class Benchmark {
         this.scripts = this.plan.files.map(file => shellFileLoader.load(file));
 
         console.assert(this.preloads === null, "This initialization should be called only once.");
-        this.preloads = Object.entries(this.plan.preload ?? {});
+        this.preloads = [];
+        this.shellBlobs = [];
+        for (let name of Object.getOwnPropertyNames(this.plan.preload)) {
+            let file = this.plan.preload[name];
+
+            const compressed = isCompressed(file);
+            if (compressed || globalThis.prefetchResources) {
+                let bytes = new Int8Array(read(file, "binary"));
+                if (compressed) {
+                    bytes = zlib.decompress(bytes);
+                }
+                // Add a `blob` prefix to the file name so that `getBinary`
+                // knows to look for a blob in JetStreamBlobs (setup by
+                // ShellScripts).
+                file = `blob://${name}`;
+                this.shellBlobs.push([file, bytes]);
+            }
+
+            this.preloads.push([name, file]);
+        }
     }
 
     scoreIdentifiers() {
@@ -1161,7 +1273,7 @@ class GroupedBenchmark extends Benchmark {
             await benchmark.prefetchResourcesForBrowser();
     }
 
-    async retryPrefetchResourcesForBrowser() {
+    async retryjForBrowser() {
         for (const benchmark of this.benchmarks)
             await benchmark.retryPrefetchResourcesForBrowser();
     }
@@ -1296,6 +1408,9 @@ class AsyncBenchmark extends DefaultBenchmark {
         } else {
             str += `
                 JetStream.getBinary = async function(path) {
+                    if (path.startsWith("blob://")) {
+                        return JetStreamBlobs[path];
+                    }
                     return new Int8Array(read(path, "binary"));
                 };
 
diff --git a/wasm-zlib.js b/wasm-zlib.js
@@ -0,0 +1,43 @@
+// zlib-based  utility for use in shells where CompressionStream and
+// DecompressionStream are not available.
+
+function module() {
+  'use strict';
+
+  let zlibPromise = null;
+  let zlibModule = null;
+
+  async function initialize() {
+    if (zlibPromise) {
+      zlibModule = await zlibPromise;
+      return zlibModule;
+    }
+    load('wasm/zlib/build/zlib.js');
+    zlibPromise = setupModule({
+      wasmBinary: new Int8Array(read('wasm/zlib/build/zlib.wasm', "binary")),
+    });
+    zlibModule = await zlibPromise;
+    return zlibModule;
+  }
+
+  function decompress(bytes) {
+    zlibModule.FS.writeFile('in', bytes);
+    const inputzStr = zlibModule.stringToNewUTF8('in');
+    const inputzoutStr = zlibModule.stringToNewUTF8('out');
+    if (zlibModule._decompressFile(inputzStr, inputzoutStr) !== 0) {
+      throw new Error();
+    }
+    const output = zlibModule.FS.readFile('out');
+    zlibModule._free(inputzStr);
+    zlibModule._free(inputzoutStr);
+    return output;
+  }
+
+  return {
+    initialize: initialize,
+    decompress: decompress,
+  };
+}
+
+globalThis.zlib = module();
+
