MyJSPIAsLibraryGlue.js

/*
 * Copyright (C) 2024 Apple Inc. All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY APPLE INC. ``AS IS'' AND ANY
 * EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
 * PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL APPLE INC. OR
 * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
 * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
 * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
 * OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */


var resolveWasmReady;
var wasmReady = new Promise(resolve => resolveWasmReady = resolve);

const ConditionCall = 1;
const ConditionReturn = 2;

const SignatureInt32Int32 = 0;
const SignatureInt32Int32Int32 = 1;
const SignatureGenericDispatch = 2;

let resolvePendingPromise = null;

Atomics.pause ??= function() { }

async function wasmToJSMain(instance) {
    var oldConditionCount = 0;
    var spinCount = 0;
    const spinLimit = 10000;
    var newConditionCount;
    while (true) {
        while ((newConditionCount = Atomics.load(Module.HEAP32, wasmToJSQueue_conditionCount_32)) == oldConditionCount) {
            if (spinCount++ < spinLimit && resolvePendingPromise) {
                Atomics.pause();
                continue;
            }

            spinCount = 0;
            var asyncResult = Atomics.waitAsync(Module.HEAP32, wasmToJSQueue_conditionCount_32, oldConditionCount);
            if (asyncResult.async)
                await asyncResult.value;
        }

        var signature = Module.HEAP32[wasmToJSQueue_signature_32];
        // TODO Handle overflow.
        var conditionCode = newConditionCount - oldConditionCount;
        if (conditionCode == ConditionCall) {
            var calleeIndex = Module.HEAP32[wasmToJSQueue_callee_32];

            switch (signature) {
            case SignatureInt32Int32Int32:
                let argument0 = Module.HEAP32[wasmToJSQueue_arguments_64 * 2];
                let argument1 = Module.HEAP32[wasmToJSQueue_arguments_64 * 2 + 2];
                var result = await wasmTableMirror[calleeIndex](argument0, argument1);
                oldConditionCount = Module.HEAP32[wasmToJSQueue_conditionCount_32];

                Module.HEAP32[jsToWasmQueue_signature_32] = signature;
                Module.HEAP32[jsToWasmQueue_arguments_64 * 2] = result;
                Atomics.add(Module.HEAP32, jsToWasmQueue_conditionCount_32, ConditionReturn);
                Atomics.notify(Module.HEAP32, jsToWasmQueue_conditionCount_32, 1);
                continue;
            default:
                break;
            }
            console.assert(false, "TODO");
            continue;
        } else if (conditionCode == ConditionReturn) {
            var result = Module.HEAP32[wasmToJSQueue_arguments_64 * 2];
            oldConditionCount = Module.HEAP32[wasmToJSQueue_conditionCount_32];
            resolvePendingPromise(result);
            resolvePendingPromise = null;
        } else {
            console.assert(false, "Wrong condition code: " + conditionCode);
            throw new Error("Wrong condition code: " + conditionCode);
        }
    }
}

async function buildGlue(imports, instanceReady) {
    let wasmBinaryFile = findWasmBinary();
    let instantiateResult = await new Promise((resolve) => {
        instantiateAsync(undefined, wasmBinaryFile, imports, resolve);
    });
    let instance = instantiateResult.instance;
    instanceReady(instance, instantiateResult.module);
    console.assert(Module.calledRun);

    globalThis.jsToWasmQueue = instantiateResult.instance.exports.initJSToWasmQueue();
    globalThis.wasmToJSQueue = instantiateResult.instance.exports.initWasmToJSQueue();

    globalThis.wasmToJSQueue_conditionCount_32 = wasmToJSQueue / 4;
    globalThis.wasmToJSQueue_signature_32 = wasmToJSQueue_conditionCount_32 + 1;
    globalThis.wasmToJSQueue_callee_32 = wasmToJSQueue_signature_32 + 1;
    globalThis.wasmToJSQueue_arguments_64 = wasmToJSQueue / 8 + 2;

    globalThis.jsToWasmQueue_conditionCount_32 = jsToWasmQueue / 4;
    globalThis.jsToWasmQueue_signature_32 = jsToWasmQueue_conditionCount_32 + 1;
    globalThis.jsToWasmQueue_callee_32 = jsToWasmQueue_signature_32 + 1;
    globalThis.jsToWasmQueue_arguments_32 = jsToWasmQueue / 4 + 4;
    globalThis.jsToWasmQueue_arguments_64 = jsToWasmQueue / 8 + 2;

    let worker = new Worker('MyJSPIAsLibraryWorkerGlue.js');
    worker.postMessage([instantiateResult.module, wasmMemory, jsToWasmQueue, wasmToJSQueue, { workerPromiseAdd: 0 }]);

    await new Promise((resolve) => worker.onmessage = resolve);

    // This part should be generalized and exports are not writable for some reason.
    let workerPromiseFibTableIndex = addFunction(instance.exports.workerPromiseFib);
    function workerPromiseFib(arg1) {
        return new Promise((resolve) => {
            // TODO: solve reentrancy
            if (resolvePendingPromise)
                throw new Error("called function with pending promise");

            resolvePendingPromise = resolve;
            Module.HEAP32[jsToWasmQueue_callee_32] = workerPromiseFibTableIndex;
            Module.HEAP32[jsToWasmQueue_signature_32] = SignatureInt32Int32;
            Module.HEAP32[jsToWasmQueue_arguments_64 * 2] = arg1;
            let old = Atomics.add(Module.HEAP32, jsToWasmQueue_conditionCount_32, ConditionCall);
            Atomics.notify(Module.HEAP32, jsToWasmQueue_conditionCount_32, 1);
        });
    }

    // We can't actually wrap this with a wasm function because that will do incorrect boxing.
    let workerPromiseAddIndex = getEmptyTableSlot();
    wasmTableMirror[workerPromiseAddIndex] = function(x,y) { return Promise.resolve(x+y) };
    console.assert(workerPromiseAddIndex == 5); // TODO: hardcoded in C++ for now.

    wasmToJSMain(instance);
    resolveWasmReady({ instance, workerPromiseFib });
}

var Module = {
    instantiateWasm: buildGlue,
};

(async function() {
    let { instance, workerPromiseFib } = await wasmReady;

    let times = new Array(1000);
    for (let i = 0; i < 1000; ++i) {
        let start = performance.now();
        let result = await workerPromiseFib(15);
        let end = performance.now();
        console.assert(result == 610);
        times[i] = (end - start);
    }


    console.log("library: ", times);
    console.log("library average: ", times.reduce((x,y) => x+y, 0) / 1000);

    if (wasmExports.promiseFib) {
        times = new Array(1000);
            for (let i = 0; i < 1000; ++i) {
            let start = performance.now();
            let result = await wasmExports.promiseFib(15);
            let end = performance.now();
            console.assert(result == 610);
            times[i] = (end - start);
        }

        console.log("JSPI: ", times);
        console.log("JSPI average: ", times.reduce((x,y) => x+y, 0) / 1000);
    }

})();