bevy_remote_wasm 0.1.0

//! A plugin for the [Bevy Remote Protocol](https://docs.rs/bevy/latest/bevy/remote/) that exposes remote methods to
//! Wasm/JavaScript.
//!
//! ## Setup
//!
//! 1. Add `bevy_remote` to your dependencies, and `bevy_remote_wasm` only for `wasm` targets.
//!
//!    ```toml
//!    [dependencies]
//!    bevy = "0.18"
//!    bevy_remote = { version = "0.18", default-features = false }
//!
//!    [target.'cfg(target_family = "wasm")'.dependencies]
//!    bevy_remote_wasm = "0.1"
//!    ```
//!
//!    Don't depend on `bevy/bevy_remote` (`bevy = { version = "0.18", features = ["bevy_remote"] }`) because it
//!    enables the default `http` transport, which does not compile on Wasm target (until
//!    [this fix](https://github.com/bevyengine/bevy/pull/23367/) is merged and released).
//!
//!    This is why it is recommended to depend on the separate crate `bevy_remote` with no default features (so
//!    `bevy_remote/http` is not enabled).
//!
//! 2. Add [`RemotePlugin`](bevy_remote::RemotePlugin) to your [`App`], and add [`RemoteWasmPlugin`] only when compiling
//!    for Wasm.
//!
//!    ```rust
//!    use bevy::prelude::*;
//!    use bevy_remote::RemotePlugin;
//!
//!    let mut app = App::new();
//!    app.add_plugins((DefaultPlugins, RemotePlugin::default()));
//!    #[cfg(target_family = "wasm")]
//!    app.add_plugins(bevy_remote_wasm::RemoteWasmPlugin);
//!    app.run();
//!    ```
//!
//! 3. Build for the `wasm32-unknown-unknown` target, then generate the JS bindings with
//!    [`wasm-bindgen`](https://wasm-bindgen.github.io/wasm-bindgen/)
//!    directly (`wasm-bindgen --target web --out-dir <output_dir> target/wasm32-unknown-unknown/debug/<crate_name>.wasm`)
//!    or through a tool such as [`wasm-pack`](https://wasm-bindgen.github.io/wasm-pack/) or [Trunk](https://trunkrs.dev/).
//!
//! 4. In JavaScript, load the generated JS module, call `init()`, then await `getBridge()`:
//!
//!    ```js
//!    import init, { getBridge } from "/example.js";
//!
//!    await init();
//!    const bridge = await getBridge();
//!    const response = await bridge.main["rpc.discover"]();
//!    console.log(response.info.version);
//!    ```
//!
//!    The generated module path and `init` function depend on your build tool, but the initialization order does not:
//!    initialize the Wasm module first, then await the bridge.
//!
//!    `getBridge()` may be called as soon as the JS module loads, but it resolves only after Bevy publishes the bridge
//!    during app startup, which, depending on your setup, may take some time.
//!
//!    The `wasm-bindgen` output also includes TypeScript declarations for the root bridge (`BrpBridge`) and the default
//!    remote methods (`BuiltInBrpBridge`):
//!
//!    ```ts
//!    import type { BuiltInBrpBridge } from "./example.js";
//!
//!    const bridge = await getBridge() as BuiltInBrpBridge;
//!    const response = await bridge.main['world.query']({ data: { option: 'all' } });
//!    // Enjoy docs and autocompletion for built-in methods!
//!    ```
//!
//! 5. Serve the generated HTML, JS, and `.wasm` files together from a web server or bundler that supports Wasm imports.
//!
//! ## Wasm API
//!
//! The plugin builds a bridge object with app-scoped method maps based on the methods registered with
//! [`RemotePlugin`](bevy_remote::RemotePlugin), more specifically through the [`RemoteMethods`]
//! resource. Each time this resource is updated, the bridge is re-published with the new methods. The `getBridge()`
//! function must be called again to get the updated bridge.
//!
//! For Bevy `0.18`, the bridge exposes only the `main` app. Method keys inside that namespace use the BRP method names
//! directly, so call them with bracket notation such as `bridge.main["world.query"]` and
//! `bridge.main["world.list_components+watch"]`.
//!
//! - Instant methods (run once, return a result):
//!
//! ```js
//! // Call the instant method with params but without callback.
//! // Returns the result wrapped in a Promise.
//! const result = await bridge.main['world.query']({ data: { option: 'all' } });
//! console.log(result);
//!
//! // Call the instant method with params and a callback that will be called with the result.
//! // Returns `undefined` wrapped in a Promise.
//! await bridge.main['world.query']({ data: { option: 'all' } }, (result) => console.log(result));
//! ```
//!
//! - Watching methods (stream results):
//!
//! ```js
//! // Call the watching method with params and a callback that will be called on each result.
//! // Returns a closer function wrapped in a Promise.
//! const close = await bridge.main['world.list_components+watch']({ entity: 123 }, (result) => console.log(result));
//!
//! // Stop the stream.
//! close();
//! ```

