JSPIAsLibrary.cpp

/*
 * 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.
 */


#include <emscripten.h>
#include <emscripten/atomic.h>
#include <atomic>
#include <cstdio>

// The hard coded signatures are encoded as <return type><arg0><arg1>...
enum class Signature : uint32_t {
    Int32Int32,
    Int32Int32Int32,
    GenericDispatch,
};

enum Operation {
    Call = 1,
    Return = 2,
};

struct BridgeQueue {
    // When the main thread wants to post a message to the worker they simply increment this count by the appropriate count for the desired operation.
    // as long as the count hasn't changed the worker knows no new call has been queued.
    // FIXME: Maybe this could be a magic callee instead? Seems like both would be the same perf wise.
    std::atomic<uint32_t> conditionCount { 0 };
    Signature signature { 0 };
    void* callee { 0 };
    uint32_t running { 0 };
    uint64_t arguments[100];
};

#define ASSERT(condition) (void)(condition)
// #define ASSERT(condition) do { \
//     if (!(condition)) { \
//         __builtin_trap(); \
//     } \
// } while (false)

#define RELEASE_ASSERT(condition, newConditionCount, oldConditionCount, result) do { \
    if (!__builtin_expect(condition, 1)) { \
        printf(#condition " failed %d %d %d\n", newConditionCount, oldConditionCount, result); \
        __builtin_trap(); \
    } \
} while (false)

// FIXME: These could probably be the same queue since changing conditionCount would effectively be passing ownership.
static BridgeQueue* jsToWasmQueue;
static BridgeQueue* wasmToJSQueue;

extern "C" {

static constexpr uint32_t spinLimit = 1000;

void wasmMain() {
    uint32_t oldConditionCount = 0;
    uint32_t spinCount = 0;

    while (true) {
        // It seems like this should be able to just be an if rather than a while loop but we might get a spurious notify since the
        // JS thread could do the Atomics.add and we reply to the main thread before the notify happens on the JS side.
        while (jsToWasmQueue->conditionCount.load(std::memory_order_relaxed) == oldConditionCount) {
            if (spinCount++ < spinLimit)
                continue;

            // Nothing has called wasm for a while, park.
            ATOMICS_WAIT_RESULT_T result = emscripten_atomic_wait_u32(&jsToWasmQueue->conditionCount, oldConditionCount, ATOMICS_WAIT_DURATION_INFINITE);
            ASSERT(result != ATOMICS_WAIT_TIMED_OUT);
        }

        // Do an acquire load here to fence the above loop.
        uint32_t newConditionCount = jsToWasmQueue->conditionCount.load(std::memory_order_acquire);
        ASSERT(newConditionCount - oldConditionCount == 1);

        // We don't have to worry about any of this being modified by the main thread since it shouldn't touch this until either we "return" back to JS
        // or the C function we're going to call calls into some other JS.
        Signature signature = jsToWasmQueue->signature;

        switch (signature) {
        case Signature::Int32Int32: {
            int32_t argument0 = jsToWasmQueue->arguments[0];
            auto callee = reinterpret_cast<int32_t (*)(int32_t)>(jsToWasmQueue->callee);
            int32_t result = callee(argument0);

            oldConditionCount = jsToWasmQueue->conditionCount.load(std::memory_order_relaxed);
            wasmToJSQueue->signature = Signature::Int32Int32;
            wasmToJSQueue->arguments[0] = result;
            wasmToJSQueue->conditionCount.fetch_add(Return, std::memory_order_release);
            emscripten_atomic_notify(&wasmToJSQueue->conditionCount, 1);
            break;
        }

        case Signature::Int32Int32Int32:
        case Signature::GenericDispatch:
            // TODO: add this in ESM_JS, we might have to assume you don't use multi-return to JS so we can set oldConditionCount.
        default:
            printf("Got signature: %d conditionCounts: (new) %d (old) %d\n", static_cast<uint32_t>(signature), newConditionCount, oldConditionCount);
            __builtin_trap();
        }
    }
}

int32_t bridgeWasmToJSInt32Int32Int32(void* callee, int32_t argument0, int32_t argument1) {
    // FIXME: Deduplicate a lot of this.
    uint32_t oldConditionCount = jsToWasmQueue->conditionCount.load(std::memory_order_relaxed);

    wasmToJSQueue->signature = Signature::Int32Int32Int32;
    wasmToJSQueue->callee = callee;
    wasmToJSQueue->arguments[0] = argument0;
    wasmToJSQueue->arguments[1] = argument1;
    wasmToJSQueue->conditionCount.fetch_add(Call, std::memory_order_release);
    emscripten_atomic_notify(&wasmToJSQueue->conditionCount, 1);

    uint32_t spinCount = 0;

    while (jsToWasmQueue->conditionCount.load(std::memory_order_relaxed) == oldConditionCount) {
        if (spinCount++ < spinLimit)
            continue;

        // Nothing has called wasm for a while, park.
        ATOMICS_WAIT_RESULT_T result = emscripten_atomic_wait_u32(&jsToWasmQueue->conditionCount, oldConditionCount, ATOMICS_WAIT_DURATION_INFINITE);
        ASSERT(result != ATOMICS_WAIT_TIMED_OUT);
    }

    uint32_t newConditionCount = jsToWasmQueue->conditionCount.load(std::memory_order_acquire);
    if (newConditionCount - oldConditionCount == Return) {
        ASSERT(jsToWasmQueue->signature == Signature::Int32Int32Int32);
        return jsToWasmQueue->arguments[0];
    }

    printf("unreachable\n");
    // TODO: support calls back to wasm.
    return 0;
}

BridgeQueue* initJSToWasmQueue() {
    jsToWasmQueue = new BridgeQueue();
    return jsToWasmQueue;
}

BridgeQueue* initWasmToJSQueue() {
    wasmToJSQueue = new BridgeQueue();
    return wasmToJSQueue;
}

EM_JS(int32_t, workerPromiseAdd, (int32_t x, int32_t y), {
    return Promise.resolve(x+y);
});

int32_t workerPromiseFib(int32_t x) {
    if (x == 0)
        return 0;
    if (x == 1)
        return 1;

    // This would be generalized from:
    // return workerPromiseAdd(workerPromiseFib(x - 1), workerPromiseFib(x - 2));
    return bridgeWasmToJSInt32Int32Int32(reinterpret_cast<void*>(5), workerPromiseFib(x - 1), workerPromiseFib(x - 2));
}

#ifdef EM_ASYNC_JS

// promise an addition
EM_ASYNC_JS(int32_t, promiseAdd, (int32_t x, int32_t y), {
    return Promise.resolve(x+y);
});

int32_t promiseFib(int32_t x) {
    if (x == 0)
        return 0;
    if (x == 1)
        return 1;
    return promiseAdd(promiseFib(x - 1), promiseFib(x - 2));
}

#endif

} // extern "C"