bevy_subworlds 0.1.0

//! Threaded multi-world management and signal transmission for bevy. 
//! 
//! A SubWorld is a Component that stores a JoinHandle to a thread running an App. Supported
//! operations include sending an exit signal, checking if the SubWorld has completed its startup 
//! schedule, and extracting data with pollable extract systems.
//! 
//! **Extraction** is when a system is sent to a SubWorld to be ran. The output of this system can 
//! be polled with an `ExtractHandle`, similar to an async task or a bevy task. 
//! 
//! **Signals** are _either_ components or resources. When a Signal is a Resource, 
//! it sends or receives signals from a parent world. When a Signal is a Component, 
//! it sends and receives signals to/from a subworld. Signals receivers/transmitters 
//! are automatically inserted onto the SubWorld as components and within the SubWorld
//! as Resources when `SubWorldBuilder::with_signal` is called.
//! 
//! # Usage
//! This simple example spawns a SubWorld and ping/pongs a signal back and forth. It also
//! demonstrates extraction by reading and mutating a resource in the SubWorld.
//! ```
//! use bevy::prelude::*;
//!
//! fn main() {
//!     App::new()
//!         .add_plugins(MinimalPlugins)
//!         .add_systems(Startup, spawn_subworld)
//!         .add_systems(Update, (
//!             send_pings,
//!             extract_data,
//!         ))
//!         .run();
//! }
//!
//! // Your custom signal type
//! enum MySignal {
//!     Ping,
//!     Pong
//! }
//!
//! #[derive(Resource)]
//! struct GameData {
//!     num: u32,
//! }
//!
//! fn spawn_subworld(
//!     world_id: WorldId,
//!     mut commands: Commands,
//! ) {
//!     SubWorld::build(world_id, &mut commands)
//!         .with_signal::<MySignal>() // < register your signal type
//!         .start(move |app| {
//!             app
//!                 .insert_resource(GameData { num: 1 })
//!                 .add_systems(Update, respond_to_pings);
//!         });
//! }
//!
//! /// Extract and increment the GameState's `num` variable from each system.
//! fn extract_data(
//!     mut query: Query<(&SubWorld, Entity, Option<&mut ExtractHandle<u32>>)>,
//!     mut commands: Commands,
//! ) {
//!     for (subworld, entity, handle) in &mut query {
//!         if let Some(mut handle) = handle {
//!             if let Ok(num) = handle.poll() {
//!                 log::info!("Extracted Number: '{num}'. Time: '{:?}'", handle.duration().unwrap());
//!                 commands.entity(entity).remove::<ExtractHandle<u32>>();
//! 
//!                 if num > 100 {
//!                     subworld.exit();
//!                 }
//!             }
//!         } else {
//!             let handle: ExtractHandle<u32> = subworld.extract(
//!                 move |mut state: ResMut<GameData>| {
//!                     state.num += 1;
//!                     state.num
//!                 }
//!             );
//!
//!             commands.entity(entity).insert(handle);
//!         }
//!     }
//! }
//!
//! /// Send pings to Subworlds
//! fn send_pings(
//!     mut query: Query<(&mut SignalRx<MySignal>, &SignalTx<MySignal>), With<SubWorld>>,
//! ) {
//!     for (mut rx, tx) in &mut query {
//!         for signal in &mut rx {
//!             match signal {
//!                 MySignal::Ping => log::info!("Ping Received!"),
//!                 MySignal::Pong => {
//!                     tx.send(MySignal::Ping)
//!                 }
//!             }
//!         }
//!
//!         tx.send(MySignal::Ping);
//!     }
//! }
//!
//! /// Respond to pings as a SubWorld
//! fn respond_to_pings(
//!     mut signal_rx: ResMut<SignalRx<MySignal>>,
//!     signal_tx: Res<SignalTx<MySignal>>,
//! ) {
//!     for signal in &mut signal_rx {
//!         match signal {
//!             MySignal::Ping => {
//!                 let _ = signal_tx.send(MySignal::Pong);
//!             },
//!             MySignal::Pong => {}
//!         }
//!     }
//! }
//! ```

use std::{mem, sync::{atomic::{AtomicBool, Ordering}, Arc}, thread::JoinHandle};