#![cfg(any(doc, target_family = "wasm"))]

use async_channel::Sender;
use bevy_app::{App, Plugin, Update};
use bevy_ecs::{
    schedule::{
        IntoScheduleConfigs,
        common_conditions::{resource_exists, resource_exists_and_changed},
    },
    system::Res,
};
use bevy_remote::{BrpMessage, BrpResult, BrpSender, RemoteMethodSystemId, RemoteMethods};
use serde::Serialize;
use serde_json::Value;
use std::{
    cell::RefCell,
    collections::BTreeMap,
    task::{Poll, Waker},
};
use wasm_bindgen::prelude::*;
use wasm_bindgen_futures::{future_to_promise, spawn_local};

/// Add this plugin to your [`App`] to allow "remote" connections over Wasm to inspect and modify the
/// [`World`](bevy_ecs::world::World).
/// It requires the [`RemotePlugin`](bevy_remote::RemotePlugin).
pub struct RemoteWasmPlugin;

impl Plugin for RemoteWasmPlugin {
    fn build(&self, app: &mut App) {
        app.add_systems(
            Update,
            // `BrpSender` and `RemoteMethods` are inserted by `RemotePlugin`.
            update_wasm_bridge
                .run_if(resource_exists::<BrpSender>)
                .run_if(resource_exists_and_changed::<RemoteMethods>),
        );
    }
}

// The TypeScript definitions for the bridge, generated by `build.rs` and included in the wasm-bindgen output.
#[wasm_bindgen(typescript_custom_section)]
const TS_TYPES: &str = include_str!(concat!(env!("OUT_DIR"), "/ts_types.d.ts"));

/// Internal state for managing the bridge promise returned by `getBridge()`.
#[derive(Default)]
struct BridgeState {
    /// The bridge, once published by the plugin.
    bridge: Option<js_sys::Object>,
    /// Waker for the most recent pending `getBridge()` call.
    waker: Option<Waker>,
}

thread_local! {
    // JS may call `getBridge()` as soon as the module is instantiated, but the bridge is only built when Bevy runs its
    // systems. This thread-local bridges that timing gap.
    static BRIDGE_STATE: RefCell<BridgeState> = RefCell::new(BridgeState::default());
}

/// Get a promise which resolves into the bridge object once it is published by the plugin.
///
/// The returned object exposes a `main` property whose functions match all methods registered with
/// [`RemoteMethods`].
#[doc(hidden)]
#[wasm_bindgen(js_name = getBridge, skip_typescript)]
pub async fn get_bridge() -> js_sys::Object {
    if let Some(bridge) = BRIDGE_STATE.with(|state| state.borrow().bridge.clone()) {
        bridge
    } else {
        std::future::poll_fn(|cx| {
            BRIDGE_STATE.with_borrow_mut(|state| {
                if let Some(bridge) = state.bridge.clone() {
                    return Poll::Ready(bridge);
                }

                state.waker = Some(cx.waker().clone());
                Poll::Pending
            })
        })
        .await
    }
}

/// Wasm-facing method wrapper signature used by the BRP bridge.
///
/// Both instant and watching BRP methods take `(params?, callback?)` and return a Promise.
/// - Instant methods without a callback resolve to the result value.
/// - Instant methods with a callback resolve to `undefined`.
/// - Watching methods resolve to a closer function that can be called to stop the stream of updates
type WasmBridgeMethod = dyn Fn(Option<JsValue>, Option<js_sys::Function>) -> js_sys::Promise + 'static;

/// Wasm-facing closer function returned by watching BRP methods.
type WasmBridgeWatchCloser = dyn Fn() + 'static;

/// Internal map of methods for a single Bevy app.
struct BrpApp(BTreeMap<String, Closure<WasmBridgeMethod>>);

/// Internal bridge representation.
struct BrpBridge {
    main: BrpApp,
}

/// Publish the bridge to Wasm and wake any pending `getBridge()` call.
fn publish_bridge(bridge: BrpBridge) {
    if let Some(waker) = BRIDGE_STATE.with_borrow_mut(|state| {
        state.bridge = Some(bridge.into());
        state.waker.take()
    }) {
        waker.wake();
    }
}

