Skip to main content

bevy_render/
extract_plugin.rs

1use crate::{
2    sync_world::{despawn_temporary_render_entities, entity_sync_system, SyncWorldPlugin},
3    Render, RenderApp, RenderSystems,
4};
5use bevy_app::{App, Plugin, SubApp};
6use bevy_derive::{Deref, DerefMut};
7use bevy_ecs::{
8    resource::Resource,
9    schedule::{IntoScheduleConfigs, Schedule, ScheduleBuildSettings, ScheduleLabel, Schedules},
10    world::{Mut, World},
11};
12use bevy_utils::default;
13
14/// Plugin that sets up the [`RenderApp`] and handles extracting data from the
15/// main world to the render world.
16pub struct ExtractPlugin {
17    /// Function that gets run at the beginning of each extraction.
18    ///
19    /// Gets the main world and render world as arguments (in that order).
20    pub pre_extract: fn(&mut World, &mut World),
21}
22
23impl Default for ExtractPlugin {
24    fn default() -> Self {
25        Self {
26            pre_extract: |_, _| {},
27        }
28    }
29}
30
31impl Plugin for ExtractPlugin {
32    fn build(&self, app: &mut App) {
33        app.add_plugins(SyncWorldPlugin);
34        app.init_resource::<ScratchMainWorld>();
35
36        let mut render_app = SubApp::new();
37
38        let mut extract_schedule = Schedule::new(ExtractSchedule);
39        // We skip applying any commands during the ExtractSchedule
40        // so commands can be applied on the render thread.
41        extract_schedule.set_build_settings(ScheduleBuildSettings {
42            auto_insert_apply_deferred: false,
43            ..default()
44        });
45        extract_schedule.set_apply_final_deferred(false);
46
47        render_app
48            .add_schedule(Render::base_schedule())
49            .add_schedule(extract_schedule)
50            .allow_ambiguous_resource::<MainWorld>()
51            .add_systems(
52                Render,
53                (
54                    // This set applies the commands from the extract schedule while the render schedule
55                    // is running in parallel with the main app.
56                    apply_extract_commands.in_set(RenderSystems::ExtractCommands),
57                    despawn_temporary_render_entities.in_set(RenderSystems::PostCleanup),
58                ),
59            );
60
61        let pre_extract = self.pre_extract;
62        render_app.set_extract(move |main_world, render_world| {
63            pre_extract(main_world, render_world);
64
65            {
66                #[cfg(feature = "trace")]
67                let _stage_span = bevy_log::info_span!("entity_sync").entered();
68                entity_sync_system(main_world, render_world);
69            }
70
71            // run extract schedule
72            extract(main_world, render_world);
73        });
74
75        app.insert_sub_app(RenderApp, render_app);
76    }
77}
78
79/// Schedule in which data from the main world is 'extracted' into the render world.
80///
81/// This step should be kept as short as possible to increase the "pipelining potential" for
82/// running the next frame while rendering the current frame.
83///
84/// This schedule is run on the render world, but it also has access to the main world.
85/// See [`MainWorld`] and [`Extract`](crate::Extract) for details on how to access main world data from this schedule.
86#[derive(ScheduleLabel, PartialEq, Eq, Debug, Clone, Hash, Default)]
87pub struct ExtractSchedule;
88
89/// Applies the commands from the extract schedule. This happens during
90/// the render schedule rather than during extraction to allow the commands to run in parallel with the
91/// main app when pipelined rendering is enabled.
92fn apply_extract_commands(render_world: &mut World) {
93    render_world.resource_scope(|render_world, mut schedules: Mut<Schedules>| {
94        schedules
95            .get_mut(ExtractSchedule)
96            .unwrap()
97            .apply_deferred(render_world);
98    });
99}
100/// The simulation [`World`] of the application, stored as a resource.
101///
102/// This resource is only available during [`ExtractSchedule`] and not
103/// during command application of that schedule.
104/// See [`Extract`](crate::Extract) for more details.
105#[derive(Resource, Default, Deref, DerefMut)]
106pub struct MainWorld(World);
107
108/// A "scratch" world used to avoid allocating new worlds every frame when
109/// swapping out the [`MainWorld`] for [`ExtractSchedule`].
110#[derive(Resource, Default)]
111struct ScratchMainWorld(World);
112
113/// Executes the [`ExtractSchedule`] step of the renderer.
114/// This updates the render world with the extracted ECS data of the current frame.
115pub fn extract(main_world: &mut World, render_world: &mut World) {
116    // temporarily add the app world to the render world as a resource
117    let scratch_world = main_world.remove_resource::<ScratchMainWorld>().unwrap();
118    let inserted_world = core::mem::replace(main_world, scratch_world.0);
119    render_world.insert_resource(MainWorld(inserted_world));
120    render_world.run_schedule(ExtractSchedule);
121
122    // move the app world back, as if nothing happened.
123    let inserted_world = render_world.remove_resource::<MainWorld>().unwrap();
124    let scratch_world = core::mem::replace(main_world, inserted_world.0);
125    main_world.insert_resource(ScratchMainWorld(scratch_world));
126}
127
128#[cfg(test)]
129mod test {
130    use bevy_app::{App, Startup};
131    use bevy_ecs::{prelude::*, schedule::ScheduleLabel};
132
133    use crate::{
134        extract_component::{ExtractComponent, ExtractComponentPlugin},
135        extract_plugin::ExtractPlugin,
136        sync_component::SyncComponent,
137        sync_world::MainEntity,
138        Render, RenderApp,
139    };
140
141    #[derive(Component, Clone, Debug)]
142    struct RenderComponent;
143
144    #[derive(Component, Clone, Debug)]
145    struct RenderComponentExtra;
146
147    #[derive(Component, Clone, Debug, ExtractComponent)]
148    struct RenderComponentSeparate;
149
150    #[derive(Component, Clone, Debug)]
151    struct RenderComponentNoExtract;
152
153    impl SyncComponent for RenderComponent {
154        type Target = (RenderComponent, RenderComponentExtra);
155    }
156
157    impl ExtractComponent for RenderComponent {
158        type QueryData = &'static Self;
159        type QueryFilter = ();
160        type Out = (RenderComponent, RenderComponentExtra);
161
162        fn extract_component(
163            _item: bevy_ecs::query::QueryItem<'_, '_, Self::QueryData>,
164        ) -> Option<Self::Out> {
165            Some((RenderComponent, RenderComponentExtra))
166        }
167    }
168
169    #[test]
170    fn extraction_works() {
171        let mut app = App::new();
172
173        app.add_plugins(ExtractPlugin::default());
174        app.add_plugins(ExtractComponentPlugin::<RenderComponent>::default());
175        app.add_plugins(ExtractComponentPlugin::<RenderComponentSeparate>::default());
176        app.add_systems(Startup, |mut commands: Commands| {
177            commands.spawn((RenderComponent, RenderComponentSeparate));
178        });
179
180        let render_app = app.get_sub_app_mut(RenderApp).unwrap();
181
182        // Normally RenderPlugin sets the RenderRecovery schedule as update, but for
183        // testing we just use the Render schedule directly.
184        render_app.update_schedule = Some(Render.intern());
185
186        render_app.world_mut().add_observer(
187            |event: On<Add, (RenderComponent, RenderComponentExtra)>, mut commands: Commands| {
188                // Simulate data that's not extracted
189                commands
190                    .entity(event.entity)
191                    .insert(RenderComponentNoExtract);
192            },
193        );
194
195        app.update();
196
197        // Check that all components have been extracted
198        {
199            let render_app = app.get_sub_app_mut(RenderApp).unwrap();
200            render_app
201                .world_mut()
202                .run_system_cached(
203                    |entity: Single<(
204                        &MainEntity,
205                        Option<&RenderComponent>,
206                        Option<&RenderComponentExtra>,
207                        Option<&RenderComponentSeparate>,
208                        Option<&RenderComponentNoExtract>,
209                    )>| {
210                        assert!(entity.1.is_some());
211                        assert!(entity.2.is_some());
212                        assert!(entity.3.is_some());
213                        assert!(entity.4.is_some());
214                    },
215                )
216                .unwrap();
217        }
218
219        // Remove RenderComponent
220        app.world_mut()
221            .run_system_cached(
222                |mut commands: Commands, query: Query<Entity, With<RenderComponent>>| {
223                    for entity in query {
224                        commands.entity(entity).remove::<RenderComponent>();
225                    }
226                },
227            )
228            .unwrap();
229
230        app.update();
231
232        // Check that the extracted components have been removed
233        {
234            let render_app = app.get_sub_app_mut(RenderApp).unwrap();
235            render_app
236                .world_mut()
237                .run_system_cached(
238                    |entity: Single<(
239                        &MainEntity,
240                        Option<&RenderComponent>,
241                        Option<&RenderComponentExtra>,
242                        Option<&RenderComponentSeparate>,
243                        Option<&RenderComponentNoExtract>,
244                    )>| {
245                        assert!(entity.1.is_none());
246                        assert!(entity.2.is_none());
247                        assert!(entity.3.is_some());
248                        assert!(entity.4.is_some());
249                    },
250                )
251                .unwrap();
252        }
253    }
254}