use parking_lot::Mutex;
use bevy::{ecs::{system::SystemState, world::WorldId}, log, prelude::*};
use crossbeam::channel::{Receiver, Sender, TryRecvError};
use web_time::{Instant, Duration};

/// A Handle to a bevy App running in its own thread. 
#[derive(Component)]
pub struct SubWorld {
    handle: Option<JoinHandle<AppExit>>,
    started: Arc<AtomicBool>,
    tx: Sender<SuperSignal>,
}

impl SubWorld {
    /// Construct a new SubWorld
    pub fn build<'c>(parent: WorldId, cmd: &'c mut Commands) -> SubWorldBuilder<'c> {
        SubWorldBuilder {
            parent,
            channel_cap: 16,
            resources: Vec::new(),
            commands: cmd.spawn_empty()
        }
    }

    /// Check if the SubWorld has exited. 
    pub fn is_exited(&self) -> bool {
        self.handle.as_ref().is_none_or(|handle| handle.is_finished())
    }

    /// Check if the SubWorld's startup sequence has ran.
    pub fn is_started(&self) -> bool {
        self.started.load(Ordering::Relaxed)
    }

    /// Take the raw handle directly by replacing the inner JoinHandle with None.
    /// This is provided so you can .join() and handle the error/exit code manually.
    pub fn take_handle(&mut self) -> Option<JoinHandle<AppExit>> {
        self.handle.take()
    }

    /// Dispatches an `AppExit` event in the SubWorld.
    /// This won't happen immediately, but it will the next
    /// time the signal receiver is read. (automatically done)
    /// 
    /// A SubWorld can also be exited by calling its `drop` destructor.
    pub fn exit(&self) {
        let _ = self.tx.send(SuperSignal::Exit);
    }

    /// Execute a function in the SubWorld with &mut Commands as input.
    /// This allows you to spawn entities or insert resources.
    pub fn run_commands<F>(&self, f: F)
    where
        F: FnOnce(&mut Commands) + Send + Sync + 'static
    {
        let _ = self.tx.send(SuperSignal::Commands(Box::new(f)));
    }

    /// Send a system to the SubWorld to be executed.
    /// 
    /// Returns a handle that can be polled for the return value.
    pub fn extract<S, Out, M>(&self, system: S) -> ExtractHandle<Out>
    where
        S: IntoSystem<(), Out, M> + Send + Sync + 'static,
        Out: Send + Sync + 'static,
    {
        let shared = Arc::new(Shared::<Out>::new());
        let shared2 = shared.clone();
        let _ = self.tx.send(SuperSignal::WorldFn(Box::new(
            move |world: &mut World| {
                match world.run_system_cached(system) {
                    Ok(out) => shared2.set(Ok(out)),
                    Err(e) => {
                        log::error!("Attempted to run a SubTask in a SubWorld, but the system failed to register with error: '{e}'.");
                        shared2.set(Err(ExtractError::RegistrationErr))
                    }
                }
            }
        )));

        ExtractHandle {
            shared,
            send_time: Instant::now(),
            recv_time: None,
        }
    }

    /// Send a system to the SubWorld to be executed, with a provided input.
    /// 
    /// Returns a handle that can be polled for the return value.
    pub fn extract_with_input<S, In, Out, M>(&self, input: In, system: S) -> ExtractHandle<Out>
    where
        S: IntoSystem<In, Out, M> + Send + Sync + 'static,
        In: SystemInput<Inner<'static> = In> + Send + Sync + 'static,
        Out: Send + Sync + 'static,
    {
        let shared = Arc::new(Shared::<Out>::new());
        let shared2 = shared.clone();
        let _ = self.tx.send(SuperSignal::WorldFn(Box::new(
            move |world: &mut World| {
                match world.run_system_cached_with(system, input) {
                    Ok(out) => shared2.set(Ok(out)),
                    Err(e) => {
                        log::error!("Attempted to run a SubTask in a SubWorld, but the system failed to register with error: '{e}'.");
                        shared2.set(Err(ExtractError::RegistrationErr))
                    }
                }
            }
        )));

        ExtractHandle {
            shared,
            send_time: Instant::now(),
            recv_time: None,
        }
    }
}

impl Drop for SubWorld {
    fn drop(&mut self) {
        let _ = self.tx.try_send(SuperSignal::Exit);
    }
}