/// Build the bridge object based on the registered remote methods.
fn update_wasm_bridge(brp_sender: Res<BrpSender>, remote_methods: Res<RemoteMethods>) {
    let main = BrpApp(
        remote_methods
            .methods()
            .into_iter()
            .filter_map(|name| remote_methods.get(&name).map(|id| (name, id)))
            .map(|(method_name, method_id)| {
                let function = build_function(&brp_sender, method_name.to_owned(), *method_id);
                (method_name, function)
            })
            .collect(),
    );

    publish_bridge(BrpBridge { main });
}

fn build_function(
    sender: &Sender<BrpMessage>,
    method_name: String,
    remote_method_system_id: RemoteMethodSystemId,
) -> Closure<WasmBridgeMethod> {
    let sender = sender.clone();
    Closure::<WasmBridgeMethod>::new(
        move |js_params: Option<JsValue>, callback: Option<js_sys::Function>| -> js_sys::Promise {
            let sender = sender.clone();
            let method = method_name.clone();
            future_to_promise(async move {
                match remote_method_system_id {
                    RemoteMethodSystemId::Instant(_) => {
                        build_instant_function(sender, method, js_params, callback).await
                    }
                    RemoteMethodSystemId::Watching(_) => {
                        build_watching_function(sender, method, js_params, callback).await
                    }
                }
            })
        },
    )
}

async fn build_instant_function(
    sender: Sender<BrpMessage>,
    method: String,
    js_params: Option<JsValue>,
    callback: Option<js_sys::Function>,
) -> Result<JsValue, JsValue> {
    let (result_tx, result_rx) = async_channel::bounded(1);

    let params = params_from_js(js_params);

    sender
        .send(BrpMessage {
            method,
            params,
            sender: result_tx,
        })
        .await
        .map_err(|e| js_sys::Error::new(&format!("Failed to send request: {e}")))?;

    let result = result_rx
        .recv()
        .await
        .map_err(|_| js_sys::Error::new("Channel closed unexpectedly"))?;
    let js_result = result_to_js(result)?;

    if let Some(callback) = callback {
        // If a callback was provided, call it with the result and resolve to undefined.
        callback.call1(&JsValue::NULL, &js_result)?;
        Ok(JsValue::UNDEFINED)
    } else {
        // If no callback, resolve the promise to the result value.
        Ok(js_result)
    }
}

async fn build_watching_function(
    sender: Sender<BrpMessage>,
    method: String,
    js_params: Option<JsValue>,
    callback: Option<js_sys::Function>,
) -> Result<JsValue, JsValue> {
    let (result_tx, result_rx) = async_channel::bounded::<BrpResult>(8);
    // A separate receiver handle for the closer to close the channel.
    let closer_rx = result_rx.clone();

    let params = params_from_js(js_params);

    sender
        .send(BrpMessage {
            method,
            params,
            sender: result_tx,
        })
        .await
        .map_err(|e| js_sys::Error::new(&format!("Failed to send request: {e}")))?;

    // Spawn the receive loop if a callback was provided.
    if let Some(callback) = callback {
        spawn_local(async move {
            while let Ok(result) = result_rx.recv().await {
                // The error is treated as a regular result to be passed to the callback.
                let arg = result_to_js(result).unwrap_or_else(std::convert::identity);
                let _ = callback.call1(&JsValue::NULL, &arg);
            }
        });
    }

    let closer = Closure::<WasmBridgeWatchCloser>::new(move || {
        // Closing the Receiver terminates the callback loop above and causes Bevy Remote to stop sending updates.
        closer_rx.close();
    });

    Ok(closer.into_js_value())
}

impl From<BrpApp> for js_sys::Object {
    fn from(app: BrpApp) -> Self {
        let object = js_sys::Object::new();

        for (method_name, function) in app.0 {
            js_sys::Reflect::set(&object, &JsValue::from_str(&method_name), &function.into_js_value()).unwrap_or(false);
        }

        object
    }
}

impl From<BrpBridge> for js_sys::Object {
    fn from(bridge: BrpBridge) -> Self {
        let object = js_sys::Object::new();
        let main = js_sys::Object::from(bridge.main);
        js_sys::Reflect::set(&object, &JsValue::from_str("main"), &main).unwrap_or(false);
        object
    }
}

