Add SceneFilter

crates/bevy_scene/src/dynamic_scene.rs

@@ -37,14 +37,13 @@ pub struct DynamicEntity {
 
 impl DynamicScene {
     /// Create a new dynamic scene from a given scene.
-    pub fn from_scene(scene: &Scene, type_registry: &AppTypeRegistry) -> Self {
-        Self::from_world(&scene.world, type_registry)
+    pub fn from_scene(scene: &Scene) -> Self {
+        Self::from_world(&scene.world)
     }
 
     /// Create a new dynamic scene from a given world.
-    pub fn from_world(world: &World, type_registry: &AppTypeRegistry) -> Self {
-        let mut builder =
-            DynamicSceneBuilder::from_world_with_type_registry(world, type_registry.clone());
+    pub fn from_world(world: &World) -> Self {
+        let mut builder = DynamicSceneBuilder::from_world(world);
 
         builder.extract_entities(world.iter_entities());
 

crates/bevy_scene/src/dynamic_scene_builder.rs

@@ -1,14 +1,23 @@
-use crate::{DynamicEntity, DynamicScene};
+use crate::{DynamicEntity, DynamicScene, SceneFilter};
 use bevy_app::AppTypeRegistry;
-use bevy_ecs::{prelude::Entity, reflect::ReflectComponent, world::World};
+use bevy_ecs::{
+    prelude::{Component, Entity},
+    reflect::ReflectComponent,
+    world::World,
+};
 use bevy_utils::default;
 use std::collections::BTreeMap;
 
 /// A [`DynamicScene`] builder, used to build a scene from a [`World`] by extracting some entities.
 ///
+/// # Component Extraction
+///
+/// By default, all components registered in a world's [`AppTypeRegistry`] resource will be extracted.
+/// This can be changed by [specifying a filter] or by explicitly [allowing]/[denying] certain components.
+///
 /// # Entity Order
 ///
-/// Extracted entities will always be stored in ascending order based on their [id](Entity::index).
+/// Extracted entities will always be stored in ascending order based on their [index](Entity::index).
 /// This means that inserting `Entity(1v0)` then `Entity(0v0)` will always result in the entities
 /// being ordered as `[Entity(0v0), Entity(1v0)]`.
 ///
@@ -30,31 +39,64 @@ use std::collections::BTreeMap;
 /// builder.extract_entity(entity);
 /// let dynamic_scene = builder.build();
 /// ```
+///
+/// [specifying a filter]: DynamicSceneBuilder::with_filter
+/// [allowing]: DynamicSceneBuilder::allow
+/// [denying]: DynamicSceneBuilder::deny
 pub struct DynamicSceneBuilder<'w> {
     entities: BTreeMap<u32, DynamicEntity>,
-    type_registry: AppTypeRegistry,
+    filter: Option<SceneFilter>,
     world: &'w World,
 }
 
 impl<'w> DynamicSceneBuilder<'w> {
     /// Prepare a builder that will extract entities and their component from the given [`World`].
-    /// All components registered in that world's [`AppTypeRegistry`] resource will be extracted.
     pub fn from_world(world: &'w World) -> Self {
         Self {
             entities: default(),
-            type_registry: world.resource::<AppTypeRegistry>().clone(),
+            filter: None,
             world,
         }
     }
 
-    /// Prepare a builder that will extract entities and their component from the given [`World`].
-    /// Only components registered in the given [`AppTypeRegistry`] will be extracted.
-    pub fn from_world_with_type_registry(world: &'w World, type_registry: AppTypeRegistry) -> Self {
-        Self {
-            entities: default(),
-            type_registry,
-            world,
-        }
+    /// Specify a custom [`SceneFilter`] to be used with this builder.
+    pub fn with_filter(&mut self, filter: Option<SceneFilter>) -> &mut Self {
+        self.filter = filter;
+        self
+    }
+
+    /// Allows the given component type, `T`, to be included in the generated scene.
+    ///
+    /// This method may be called for multiple components.
+    ///
+    /// Please note that this method is mutually exclusive with the [`deny`](Self::deny) method.
+    /// Calling this one will replace the denylist with a new allowlist.
+    pub fn allow<T: Component>(&mut self) -> &mut Self {
+        self.filter = Some(
+            self.filter
+                .take()
+                .unwrap_or_else(SceneFilter::new_allowlist)
+                .allow::<T>(),
+        );
+
+        self
+    }
+
+    /// Denies the given component type, `T`, from being included in the generated scene.
+    ///
+    /// This method may be called for multiple components.
+    ///
+    /// Please note that this method is mutually exclusive with the [`allow`](Self::allow) method.
+    /// Calling this one will replace the allowlist with a new denylist.
+    pub fn deny<T: Component>(&mut self) -> &mut Self {
+        self.filter = Some(
+            self.filter
+                .take()
+                .unwrap_or_else(SceneFilter::new_denylist)
+                .deny::<T>(),
+        );
+
+        self
     }
 
     /// Consume the builder, producing a [`DynamicScene`].
@@ -97,7 +139,7 @@ impl<'w> DynamicSceneBuilder<'w> {
     /// let scene = builder.build();
     /// ```
     pub fn extract_entities(&mut self, entities: impl Iterator<Item = Entity>) -> &mut Self {
-        let type_registry = self.type_registry.read();
+        let type_registry = self.world.resource::<AppTypeRegistry>().read();
 
         for entity in entities {
             let index = entity.index();
@@ -112,23 +154,38 @@ impl<'w> DynamicSceneBuilder<'w> {
             };
 
             for component_id in self.world.entity(entity).archetype().components() {
-                let reflect_component = self
+                let Some(type_id) = self
                     .world
                     .components()
                     .get_info(component_id)
-                    .and_then(|info| type_registry.get(info.type_id().unwrap()))
-                    .and_then(|registration| registration.data::<ReflectComponent>());
-
-                if let Some(reflect_component) = reflect_component {
-                    if let Some(component) = reflect_component.reflect(self.world, entity) {
-                        entry.components.push(component.clone_value());
-                    }
+                    .and_then(|info| info.type_id()) else {
+                    continue;
+                };
+
+                let is_denied = self
+                    .filter
+                    .as_ref()
+                    .map(|filter| filter.is_denied_by_id(type_id))
+                    .unwrap_or_default();
+
+                if is_denied {
+                    // Component is either in the denylist or _not_ in the allowlist
+                    continue;
                 }
+
+                let Some(component) = type_registry
+                    .get(type_id)
+                    .and_then(|registration| registration.data::<ReflectComponent>())
+                    .and_then(|reflect_component| reflect_component.reflect(self.world, entity)) else {
+                    continue;
+                };
+
+                entry.components.push(component.clone_value());
             }
+
             self.entities.insert(index, entry);
         }
 
-        drop(type_registry);
         self
     }
 }
@@ -270,4 +327,60 @@ mod tests {
         scene_entities.sort();
         assert_eq!(scene_entities, [entity_a_b.index(), entity_a.index()]);
     }
+
+    #[test]
+    fn should_extract_allowed_components() {
+        let mut world = World::default();
+
+        let atr = AppTypeRegistry::default();
+        {
+            let mut register = atr.write();
+            register.register::<ComponentA>();
+            register.register::<ComponentB>();
+        }
+        world.insert_resource(atr);
+
+        let entity_a_b = world.spawn((ComponentA, ComponentB)).id();
+        let entity_a = world.spawn(ComponentA).id();
+        let entity_b = world.spawn(ComponentB).id();
+
+        let mut builder = DynamicSceneBuilder::from_world(&world);
+        builder
+            .allow::<ComponentA>()
+            .extract_entities([entity_a_b, entity_a, entity_b].into_iter());
+        let scene = builder.build();
+
+        assert_eq!(scene.entities.len(), 3);
+        assert!(scene.entities[0].components[0].represents::<ComponentA>());
+        assert!(scene.entities[1].components[0].represents::<ComponentA>());
+        assert_eq!(scene.entities[2].components.len(), 0);
+    }
+
+    #[test]
+    fn should_not_extract_denied_components() {
+        let mut world = World::default();
+
+        let atr = AppTypeRegistry::default();
+        {
+            let mut register = atr.write();
+            register.register::<ComponentA>();
+            register.register::<ComponentB>();
+        }
+        world.insert_resource(atr);
+
+        let entity_a_b = world.spawn((ComponentA, ComponentB)).id();
+        let entity_a = world.spawn(ComponentA).id();
+        let entity_b = world.spawn(ComponentB).id();
+
+        let mut builder = DynamicSceneBuilder::from_world(&world);
+        builder
+            .deny::<ComponentA>()
+            .extract_entities([entity_a_b, entity_a, entity_b].into_iter());
+        let scene = builder.build();
+
+        assert_eq!(scene.entities.len(), 3);
+        assert!(scene.entities[0].components[0].represents::<ComponentB>());
+        assert_eq!(scene.entities[1].components.len(), 0);
+        assert!(scene.entities[2].components[0].represents::<ComponentB>());
+    }
 }

crates/bevy_scene/src/lib.rs

@@ -2,6 +2,7 @@ mod bundle;
 mod dynamic_scene;
 mod dynamic_scene_builder;
 mod scene;
+mod scene_filter;
 mod scene_loader;
 mod scene_spawner;
 
@@ -12,13 +13,15 @@ pub use bundle::*;
 pub use dynamic_scene::*;
 pub use dynamic_scene_builder::*;
 pub use scene::*;
+pub use scene_filter::*;
 pub use scene_loader::*;
 pub use scene_spawner::*;
 
 pub mod prelude {
     #[doc(hidden)]
     pub use crate::{
-        DynamicScene, DynamicSceneBuilder, DynamicSceneBundle, Scene, SceneBundle, SceneSpawner,
+        DynamicScene, DynamicSceneBuilder, DynamicSceneBundle, Scene, SceneBundle, SceneFilter,
+        SceneSpawner,
     };
 }
 

crates/bevy_scene/src/scene_filter.rs

@@ -0,0 +1,104 @@
+use bevy_utils::HashSet;
+use std::any::{Any, TypeId};
+
+/// A filter used to control which types can be added to a [`DynamicScene`].
+///
+/// This scene filter _can_ be used more generically to represent a filter for any given type;
+/// however, note that its intended usage with `DynamicScene` only considers [components].
+/// Adding types that are not components will have no effect when used with `DynamicScene`.
+///
+/// [`DynamicScene`]: crate::DynamicScene
+/// [components]: bevy_ecs::prelude::Component
+pub enum SceneFilter {
+    /// Contains the set of permitted types by their [`TypeId`].
+    ///
+    /// Types not contained within this set should not be allowed to be saved to an associated [`DynamicScene`].
+    ///
+    /// [`DynamicScene`]: crate::DynamicScene
+    Allowlist(HashSet<TypeId>),
+    /// Contains the set of prohibited types by their [`TypeId`].
+    ///
+    /// Types contained within this set should not be allowed to be saved to an associated [`DynamicScene`].
+    ///
+    /// [`DynamicScene`]: crate::DynamicScene
+    Denylist(HashSet<TypeId>),
+}
+
+impl SceneFilter {
+    /// Create a new [allowlist](Self::Allowlist).
+    pub fn new_allowlist() -> Self {
+        Self::Allowlist(HashSet::new())
+    }
+
+    /// Create a new [denylist](Self::Denylist).
+    pub fn new_denylist() -> Self {
+        Self::Denylist(HashSet::new())
+    }
+
+    /// Allow the given type, `T`.
+    ///
+    /// If this filter is already set as a [denylist](Self::Denylist),
+    /// then it will be completely replaced by a new [allowlist](Self::Allowlist).
+    pub fn allow<T: Any>(self) -> Self {
+        self.allow_by_id(TypeId::of::<T>())
+    }
+
+    /// Allow the given type.
+    ///
+    /// If this filter is already set as a [denylist](Self::Denylist),
+    /// then it will be completely replaced by a new [allowlist](Self::Allowlist).
+    pub fn allow_by_id(mut self, type_id: TypeId) -> Self {
+        match self {
+            Self::Allowlist(ref mut list) => {
+                list.insert(type_id);
+                self
+            }
+            Self::Denylist(_) => Self::Allowlist(HashSet::from([type_id])),
+        }
+    }
+
+    /// Deny the given type, `T`.
+    ///
+    /// If this filter is already set as an [allowlist](Self::Allowlist),
+    /// then it will be completely replaced by a new [denylist](Self::Denylist).
+    pub fn deny<T: Any>(self) -> Self {
+        self.deny_by_id(TypeId::of::<T>())
+    }
+
+    /// Deny the given type.
+    ///
+    /// If this filter is already set as an [allowlist](Self::Allowlist),
+    /// then it will be completely replaced by a new [denylist](Self::Denylist).
+    pub fn deny_by_id(mut self, type_id: TypeId) -> Self {
+        match self {
+            Self::Allowlist(_) => Self::Denylist(HashSet::from([type_id])),
+            Self::Denylist(ref mut list) => {
+                list.insert(type_id);
+                self
+            }
+        }
+    }
+
+    /// Returns true if the given type, `T`, is allowed by the filter.
+    pub fn is_allowed<T: Any>(&self) -> bool {
+        self.is_allowed_by_id(TypeId::of::<T>())
+    }
+
+    /// Returns true if the given type is allowed by the filter.
+    pub fn is_allowed_by_id(&self, type_id: TypeId) -> bool {
+        match self {
+            SceneFilter::Allowlist(list) => list.contains(&type_id),
+            SceneFilter::Denylist(list) => !list.contains(&type_id),
+        }
+    }
+
+    /// Returns true if the given type, `T`, is denied by the filter.
+    pub fn is_denied<T: Any>(&self) -> bool {
+        self.is_denied_by_id(TypeId::of::<T>())
+    }
+
+    /// Returns true if the given type is denied by the filter.
+    pub fn is_denied_by_id(&self, type_id: TypeId) -> bool {
+        !self.is_allowed_by_id(type_id)
+    }
+}