pub struct SubWorldBuilder<'c> {
    parent: WorldId,
    channel_cap: usize,
    resources: Vec<Box<dyn FnOnce(&mut App) -> &mut App + Send + Sync + 'static>>,
    commands: EntityCommands<'c>,
}

impl<'c> SubWorldBuilder<'c> {
    /// Set the capacity of the signal channels.
    pub fn with_channel_cap(mut self, cap: usize) -> Self {
        self.channel_cap = cap;
        self
    }

    /// The SubWorld will have a `SignalRx<T>` resource, and the
    /// entity this SubWorld is attached to will have a `SignalTx<T>` 
    /// component.
    pub fn with_signal_rx<T>(mut self) -> Self 
    where
        T: Send + Sync + 'static
    {
        let (tx, rx) = crossbeam::channel::bounded::<T>(self.channel_cap);
        self.commands.insert(SignalTx { tx });
        self.resources.push(Box::new(move |app| app.insert_resource(SignalRx { rx })));
        self
    }

    /// The SubWorld will have a `SignalTx<T>` resource, and the 
    /// entity this SubWorld is attached to will have a `SignalRx<T>` 
    /// component.
    pub fn with_signal_tx<T>(mut self) -> Self 
    where
        T: Send + Sync + 'static,
    {
        let (tx, rx) = crossbeam::channel::bounded::<T>(self.channel_cap);
        self.commands.insert(SignalRx { rx });
        self.resources.push(Box::new(move |app| app.insert_resource(SignalTx { tx })));
        self
    }

    /// The SubWorld will have a `SignalTx<T>` and `SignalRx<T>` resource, 
    /// and the entity this SubWorld is attached to will have a `SignalTx<T>` 
    /// and `SignalRx<T>` component.
    pub fn with_signal<T>(self) -> Self 
    where
        T: Send + Sync + 'static
    {
        self.with_signal_rx::<T>()
            .with_signal_tx::<T>()
    }

    /// Start the subworld with the App returned by the closure. This closure
    /// will run in the spawned thread, so it won't cause the calling system
    /// to be unbearably slow. 
    pub fn start<F>(mut self, f: F) -> EntityCommands<'c>
    where
        F: FnOnce(&mut App) + Send + Sync + 'static
    {
        let mut resources = std::mem::take(&mut self.resources);
        let (tx, rx) = crossbeam::channel::bounded::<SuperSignal>(4);
        let start = Arc::new(AtomicBool::new(false));
        let started = start.clone();
        let handle = std::thread::spawn(move || {
            let mut app = App::new();
            (f)(&mut app);
            while let Some(res) = resources.pop() {
                (res)(&mut app);
            }
            app
                .add_systems(PreUpdate, process_super_signals)
                .insert_resource(SuperWorld {
                    parent: self.parent,
                    started,
                    rx,
                })
                .run()
        });

        self.commands.insert(
            SubWorld {
                handle: Some(handle),
                started: start,
                tx
            }
        );
        self.commands
    }
}

/// A signal transmitter between a SubWorld and SuperWorld.
/// When used as a Resource, these are signals to the SuperWorld it belongs to.
/// When used as a Component, these are signals to a SubWorld it is attached to.
#[derive(Component, Resource)]
#[component(immutable)]
pub struct SignalTx<T> 
where
    T: Send + Sync + 'static
{
    tx: Sender<T>,
}

impl<T> SignalTx<T> 
where
    T: Send + Sync + 'static,
{
    pub fn send(&self, signal: T) {
        let _ = self.tx.try_send(signal);
    }
}

/// A signal receiver between a SubWorld and SuperWorld.
/// When used as a Resource, these are signals from the SuperWorld it belongs to.
/// When used as a Component, these are signals from the SubWorld it is attached to.
#[derive(Component, Resource)]
pub struct SignalRx<T> 
where
    T: Send + Sync + 'static,
{
    rx: Receiver<T>,
}

impl<T> Iterator for &mut SignalRx<T>
where
    T: Send + Sync + 'static
{
    type Item = T;

    fn next(&mut self) -> Option<Self::Item> {
        self.rx.try_recv().ok()
    }
}

/// A handle to the parent World of a SubWorld.
#[derive(Resource)]
pub struct SuperWorld {
    parent: WorldId,
    started: Arc<AtomicBool>,
    rx: Receiver<SuperSignal>,
}

