Zlib compress all wasm files and decompress them during prefetch

.github/workflows/test.yml

@@ -39,6 +39,9 @@ jobs:
             - name: Install Node Packages
               run: npm ci
 
+            - name: Decompress compressed files
+              run: npm run decompress
+
             - name: Cache jsvu Binaries 
               uses: actions/cache@v4
               with:

.gitignore

@@ -8,3 +8,8 @@ node_modules
 
 # v8.log is generated by the d8-shell if profiling is enabled
 v8.log
+
+# Decompressed files go here
+RexBench/FlightPlanner/waypoints.js
+SeaMonster/inspector-json-payload.js
+wasm/argon2/build/argon2.wasm

JetStreamDriver.js

@@ -31,7 +31,17 @@ const defaultIterationCount = 120;
 const defaultWorstCaseCount = 4;
 
 if (!JetStreamParams.prefetchResources)
-    console.warn("Disabling resource prefetching!");
+    console.warn("Disabling resource prefetching! All compressed files must have been decompressed using `npm run decompress`");
+
+if (!isInBrowser && JetStreamParams.prefetchResources) {
+    // Use the wasm compiled zlib as a polyfill when decompression stream is
+    // not available in JS shells.
+    load("./wasm/zlib/shell.js");
+
+    // Load a polyfill for TextEncoder/TextDecoder in shells. Used when
+    // decompressing a prefetched resource and converting it to text.
+    load("./polyfills/fast-text-encoding/1.0.3/text.js");
+}
 
 // Used for the promise representing the current benchmark run.
 this.currentResolve = null;
@@ -138,6 +148,20 @@ function uiFriendlyDuration(time) {
     return `${time.toFixed(3)} ms`;
 }
 
+// Files can be zlib compressed to reduce the size of the JetStream source code.
+// 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");
+}
+
+function uncompressedName(name) {
+    console.assert(isCompressed(name));
+    return name.slice(0, -2);
+}
+
 // 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.
