bevy_background_compute/
lib.rs

1//! A bevy plugin streamlining task handling.
2
3use std::{future::Future, marker::PhantomData};
4
5use bevy_app::{App, Plugin, Update};
6use bevy_ecs::{
7    event::{Event, EventWriter},
8    schedule::{IntoSystemConfigs, SystemSet},
9    system::{Commands, ResMut, Resource},
10    world::{Command, Mut, World},
11};
12use bevy_tasks::AsyncComputeTaskPool;
13use pollable::{PollableTask, SpawnPollableExt};
14
15mod pollable;
16
17/// Extension trait to add compute_in_background to Commands
18pub trait ComputeInBackgroundCommandExt {
19    fn compute_in_background<F, T>(&mut self, future: F)
20    where
21        F: Future<Output = T> + Send + 'static,
22        T: Send + Sync + 'static;
23}
24
25/// The plugin to make a type computable in background
26pub struct BackgroundComputePlugin<T>(PhantomData<fn() -> T>);
27
28/// The set for the internal task checking system
29#[derive(SystemSet)]
30pub struct BackgroundComputeCheck<T>(PhantomData<T>);
31
32/// Command struct for holding a future
33struct ComputeInBackground<F, T>(F)
34where
35    F: Future<Output = T> + Send + 'static,
36    T: Send + Sync + 'static;
37
38/// Resource for holding background tasks to check for completion
39#[derive(Resource)]
40struct BackgroundTasks<T> {
41    tasks: Vec<PollableTask<T>>,
42}
43
44/// Event sent once a background compute completes
45#[derive(Event)]
46pub struct BackgroundComputeComplete<T>(pub T)
47where
48    T: Send + Sync + 'static;
49
50impl<'w, 's> ComputeInBackgroundCommandExt for Commands<'w, 's> {
51    fn compute_in_background<F, T>(&mut self, future: F)
52    where
53        F: Future<Output = T> + Send + 'static,
54        T: Send + Sync + 'static,
55    {
56        self.queue(ComputeInBackground(future))
57    }
58}
59
60impl<T> Default for BackgroundComputePlugin<T> {
61    fn default() -> Self {
62        Self(PhantomData)
63    }
64}
65
66impl<T> Default for BackgroundComputeCheck<T> {
67    fn default() -> Self {
68        Self(PhantomData)
69    }
70}
71
72// These impls have to be written manually for BackgroundComputeCheck<T>
73// instead of being derived because the #derive[] doesn't ignore the
74// T bounds in a PhantomData<T> member
75// TODO see: https://github.com/rust-lang/rust/issues/26925 wait 4 fix
76mod impls {
77    use super::BackgroundComputeCheck;
78    use std::fmt::Debug;
79    use std::hash::Hash;
80
81    // TODO this derive is undocumented as a requirement in the migration guide
82    // try to remove it, make an issue
83    impl<T> Clone for BackgroundComputeCheck<T> {
84        fn clone(&self) -> Self {
85            Self(self.0)
86        }
87    }
88
89    impl<T> Debug for BackgroundComputeCheck<T> {
90        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
91            write!(f, "BackgroundComputeCheck<{:?}>", self.0)
92        }
93    }
94
95    impl<T> PartialEq for BackgroundComputeCheck<T> {
96        fn eq(&self, other: &Self) -> bool {
97            self.0 == other.0
98        }
99    }
100
101    impl<T> Eq for BackgroundComputeCheck<T> {}
102
103    impl<T> Hash for BackgroundComputeCheck<T> {
104        fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
105            self.0.hash(state);
106        }
107    }
108}
109
110impl<T> Plugin for BackgroundComputePlugin<T>
111where
112    T: Send + Sync + 'static,
113{
114    fn build(&self, app: &mut App) {
115        app.add_event::<BackgroundComputeComplete<T>>()
116            .insert_resource(BackgroundTasks::<T> { tasks: vec![] })
117            .add_systems(
118                Update,
119                background_compute_check_system::<T>.in_set(BackgroundComputeCheck::<T>::default()),
120            );
121    }
122}
123
124impl<F, T> Command for ComputeInBackground<F, T>
125where
126    F: Future<Output = T> + Send + 'static,
127    T: Send + Sync + 'static,
128{
129    fn apply(self, world: &mut World) {
130        world.resource_scope(|_, mut holder: Mut<BackgroundTasks<T>>| {
131            let func = self.0;
132            holder
133                .tasks
134                .push(AsyncComputeTaskPool::get().spawn_pollable(func));
135        });
136    }
137}
138
139/// System responsible for checking tasks being computed in background and sending completion events
140fn background_compute_check_system<T>(
141    mut holder: ResMut<BackgroundTasks<T>>,
142    mut event_writer: EventWriter<BackgroundComputeComplete<T>>,
143) where
144    T: Send + Sync + 'static,
145{
146    holder.tasks.retain(|pollable| {
147        if let Some(value) = pollable.poll() {
148            event_writer.send(BackgroundComputeComplete(value));
149            false
150        } else {
151            true
152        }
153    })
154}
bevy_background_compute/lib.rs

bevy_background_compute/
lib.rs