impl SuperWorld {
    pub fn parent(&self) -> WorldId {
        self.parent
    }
}

enum SuperSignal {
    Commands(Box<dyn FnOnce(&mut Commands) + Send + Sync + 'static>),
    WorldFn(Box<dyn FnOnce(&mut World) + Send + Sync + 'static>),
    Exit,
}

fn process_super_signals(
    world: &mut World,
    has_ran: Local<bool>,
    params: &mut SystemState<(
        ResMut<SuperWorld>,
        Commands,
        EventWriter<AppExit>,
    )>,
) {
    let mut systems = Vec::new();

    {
        let (signals, mut commands, mut exit_ev) = params.get_mut(world);

        if !*has_ran {
            signals.started.store(true, Ordering::Relaxed);
        }

        loop {
            match signals.rx.try_recv() {
                Ok(signal) => {
                    match signal {
                        SuperSignal::Exit => {
                            exit_ev.write(AppExit::Success);
                        },
                        SuperSignal::Commands(f) => {
                            (f)(&mut commands)
                        },
                        SuperSignal::WorldFn(f) => {
                            systems.push(f);
                        }
                    }
                },
                Err(e) => {
                    match e {
                        TryRecvError::Empty => break,
                        TryRecvError::Disconnected => {
                            exit_ev.write(AppExit::Success);
                        },
                    }
                },
            }
        }
    }

    while let Some(system) = systems.pop() {
        (system)(world)
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ExtractError {
    /// The extraction request did not make it to the SubWorld.
    /// This occurs when either the signal channel is full, or the SubWorld has exited.
    SendFailed,

    /// The provided system failed validation, possible causes include:
    ///  - Unregistered Resource
    ///  - Parameter Mutability Conflict
    /// The specific error is logged to console at the time of execution.
    RegistrationErr,

    /// The ExtractHandle has finished, and its output has already been consumed
    /// by an earlier call to `poll`. No more messages will be received from the
    /// ExtractHandle.
    AlreadyReceived,

    /// The Extraction has not yet finished.
    NotFinished,
}

/// A Handle to a system running in a SubWorld.
/// This handle can be polled to get the output.
#[derive(Component)]
pub struct ExtractHandle<Out=()> {
    shared: Arc<Shared<Out>>,
    send_time: Instant,
    recv_time: Option<Instant>,
}

impl<Out> ExtractHandle<Out> {
    /// The duration between when the task was started
    /// to when it was received. 
    pub fn duration(&self) -> Option<Duration> {
        self.recv_time.map(|time| {
            time.duration_since(self.send_time)
        })
    }

    /// The time the Extraction was dispatched.
    pub fn send_time(&self) -> Instant {
        self.send_time
    }

    /// The time the Extraction was received, if it has happened.
    pub fn recv_time(&self) -> Option<Instant> {
        self.recv_time
    }

    /// Whether the output of the task has already been received.
    pub fn is_received(&self) -> bool {
        self.recv_time.is_some()
    }

    /// Check if the output is ready.
    pub fn poll(&mut self) -> Result<Out, ExtractError> {
        match self.shared.get() {
            Ok(out) => {
                self.recv_time = Some(Instant::now());
                Ok(out)
            },
            Err(e) => Err(e)
        }
    }
}

struct Shared<Out> {
    inner: Mutex<Result<Out, ExtractError>>,  
}

impl<Out> Shared<Out> {
    fn new() -> Self {
        Self {
            inner: Mutex::new(Err(ExtractError::NotFinished))
        }
    }

    fn get(self: &Arc<Self>) -> Result<Out, ExtractError> {
        let mut guard = self.inner.lock();
        if guard.is_ok() {
            mem::replace(&mut*guard, Err(ExtractError::AlreadyReceived))
        } else {
            let Err(e) = &*guard else { unreachable!() };
            // arc dropped before a response was written, indicating the
            // channel was full or the subworld exited.
            if *e == ExtractError::NotFinished && Arc::strong_count(&self) == 1 {
                return Err(ExtractError::SendFailed);
            }
            Err(e.clone())
        }
    }

    fn set(&self, value: Result<Out, ExtractError>) {
        *self.inner.lock() = value;
    }
}