@@ -150,14 +174,28 @@ class ShellFileLoader {
     // share common code.
     load(url) {
         console.assert(!isInBrowser);
+
+        let compressed = isCompressed(url);
+        if (compressed && !JetStreamParams.prefetchResources) {
+            url = uncompressedName(url);
+        }
+
+        // If we aren't supposed to prefetch this then return code snippet that will load the url on-demand.
         if (!JetStreamParams.prefetchResources)
             return `load("${url}");`
 
         if (this.requests.has(url)) {
             return this.requests.get(url);
         }
 
-        const contents = readFile(url);
+        let contents;
+        if (compressed) {
+            const compressedBytes = new Int8Array(read(url, "binary"));
+            const decompressedBytes = zlib.decompress(compressedBytes);
+            contents = new TextDecoder().decode(decompressedBytes);
+        } else {
+            contents = readFile(url);
+        }
         this.requests.set(url, contents);
         return contents;
     }
@@ -209,10 +247,14 @@ class Driver {
             performance.mark("update-ui");
             benchmark.updateUIAfterRun();
 
-            if (isInBrowser && JetStreamParams.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;
@@ -361,6 +403,9 @@ class Driver {
 
     async prefetchResources() {
         if (!isInBrowser) {
+            if (JetStreamParams.prefetchResources) {
+                await zlib.initialize();
+            }
             for (const benchmark of this.benchmarks)
                 benchmark.prefetchResourcesForShell();
             return;
@@ -576,6 +621,11 @@ class Scripts {
 }
 
 class ShellScripts extends Scripts {
+    constructor() {
+        super();
+        this.prefetchedResources = Object.create(null);;
+    }
+
     run() {
         let globalObject;
         let realm;
@@ -602,13 +652,32 @@ class ShellScripts extends Scripts {
             currentReject
         };
 
+        // Pass the prefetched resources to the benchmark global.
+        if (JetStreamParams.prefetchResources) {
+            // Pass the 'TextDecoder' polyfill into the benchmark global. Don't
+            // use 'TextDecoder' as that will get picked up in the kotlin test
+            // without full support.
+            globalObject.ShellTextDecoder = TextDecoder;
+            // Store shellPrefetchedResources on ShellPrefetchedResources so that
+            // getBinary and getString can find them.
+            globalObject.ShellPrefetchedResources = this.prefetchedResources;
+        } else {
+            console.assert(Object.values(this.prefetchedResources).length === 0, "Unexpected prefetched resources");
+        }
+
         globalObject.performance ??= performance;
         for (const script of this.scripts)
             globalObject.loadString(script);
 
         return isD8 ? realm : globalObject;
     }
 
+    addPrefetchedResources(prefetchedResources) {
+        for (let [file, bytes] of Object.entries(prefetchedResources)) {
+            this.prefetchedResources[file] = bytes;
+        }
+    }
+
     add(text) {
         this.scripts.push(text);
     }
@@ -641,7 +710,6 @@ class BrowserScripts extends Scripts {
         return magicFrame;
     }
 
-
     add(text) {
         this.scripts.push(`<script>${text}</script>`);
     }
@@ -661,6 +729,7 @@ class Benchmark {
         this.allowUtf16 = !!plan.allowUtf16;
         this.scripts = null;
         this.preloads = null;
+        this.shellPrefetchedResources = null;
         this.results = [];
         this._state = BenchmarkState.READY;
     }
@@ -774,6 +843,9 @@ class Benchmark {
         if (!!this.plan.exposeBrowserTest)
             scripts.addBrowserTest();
 
+        if (this.shellPrefetchedResources) {
+            scripts.addPrefetchedResources(this.shellPrefetchedResources);
+        }
         if (this.plan.preload) {
             let preloadCode = "";
             for (let [ variableName, blobURLOrPath ] of this.preloads)
@@ -792,7 +864,7 @@ class Benchmark {
         } else {
             const cache = JetStream.blobDataCache;
             for (const file of this.plan.files) {
-                scripts.addWithURL(JetStreamParams.prefetchResources ? cache[file].blobURL : file);
+                scripts.addWithURL(cache[file].blobURL);
             }
         }
 
@@ -843,10 +915,19 @@ class Benchmark {
 
     async doLoadBlob(resource) {
         const blobData = JetStream.blobDataCache[resource];
+
+        const compressed = isCompressed(resource);
+        if (compressed && !JetStreamParams.prefetchResources) {
+            resource = uncompressedName(resource);
+        }
+
+        // If we aren't supposed to prefetch this then set the blobURL to just
+        // be the resource URL.
         if (!JetStreamParams.prefetchResources) {
             blobData.blobURL = resource;
             return blobData;
         }
+
         let response;
         let tries = 3;
         while (tries--) {
@@ -862,7 +943,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;
@@ -998,7 +1087,27 @@ 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.shellPrefetchedResources = Object.create(null);
+        if (!this.plan.preload) {
+            return;
+        }
+        for (let [name, file] of Object.entries(this.plan.preload)) {
+            const compressed = isCompressed(file);
+            if (compressed && !JetStreamParams.prefetchResources) {
+                file = uncompressedName(file);
+            }
+
+            if (JetStreamParams.prefetchResources) {
+                let bytes = new Int8Array(read(file, "binary"));
+                if (compressed) {
+                    bytes = zlib.decompress(bytes);
+                }
+                this.shellPrefetchedResources[file] = bytes;
+            }
+
+            this.preloads.push([name, file]);
+        }
     }
 
     scoreIdentifiers() {
@@ -1242,15 +1351,23 @@ class AsyncBenchmark extends DefaultBenchmark {
         } else {
             str += `
                 JetStream.getBinary = async function(path) {
+                    if ("ShellPrefetchedResources" in globalThis) {
+                        return ShellPrefetchedResources[path];
+                    }
                     return new Int8Array(read(path, "binary"));
                 };
 
                 JetStream.getString = async function(path) {
+                    if ("ShellPrefetchedResources" in globalThis) {
+                        return new ShellTextDecoder().decode(ShellPrefetchedResources[path]);
+                    }
                     return read(path);
                 };
 
                 JetStream.dynamicImport = async function(path) {
                     try {
+                        // TODO: this skips the prefetched resources, but I'm
+                        // not sure of a way around that.
                         return await import(path);
                     } catch (e) {
                         // In shells, relative imports require different paths, so try with and
@@ -1467,7 +1584,11 @@ class WasmLegacyBenchmark extends Benchmark {
             `;
         } else {
             str += `
-            Module[key] = new Int8Array(read(path, "binary"));
+            if (ShellPrefetchedResources) {
+                Module[key] = ShellPrefetchedResources[path];
+            } else {
+                Module[key] = new Int8Array(read(path, "binary"));
+            }
             if (andThen == doRun) {
                 globalObject.read = (...args) => {
                     console.log("should not be inside read: ", ...args);
@@ -1777,7 +1898,7 @@ let BENCHMARKS = [
         name: "FlightPlanner",
         files: [
             "./RexBench/FlightPlanner/airways.js",
-            "./RexBench/FlightPlanner/waypoints.js",
+            "./RexBench/FlightPlanner/waypoints.js.z",
             "./RexBench/FlightPlanner/flight_planner.js",
             "./RexBench/FlightPlanner/expectations.js",
             "./RexBench/FlightPlanner/benchmark.js",
@@ -1878,7 +1999,7 @@ let BENCHMARKS = [
     new DefaultBenchmark({
         name: "json-stringify-inspector",
         files: [
-            "./SeaMonster/inspector-json-payload.js",
+            "./SeaMonster/inspector-json-payload.js.z",
             "./SeaMonster/json-stringify-inspector.js",
         ],
         iterations: 20,
@@ -1888,7 +2009,7 @@ let BENCHMARKS = [
     new DefaultBenchmark({
         name: "json-parse-inspector",
         files: [
-            "./SeaMonster/inspector-json-payload.js",
+            "./SeaMonster/inspector-json-payload.js.z",
             "./SeaMonster/json-parse-inspector.js",
         ],
         iterations: 20,
@@ -2295,7 +2416,7 @@ let BENCHMARKS = [
             "./wasm/argon2/benchmark.js",
         ],
         preload: {
-            wasmBinary: "./wasm/argon2/build/argon2.wasm",
+            wasmBinary: "./wasm/argon2/build/argon2.wasm.z",
         },
         iterations: 30,
         worstCaseCount: 3,