fn params_from_js(params: Option<JsValue>) -> Option<Value> {
    serde_wasm_bindgen::from_value(params?).ok()
}

fn result_to_js(result: BrpResult) -> Result<JsValue, JsValue> {
    result
        .map(|value| {
            let serializer = serde_wasm_bindgen::Serializer::new().serialize_maps_as_objects(true);
            value.serialize(&serializer).unwrap_or(JsValue::UNDEFINED)
        })
        .map_err(|err| {
            let msg = format!("[{}] {}", err.code, err.message);
            js_sys::Error::new(&msg).into()
        })
}

#[cfg(test)]
mod tests {
    use super::*;
    use bevy::prelude::*;
    use bevy_remote::RemotePlugin;
    use wasm_bindgen_test::wasm_bindgen_test;

    #[wasm_bindgen_test(async)]
    async fn get_bridge_with_default_methods() {
        let mut app = App::new();
        app.add_plugins(MinimalPlugins);
        app.add_plugins((RemotePlugin::default(), RemoteWasmPlugin));
        app.update();

        let bridge = get_bridge().await;
        assert_eq!(js_sys::Object::keys(&bridge).length(), 1);

        let main = js_sys::Reflect::get(&bridge, &JsValue::from_str("main")).unwrap();
        let main = main.dyn_into::<js_sys::Object>().unwrap();

        // There are more than 10 default remote methods, actual number doesn't matter.
        assert!(js_sys::Object::keys(&main).length() > 10);

        let method = js_sys::Reflect::get(&main, &JsValue::from_str("rpc.discover")).unwrap();
        let method = method.dyn_into::<js_sys::Function>().unwrap();
        let promise = method
            .call0(&JsValue::NULL)
            .unwrap()
            .dyn_into::<js_sys::Promise>()
            .unwrap();

        let result = wasm_bindgen_futures::JsFuture::from(promise).await.unwrap();
        let info = js_sys::Reflect::get(&result, &JsValue::from_str("info")).unwrap();
        let version = js_sys::Reflect::get(&info, &JsValue::from_str("version")).unwrap();
        assert!(
            !version.is_undefined(),
            "Expected version to be present, got {:?}",
            version
        );
    }

    #[wasm_bindgen_test(async)]
    async fn get_bridge_without_methods() {
        let mut app = App::new();
        app.add_plugins(MinimalPlugins);
        app.add_plugins((RemotePlugin::default(), RemoteWasmPlugin));
        app.insert_resource(RemoteMethods::new()); // Remove the default remote methods.
        app.update();

        let bridge = get_bridge().await;
        assert_eq!(js_sys::Object::keys(&bridge).length(), 1);

        let main = js_sys::Reflect::get(&bridge, &JsValue::from_str("main")).unwrap();
        let main = main.dyn_into::<js_sys::Object>().unwrap();

        assert_eq!(js_sys::Object::keys(&main).length(), 0);
    }

    #[wasm_bindgen_test(async)]
    async fn get_bridge_with_updated_methods() {
        let mut app = App::new();
        app.add_plugins(MinimalPlugins);
        app.add_plugins((RemotePlugin::default(), RemoteWasmPlugin));
        app.update();

        let default_bridge = get_bridge().await;
        let default_main = js_sys::Reflect::get(&default_bridge, &JsValue::from_str("main")).unwrap();
        let default_main = default_main.dyn_into::<js_sys::Object>().unwrap();

        let mut updated_methods = RemoteMethods::new();
        let my_handler = |In(_params): In<Option<Value>>| -> BrpResult { Ok(Value::Null) };
        let my_handler_id = app.register_system(my_handler);
        updated_methods.insert("my_method", RemoteMethodSystemId::Instant(my_handler_id));
        app.insert_resource(updated_methods);
        app.update();

        let updated_bridge = get_bridge().await;
        let updated_main = js_sys::Reflect::get(&updated_bridge, &JsValue::from_str("main")).unwrap();
        let updated_main = updated_main.dyn_into::<js_sys::Object>().unwrap();

        // Default bridge object is still valid and has the default remote methods.
        assert_eq!(js_sys::Object::keys(&default_bridge).length(), 1);
        assert!(js_sys::Object::keys(&default_main).length() > 10);
        // Updated bridge should only have "my_method".
        assert_eq!(js_sys::Object::keys(&updated_bridge).length(), 1);
        assert_eq!(js_sys::Object::keys(&updated_main).length(), 1);
        assert!(js_sys::Reflect::has(&updated_main, &JsValue::from_str("my_method")).unwrap());
    }
}