Use the WebAssembly implementation unless requested.

this resulted in a massive performance boost, ranging from
20% (node.js) to 50% (Chrome) to more than 2x (Firefox).

index.d.ts

@@ -108,6 +108,7 @@ export interface CompressOptions extends AnsOptions {
 }
 
 export interface DefaultModelCompressOptions extends CompressOptions, DefaultModelOptions {
+    disableWasm?: boolean;
 }
 
 export interface Output {

index.mjs

@@ -18,6 +18,8 @@ import {
 
 import { estimateDeflatedSize } from './deflate.mjs';
 
+import { getContextItemShift, makeDefaultModelRunner } from './wasm.mjs';
+
 //------------------------------------------------------------------------------
 
 // returns clamp(0, floor(log2(x/y)), 31) where x and y are integers.
@@ -74,6 +76,7 @@ export class ResourcePool {
     constructor() {
         // arrayBuffers.get(size) is an array of ArrayBuffer of given size
         this.arrayBuffers = new Map();
+        this.wasmRunners = new Map();
     }
 
     allocate(parent, size) {
@@ -93,6 +96,13 @@ export class ResourcePool {
         }
         available.push(buf);
     }
+
+    wasmDefaultModelRunner(contextItemShift) {
+        if (!this.wasmRunners.get(contextItemShift)) {
+            this.wasmRunners.set(contextItemShift, makeDefaultModelRunner(contextItemShift));
+        }
+        return this.wasmRunners.get(contextItemShift);
+    }
 }
 
 export const ArrayBufferPool = ResourcePool;
@@ -535,10 +545,52 @@ export const compressWithModel = (input, model, options) => {
 };
 
 export const compressWithDefaultModel = (input, options) => {
-    const model = new DefaultModel(options);
-    const ret = compressWithModel(input, model, options);
-    ret.quotesSeen = model.quotesSeen;
-    return ret;
+    // if the model is _exactly_ a DefaultModel and no fancy options are in use,
+    // we have a faster implementation using JIT-compiled WebAssembly instances.
+    // we only use it when we have a guarantee that the wasm instance can be cached.
+    let runDefaultModel;
+    const resourcePool = options.resourcePool || options.arrayBufferPool;
+    if (
+        resourcePool &&
+        (options.preset || []).length === 0 &&
+        !options.disableWasm &&
+        !options.calculateByteEntropy &&
+        options.sparseSelectors.length <= 64 &&
+        options.sparseSelectors.every(sel => sel < 0x8000)
+    ) {
+        const contextItemShift = getContextItemShift(options);
+        const resourcePool = options.resourcePool || options.arrayBufferPool;
+        try {
+            runDefaultModel = resourcePool.wasmDefaultModelRunner(contextItemShift);
+        } catch (e) {
+            // WebAssembly is probably not supported, fall back to the pure JS impl
+        }
+    }
+
+    if (runDefaultModel) {
+        const { predictions, quotesSeen } = runDefaultModel(input, options);
+
+        const { inBits, outBits } = options;
+        const encoder = new AnsEncoder(options);
+        for (let offset = 0, bitOffset = 0; offset < input.length; ++offset) {
+            const code = input[offset];
+            for (let i = inBits - 1; i >= 0; --i) {
+                const bit = (code >> i) & 1;
+                const prob = predictions[bitOffset++];
+                encoder.writeBit(bit, prob);
+            }
+        }
+        const { state, buf } = encoder.finish();
+
+        const bufLengthInBytes = Math.ceil(buf.length * (outBits < 0 ? Math.log2(-outBits) : outBits) / 8);
+        const inputLength = input.length;
+        return { state, buf, inputLength, bufLengthInBytes, quotesSeen };
+    } else {
+        const model = new DefaultModel(options);
+        const ret = compressWithModel(input, model, options);
+        ret.quotesSeen = model.quotesSeen;
+        return ret;
+    }
 };
 
 export const decompressWithModel = ({ state, buf, inputLength }, model, options) => {
@@ -642,6 +694,7 @@ export class Packer {
             resourcePool: options.resourcePool || options.arrayBufferPool,
             numAbbreviations: typeof options.numAbbreviations === 'number' ? options.numAbbreviations : 64,
             allowFreeVars: options.allowFreeVars,
+            disableWasm: options.disableWasm,
         };
 
         this.inputsByType = {};

test.mjs

@@ -295,10 +295,16 @@ function testDefaultModel(title, modelQuotes) {
         t.deepEqual(model.quotesSeen, new Set(modelQuotes ? [39, 96] : []));
 
         // test compressWithDefaultModel as well
-        const compressed2 = compressWithDefaultModel(input, options);
-        t.deepEqual(compressed2.state, compressed.state);
-        t.deepEqual(compressed2.buf, compressed.buf);
-        t.deepEqual(compressed2.quotesSeen, model.quotesSeen);
+        for (const disableWasm of [true, false]) {
+            // the wasm result should not only decompress correctly, but also
+            // be identical to the non-wasm result so that the compression
+            // ratio is never affected.
+            options.disableWasm = disableWasm;
+            const compressed2 = compressWithDefaultModel(input, options);
+            t.deepEqual(compressed2.state, compressed.state);
+            t.deepEqual(compressed2.buf, compressed.buf);
+            t.deepEqual(compressed2.quotesSeen, model.quotesSeen);
+        }
     });
 }
 
@@ -333,6 +339,32 @@ test('DirectContextModel.confirmations', t => {
     }
 });
 
+test('DefaultModel.confirmations', t => {
+    // same as above, but tests the wasm version
+    const resourcePool = new ResourcePool();
+    const options = {
+        inBits: 8,
+        outBits: 8,
+        precision: 8, // additionally required to trigger the mark overflow within 256 tries
+        contextBits: 5, // 32 elements
+        modelMaxCount: 63,
+        sparseSelectors: [0, 1, 2, 3, 4],
+        modelRecipBaseCount: 20,
+        recipLearningRate: 256,
+        modelQuotes: false,
+        resourcePool,
+    };
+
+    for (const size of [1, 10]) {
+        for (let i = 0; i < 1000; ++i) {
+            const input = [...crypto.randomBytes(size)];
+            const compressed = compressWithDefaultModel(input, options);
+            const decompressed = decompressWithModel(compressed, new DefaultModel(options), options);
+            t.deepEqual(decompressed, input);
+        }
+    }
+});
+
 //------------------------------------------------------------------------------
 
 test('prepareJs without abbrs', t => {

tools/compress-demo.mjs

@@ -17,12 +17,14 @@ if (!m) throw 'demo.html does not match the expected template';
 const [, preamble, style, body, script] = m;
 
 const roadrollerCode = fs.readFileSync(await resolve('../index.mjs'), { encoding: 'utf-8' });
+const wasmCode = fs.readFileSync(await resolve('../wasm.mjs'), { encoding: 'utf-8' });
 const deflateCode = fs.readFileSync(await resolve('../deflate.mjs'), { encoding: 'utf-8' });
 const jsTokensCode = fs.readFileSync(await resolve('../js-tokens.mjs'), { encoding: 'utf-8' });
 
 const { code, vars } = await minifyJs(
     stripModule(deflateCode) +
     stripModule(jsTokensCode) +
+    stripModule(wasmCode) +
     stripModule(roadrollerCode) +
     stripModule(script, true)
 );