ranim_render/
extract.rs

1//! Composable extraction of evaluated CPU values into the render world.
2
3use std::{
4    any::TypeId,
5    collections::{HashMap, HashSet},
6    marker::PhantomData,
7};
8
9use bevy_app::{App, Plugin};
10use bevy_ecs::{prelude::*, world::EntityWorldMut};
11use ranim_core::{
12    frame::{EvaluatedFrame, FrameItemId},
13    glam::DAffine3,
14    transformed::Transformed,
15};
16
17use crate::{schedule::RenderSchedule, world::SourceOrder};
18
19/// Extract an evaluated value into render-world entities or resources.
20///
21/// Implementations can recurse into other values using the same context. Each
22/// call to [`ExtractContext::entity`] claims the next persistent output of the
23/// source frame item. Outputs no longer claimed on the next extraction are
24/// despawned automatically. Entities created directly through
25/// [`ExtractContext::world_mut`] remain the implementation's responsibility.
26pub trait Extract {
27    /// Synchronize this value's render-world outputs.
28    fn extract(&self, ctx: &mut ExtractContext<'_>);
29}
30
31/// Installs extraction for one exact evaluated frame payload type.
32///
33/// For example, a frame containing `Vec<T>` needs `ExtractPlugin<Vec<T>>`;
34/// installing `ExtractPlugin<T>` alone does not visit that batch. Recursive
35/// [`Extract`] implementations handle the values inside the registered type.
36/// Handwritten systems can also be added to [`RenderSchedule::Extract`].
37pub struct ExtractPlugin<T>(PhantomData<fn() -> T>);
38
39impl<T> Default for ExtractPlugin<T> {
40    fn default() -> Self {
41        Self(PhantomData)
42    }
43}
44
45impl<T: Extract + Send + Sync + 'static> Plugin for ExtractPlugin<T> {
46    fn build(&self, app: &mut App) {
47        app.init_resource::<ExtractedEntities<T>>()
48            .add_systems(RenderSchedule::Extract, extract_items::<T>);
49    }
50}
51
52type OutputId = (FrameItemId, usize);
53
54struct OutputEntity {
55    entity: Entity,
56    bundle_type: TypeId,
57}
58
59#[derive(Resource)]
60struct ExtractedEntities<T: Send + Sync + 'static> {
61    entities: HashMap<OutputId, OutputEntity>,
62    seen: HashSet<OutputId>,
63    inputs: HashSet<FrameItemId>,
64    marker: PhantomData<fn() -> T>,
65}
66
67impl<T: Send + Sync + 'static> Default for ExtractedEntities<T> {
68    fn default() -> Self {
69        Self {
70            entities: HashMap::new(),
71            seen: HashSet::new(),
72            inputs: HashSet::new(),
73            marker: PhantomData,
74        }
75    }
76}
77
78/// World access, output identity, order, and inherited transform for extraction.
79///
80/// A context belongs to one [`ranim_core::frame::FrameItem`]. Recursive values
81/// share its output sequence; transform scopes compose without changing that
82/// sequence. Input data remains unchanged when a frame is rendered again.
83pub struct ExtractContext<'w> {
84    world: &'w mut World,
85    entities: &'w mut HashMap<OutputId, OutputEntity>,
86    seen: &'w mut HashSet<OutputId>,
87    source: FrameItemId,
88    order: usize,
89    next_part: &'w mut usize,
90    transform: DAffine3,
91}
92
93impl ExtractContext<'_> {
94    /// Access the render world to update resources or perform custom operations.
95    ///
96    /// The current frame and this plugin's identity state are temporarily scoped
97    /// out of the world while extracting. The input value is available as `self`
98    /// in [`Extract::extract`]. Directly spawned entities are not tracked by the
99    /// context; manage their lifetimes in the implementation.
100    pub fn world_mut(&mut self) -> &mut World {
101        self.world
102    }
103
104    /// Claim the source's next persistent output entity.
105    ///
106    /// `B` identifies the output bundle's shape. The entity is reused while this
107    /// shape stays the same; changing it replaces the entity so old components
108    /// cannot remain. Values are inserted by the caller. Components managed
109    /// outside `B` remain the implementation's responsibility.
110    /// An output not claimed on a later frame is despawned.
111    pub fn entity<B: Bundle>(&mut self) -> EntityWorldMut<'_> {
112        self.next_entity(TypeId::of::<B>())
113    }
114
115    /// The accumulated local-to-world transform for the current value.
116    pub fn transform(&self) -> DAffine3 {
117        self.transform
118    }
119
120    /// Extract a subtree under `current * transform`, sharing output identities.
121    pub fn with_transform<R>(
122        &mut self,
123        transform: DAffine3,
124        extract: impl FnOnce(&mut ExtractContext<'_>) -> R,
125    ) -> R {
126        let mut child = ExtractContext {
127            world: self.world,
128            entities: self.entities,
129            seen: self.seen,
130            source: self.source,
131            order: self.order,
132            next_part: self.next_part,
133            transform: self.transform * transform,
134        };
135        extract(&mut child)
136    }
137
138    fn next_entity(&mut self, bundle_type: TypeId) -> EntityWorldMut<'_> {
139        let part = *self.next_part;
140        *self.next_part += 1;
141        let id = (self.source, part);
142        self.seen.insert(id);
143
144        let output = self.entities.entry(id).or_insert_with(|| OutputEntity {
145            entity: self.world.spawn_empty().id(),
146            bundle_type,
147        });
148        if output.bundle_type != bundle_type || !self.world.entities().contains(output.entity) {
149            self.world.despawn(output.entity);
150            output.entity = self.world.spawn_empty().id();
151            output.bundle_type = bundle_type;
152        }
153
154        let mut entity = self.world.entity_mut(output.entity);
155        let order = SourceOrder {
156            order: self.order,
157            part,
158        };
159        if entity.get::<SourceOrder>() != Some(&order) {
160            entity.insert(order);
161        }
162        entity
163    }
164}
165
166fn extract_items<T: Extract + Send + Sync + 'static>(world: &mut World) {
167    let _span = crate::cpu_probe::span("reconcile");
168    world.resource_scope(|world, mut state: Mut<ExtractedEntities<T>>| {
169        world.resource_scope(|world, frame: Mut<EvaluatedFrame>| {
170            let state = &mut *state;
171            state.seen.clear();
172            state.inputs.clear();
173            for item in frame.items::<T>() {
174                assert!(
175                    state.inputs.insert(item.id),
176                    "duplicate frame item id {:?} for `{}`",
177                    item.id,
178                    std::any::type_name::<T>()
179                );
180                let mut next_part = 0;
181                item.value.extract(&mut ExtractContext {
182                    world,
183                    entities: &mut state.entities,
184                    seen: &mut state.seen,
185                    source: item.id,
186                    order: item.order,
187                    next_part: &mut next_part,
188                    transform: DAffine3::IDENTITY,
189                });
190            }
191            state.entities.retain(|id, output| {
192                if state.seen.contains(id) {
193                    true
194                } else {
195                    world.despawn(output.entity);
196                    false
197                }
198            });
199        });
200    });
201}
202
203impl<T: Extract, G: Clone + Into<DAffine3>> Extract for Transformed<T, G> {
204    fn extract(&self, ctx: &mut ExtractContext<'_>) {
205        ctx.with_transform(self.transform.clone().into(), |ctx| self.inner.extract(ctx));
206    }
207}
208
209impl<T: Extract> Extract for Vec<T> {
210    fn extract(&self, ctx: &mut ExtractContext<'_>) {
211        for item in self {
212            item.extract(ctx);
213        }
214    }
215}
216
217impl<T: Extract> Extract for Option<T> {
218    fn extract(&self, ctx: &mut ExtractContext<'_>) {
219        if let Some(item) = self {
220            item.extract(ctx);
221        }
222    }
223}
224
225impl<T: Extract, const N: usize> Extract for [T; N] {
226    fn extract(&self, ctx: &mut ExtractContext<'_>) {
227        for item in self {
228            item.extract(ctx);
229        }
230    }
231}
232
233impl Extract for () {
234    fn extract(&self, _ctx: &mut ExtractContext<'_>) {}
235}
236
237macro_rules! impl_extract_for_tuple {
238    ($(($T:ident, $value:ident)),*) => {
239        impl<$($T: Extract),*> Extract for ($($T,)*) {
240            fn extract(&self, ctx: &mut ExtractContext<'_>) {
241                let ($($value,)*) = self;
242                $($value.extract(ctx);)*
243            }
244        }
245    };
246}
247
248variadics_please::all_tuples!(impl_extract_for_tuple, 1, 15, T, value);
249
250#[cfg(test)]
251mod tests {
252    use bevy_app::SubApp;
253    use bevy_ecs::{query::Changed, schedule::ScheduleLabel};
254    use ranim_core::glam::{DMat4, DVec3};
255
256    use super::*;
257
258    #[derive(Component, Default, Clone, Debug, PartialEq)]
259    struct TestItem {
260        value: f32,
261        transform: ranim_core::glam::Mat4,
262    }
263
264    impl Extract for TestItem {
265        fn extract(&self, ctx: &mut ExtractContext<'_>) {
266            let mut output = self.clone();
267            output.transform = DMat4::from(ctx.transform()).as_mat4() * output.transform;
268            let mut entity = ctx.entity::<Self>();
269            if entity.get::<Self>() != Some(&output) {
270                entity.insert(output);
271            }
272        }
273    }
274
275    #[derive(Component, Default, Clone)]
276    struct OtherItem {
277        _value: u8,
278    }
279
280    impl Extract for OtherItem {
281        fn extract(&self, ctx: &mut ExtractContext<'_>) {
282            ctx.entity::<Self>().insert(self.clone());
283        }
284    }
285
286    fn extraction_app<T: Extract + Send + Sync + 'static>() -> SubApp {
287        let mut app = SubApp::new();
288        app.update_schedule = Some(RenderSchedule::Extract.intern());
289        app.init_resource::<EvaluatedFrame>()
290            .add_plugins(ExtractPlugin::<T>::default());
291        app
292    }
293
294    fn entity<T: Send + Sync + 'static>(app: &SubApp, source: FrameItemId, part: usize) -> Entity {
295        app.world().resource::<ExtractedEntities<T>>().entities[&(source, part)].entity
296    }
297
298    #[test]
299    fn raw_compositions_extract_in_order_and_empty_vectors_remove_outputs() {
300        type Input = Vec<(TestItem, Option<OtherItem>)>;
301        let mut app = extraction_app::<Input>();
302        app.world_mut().resource_mut::<EvaluatedFrame>().push(
303            (7, 0),
304            vec![
305                (TestItem::default(), Some(OtherItem::default())),
306                (TestItem::default(), None),
307            ],
308        );
309        app.update();
310        for part in 0..3 {
311            assert_eq!(
312                app.world()
313                    .get::<SourceOrder>(entity::<Input>(&app, (7, 0), part)),
314                Some(&SourceOrder { order: 0, part }),
315            );
316        }
317        let previous = entity::<Input>(&app, (7, 0), 0);
318        app.world_mut()
319            .resource_mut::<EvaluatedFrame>()
320            .update([((7, 0), Input::new())]);
321        app.update();
322        assert!(!app.world().entities().contains(previous));
323        assert!(
324            app.world()
325                .resource::<ExtractedEntities<Input>>()
326                .entities
327                .is_empty()
328        );
329    }
330
331    #[test]
332    fn unchanged_component_outputs_keep_entities_and_changed_tracking() {
333        #[derive(Resource, Default)]
334        struct Changes(Vec<usize>);
335
336        let mut app = extraction_app::<Vec<TestItem>>();
337        app.init_resource::<Changes>().add_systems(
338            RenderSchedule::Extract,
339            (|items: Query<Entity, Changed<TestItem>>, mut changes: ResMut<Changes>| {
340                changes.0.push(items.iter().count());
341            })
342            .after(extract_items::<Vec<TestItem>>),
343        );
344        app.world_mut()
345            .resource_mut::<EvaluatedFrame>()
346            .push((1, 0), vec![TestItem::default()]);
347        app.update();
348        let original = entity::<Vec<TestItem>>(&app, (1, 0), 0);
349        app.update();
350        assert_eq!(entity::<Vec<TestItem>>(&app, (1, 0), 0), original);
351        assert_eq!(app.world().resource::<Changes>().0, [1, 0]);
352
353        let changed = TestItem {
354            value: 2.0,
355            ..Default::default()
356        };
357        app.world_mut()
358            .resource_mut::<EvaluatedFrame>()
359            .update([((1, 0), vec![changed])]);
360        app.update();
361        assert_eq!(app.world().resource::<Changes>().0, [1, 0, 1]);
362    }
363
364    #[test]
365    fn absent_optional_output_cannot_leave_a_previous_bundle_shape() {
366        type Input = (Option<TestItem>, OtherItem);
367        let mut app = extraction_app::<Input>();
368        app.world_mut()
369            .resource_mut::<EvaluatedFrame>()
370            .push((1, 0), (Some(TestItem::default()), OtherItem::default()));
371        app.update();
372        let first = entity::<Input>(&app, (1, 0), 0);
373        let second = entity::<Input>(&app, (1, 0), 1);
374
375        app.world_mut()
376            .resource_mut::<EvaluatedFrame>()
377            .update([((1, 0), (None::<TestItem>, OtherItem::default()))]);
378        app.update();
379        let current = entity::<Input>(&app, (1, 0), 0);
380        assert!(!app.world().entities().contains(first));
381        assert!(!app.world().entities().contains(second));
382        assert!(app.world().get::<OtherItem>(current).is_some());
383        assert!(app.world().get::<TestItem>(current).is_none());
384    }
385
386    #[test]
387    fn nested_transforms_compose_during_extraction_without_changing_raw_input() {
388        type Input = Transformed<Transformed<TestItem, DAffine3>, DAffine3>;
389        let outer = DAffine3::from_translation(DVec3::X * 3.0);
390        let inner = DAffine3::from_scale(DVec3::splat(2.0));
391        let input = Transformed::new(Transformed::new(TestItem::default(), inner), outer);
392        let mut app = extraction_app::<Input>();
393        app.world_mut()
394            .resource_mut::<EvaluatedFrame>()
395            .push((1, 0), input.clone());
396        app.update();
397        app.update();
398
399        let output = app
400            .world()
401            .get::<TestItem>(entity::<Input>(&app, (1, 0), 0))
402            .unwrap();
403        assert_eq!(output.transform, DMat4::from(outer * inner).as_mat4());
404        assert_eq!(
405            app.world().resource::<EvaluatedFrame>().items::<Input>()[0].value,
406            input
407        );
408    }
409
410    #[test]
411    fn a_missing_whole_batch_cleans_only_its_own_entities() {
412        let mut app = extraction_app::<TestItem>();
413        app.add_plugins(ExtractPlugin::<OtherItem>::default());
414        app.world_mut()
415            .resource_mut::<EvaluatedFrame>()
416            .push((1, 0), TestItem::default());
417        app.world_mut()
418            .resource_mut::<EvaluatedFrame>()
419            .push((2, 0), OtherItem::default());
420        app.update();
421        let removed = entity::<TestItem>(&app, (1, 0), 0);
422        let preserved = entity::<OtherItem>(&app, (2, 0), 0);
423
424        let mut frame = EvaluatedFrame::new();
425        frame.push((2, 0), OtherItem::default());
426        app.insert_resource(frame);
427        app.update();
428        assert!(!app.world().entities().contains(removed));
429        assert_eq!(entity::<OtherItem>(&app, (2, 0), 0), preserved);
430    }
431
432    #[test]
433    #[should_panic(expected = "duplicate frame item id (1, 0)")]
434    fn duplicate_input_ids_are_rejected() {
435        let mut app = extraction_app::<TestItem>();
436        app.world_mut()
437            .resource_mut::<EvaluatedFrame>()
438            .update([((1, 0), TestItem::default()), ((1, 0), TestItem::default())]);
439        app.update();
440    }
441}
ranim_render/extract.rs

ranim_render/
extract.rs
