From 6b546414457938b49e74249cb65c4ae1e800fc6d Mon Sep 17 00:00:00 2001 From: Gino Valente Date: Mon, 28 Nov 2022 21:47:51 -0700 Subject: [PATCH] Add SceneFilter --- crates/bevy_scene/src/dynamic_scene.rs | 9 +- .../bevy_scene/src/dynamic_scene_builder.rs | 161 +++++++++++++++--- crates/bevy_scene/src/lib.rs | 5 +- crates/bevy_scene/src/scene_filter.rs | 104 +++++++++++ crates/bevy_scene/src/serde.rs | 22 +-- examples/scene/scene.rs | 20 ++- 6 files changed, 271 insertions(+), 50 deletions(-) create mode 100644 crates/bevy_scene/src/scene_filter.rs diff --git a/crates/bevy_scene/src/dynamic_scene.rs b/crates/bevy_scene/src/dynamic_scene.rs index 92f8cae312f4e4..65e19299148f96 100644 --- a/crates/bevy_scene/src/dynamic_scene.rs +++ b/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()); diff --git a/crates/bevy_scene/src/dynamic_scene_builder.rs b/crates/bevy_scene/src/dynamic_scene_builder.rs index 841af6756f3ba3..e7ba8ae9f239cd 100644 --- a/crates/bevy_scene/src/dynamic_scene_builder.rs +++ b/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, - type_registry: AppTypeRegistry, + filter: Option, 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::().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) -> &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(&mut self) -> &mut Self { + self.filter = Some( + self.filter + .take() + .unwrap_or_else(SceneFilter::new_allowlist) + .allow::(), + ); + + 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(&mut self) -> &mut Self { + self.filter = Some( + self.filter + .take() + .unwrap_or_else(SceneFilter::new_denylist) + .deny::(), + ); + + 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) -> &mut Self { - let type_registry = self.type_registry.read(); + let type_registry = self.world.resource::().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::()); - - 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::()) + .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::(); + register.register::(); + } + 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::() + .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::()); + assert!(scene.entities[1].components[0].represents::()); + 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::(); + register.register::(); + } + 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::() + .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::()); + assert_eq!(scene.entities[1].components.len(), 0); + assert!(scene.entities[2].components[0].represents::()); + } } diff --git a/crates/bevy_scene/src/lib.rs b/crates/bevy_scene/src/lib.rs index 184c2490d8718e..da4ba6605cfa47 100644 --- a/crates/bevy_scene/src/lib.rs +++ b/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, }; } diff --git a/crates/bevy_scene/src/scene_filter.rs b/crates/bevy_scene/src/scene_filter.rs new file mode 100644 index 00000000000000..d488f744b14e35 --- /dev/null +++ b/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), + /// 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), +} + +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(self) -> Self { + self.allow_by_id(TypeId::of::()) + } + + /// 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(self) -> Self { + self.deny_by_id(TypeId::of::()) + } + + /// 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(&self) -> bool { + self.is_allowed_by_id(TypeId::of::()) + } + + /// 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(&self) -> bool { + self.is_denied_by_id(TypeId::of::()) + } + + /// 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) + } +} diff --git a/crates/bevy_scene/src/serde.rs b/crates/bevy_scene/src/serde.rs index 0a225690168d2a..231c75f0c0984b 100644 --- a/crates/bevy_scene/src/serde.rs +++ b/crates/bevy_scene/src/serde.rs @@ -541,7 +541,7 @@ mod tests { let registry = world.resource::(); - let scene = DynamicScene::from_world(&world, registry); + let scene = DynamicScene::from_world(&world); let scene_serializer = SceneSerializer::new(&scene, ®istry.0); let serialized_scene = postcard::to_allocvec(&scene_serializer).unwrap(); @@ -579,7 +579,7 @@ mod tests { let registry = world.resource::(); - let scene = DynamicScene::from_world(&world, registry); + let scene = DynamicScene::from_world(&world); let scene_serializer = SceneSerializer::new(&scene, ®istry.0); let serialized_scene = bincode::serialize(&scene_serializer).unwrap(); @@ -655,8 +655,7 @@ mod tests { #[should_panic(expected = "entity count did not match")] fn should_panic_when_entity_count_not_eq() { let mut world = create_world(); - let registry = world.resource::(); - let scene_a = DynamicScene::from_world(&world, registry); + let scene_a = DynamicScene::from_world(&world); world.spawn(MyComponent { foo: [1, 2, 3], @@ -664,8 +663,7 @@ mod tests { baz: MyEnum::Unit, }); - let registry = world.resource::(); - let scene_b = DynamicScene::from_world(&world, registry); + let scene_b = DynamicScene::from_world(&world); assert_scene_eq(&scene_a, &scene_b); } @@ -683,8 +681,7 @@ mod tests { }) .id(); - let registry = world.resource::(); - let scene_a = DynamicScene::from_world(&world, registry); + let scene_a = DynamicScene::from_world(&world); world.entity_mut(entity).insert(MyComponent { foo: [3, 2, 1], @@ -692,8 +689,7 @@ mod tests { baz: MyEnum::Unit, }); - let registry = world.resource::(); - let scene_b = DynamicScene::from_world(&world, registry); + let scene_b = DynamicScene::from_world(&world); assert_scene_eq(&scene_a, &scene_b); } @@ -711,13 +707,11 @@ mod tests { }) .id(); - let registry = world.resource::(); - let scene_a = DynamicScene::from_world(&world, registry); + let scene_a = DynamicScene::from_world(&world); world.entity_mut(entity).remove::(); - let registry = world.resource::(); - let scene_b = DynamicScene::from_world(&world, registry); + let scene_b = DynamicScene::from_world(&world); assert_scene_eq(&scene_a, &scene_b); } diff --git a/examples/scene/scene.rs b/examples/scene/scene.rs index a3210214ad7b9d..5ea083197fd785 100644 --- a/examples/scene/scene.rs +++ b/examples/scene/scene.rs @@ -85,9 +85,18 @@ fn log_system(query: Query<(Entity, &ComponentA), Changed>) { } fn save_scene_system(world: &mut World) { - // Scenes can be created from any ECS World. You can either create a new one for the scene or - // use the current World. + // Scenes can be created from any ECS World. + // You can either create a new one for the scene or use the current World. + // For demonstration purposes, we'll create a new one. let mut scene_world = World::new(); + + // The `TypeRegistry` resource contains information about all registered types (including components). + // This is used to construct scenes, so we'll want to ensure that our previous type registrations + // exist in this new scene world as well. + // To do this, we can simply clone the `AppTypeRegistry` resource. + let type_registry = world.resource::().clone(); + scene_world.insert_resource(type_registry); + let mut component_b = ComponentB::from_world(world); component_b.value = "hello".to_string(); scene_world.spawn(( @@ -97,12 +106,11 @@ fn save_scene_system(world: &mut World) { )); scene_world.spawn(ComponentA { x: 3.0, y: 4.0 }); - // The TypeRegistry resource contains information about all registered types (including - // components). This is used to construct scenes. - let type_registry = world.resource::(); - let scene = DynamicScene::from_world(&scene_world, type_registry); + // With our sample world ready to go, we can now create our scene: + let scene = DynamicScene::from_world(&scene_world); // Scenes can be serialized like this: + let type_registry = world.resource::(); let serialized_scene = scene.serialize_ron(type_registry).unwrap(); // Showing the scene in the console