From 99983b40a59e1f0a7a56142a643f67ec71200d56 Mon Sep 17 00:00:00 2001 From: Carter Anderson Date: Thu, 19 Mar 2020 13:21:55 -0700 Subject: [PATCH] AssetBatcher --- src/app/app_builder.rs | 2 + src/asset/mod.rs | 14 +- src/render/render_resource/asset_batcher.rs | 306 ++++++++++++++++++++ src/render/render_resource/mod.rs | 2 + 4 files changed, 318 insertions(+), 6 deletions(-) create mode 100644 src/render/render_resource/asset_batcher.rs diff --git a/src/app/app_builder.rs b/src/app/app_builder.rs index dc207b22c1..80ce410b87 100644 --- a/src/app/app_builder.rs +++ b/src/app/app_builder.rs @@ -17,6 +17,7 @@ use pipeline::PipelineDescriptor; use render_graph::RenderGraphBuilder; use shader::Shader; use std::collections::HashMap; +use render_resource::AssetBatchers; pub struct AppBuilder { pub world: World, @@ -131,6 +132,7 @@ impl AppBuilder { .insert(AssetStorage::::new()); self.resources.insert(ShaderPipelineAssignments::new()); self.resources.insert(CompiledShaderMap::new()); + self.resources.insert(AssetBatchers::default()); self } diff --git a/src/asset/mod.rs b/src/asset/mod.rs index ce01779d71..02895f9039 100644 --- a/src/asset/mod.rs +++ b/src/asset/mod.rs @@ -8,13 +8,15 @@ use std::{ use std::{collections::HashMap, marker::PhantomData}; +pub type HandleId = usize; + pub struct Handle { - pub id: usize, + pub id: HandleId, marker: PhantomData, } impl Handle { - pub fn new(id: usize) -> Self { + pub fn new(id: HandleId) -> Self { Handle { id, marker: PhantomData, @@ -68,9 +70,9 @@ pub trait Asset { } pub struct AssetStorage { - assets: HashMap, + assets: HashMap, names: HashMap>, - current_index: usize, + current_index: HandleId, } impl AssetStorage { @@ -100,11 +102,11 @@ impl AssetStorage { self.names.insert(name.to_string(), handle); } - pub fn get_id(&self, id: usize) -> Option<&T> { + pub fn get_id(&self, id: HandleId) -> Option<&T> { self.assets.get(&id) } - pub fn get_id_mut(&mut self, id: usize) -> Option<&mut T> { + pub fn get_id_mut(&mut self, id: HandleId) -> Option<&mut T> { self.assets.get_mut(&id) } diff --git a/src/render/render_resource/asset_batcher.rs b/src/render/render_resource/asset_batcher.rs new file mode 100644 index 0000000000..0030c8d27d --- /dev/null +++ b/src/render/render_resource/asset_batcher.rs @@ -0,0 +1,306 @@ +use super::RenderResource; +use crate::asset::{Handle, HandleId}; +use legion::prelude::Entity; +use std::{any::TypeId, collections::HashMap, hash::Hash}; + +// TODO: if/when const generics land, revisit this design + +#[derive(Hash, Eq, PartialEq, Debug)] +pub struct BatchKey2 { + handle1: HandleId, + handle2: HandleId, +} + +#[derive(Hash, Eq, PartialEq, Clone, Debug)] +pub struct AssetSetBatcherKey2 { + handle1_type: TypeId, + handle2_type: TypeId, +} + +struct EntitySetState2 { + handle1: Option, + handle2: Option, +} + +impl EntitySetState2 { + fn is_full(&self) -> bool { + self.handle1.is_some() && self.handle2.is_some() + } +} + +#[derive(Hash, PartialEq, Debug)] +pub struct Batch2 { + pub entities: Vec, + pub buffer1: Option, + pub buffer2: Option, +} + +pub struct AssetSetBatcher2 { + key: AssetSetBatcherKey2, + set_batches: HashMap, + entity_set_states: HashMap, +} + +impl AssetSetBatcher2 { + fn new(key: AssetSetBatcherKey2) -> Self { + AssetSetBatcher2 { + key, + set_batches: HashMap::new(), + entity_set_states: HashMap::new(), + } + } + + fn add_entity_to_set(&mut self, entity: Entity) { + // these unwraps are safe because this function is only called from set_entity_handle on a "full" state + let state = self.entity_set_states.get(&entity).unwrap(); + let key = BatchKey2 { + handle1: state.handle1.unwrap(), + handle2: state.handle2.unwrap(), + }; + + match self.set_batches.get_mut(&key) { + Some(instance_set) => { + instance_set.entities.push(entity); + } + None => { + self.set_batches.insert( + key, + Batch2 { + entities: vec![entity], + buffer1: None, + buffer2: None, + }, + ); + } + } + } + + pub fn set_entity_handle1(&mut self, entity: Entity, handle_id: HandleId) { + match self.entity_set_states.get_mut(&entity) { + None => { + // TODO: when generalizing to set size 1, ensure you treat set as "full" here + self.entity_set_states.insert( + entity, + EntitySetState2 { + handle1: Some(handle_id), + handle2: None, + }, + ); + } + Some(state) => { + state.handle1 = Some(handle_id); + if state.is_full() { + self.add_entity_to_set(entity); + } + } + } + } + + pub fn set_entity_handle2(&mut self, entity: Entity, handle_id: HandleId) { + match self.entity_set_states.get_mut(&entity) { + None => { + // TODO: when generalizing to set size 1, ensure you treat set as "full" here + self.entity_set_states.insert( + entity, + EntitySetState2 { + handle1: None, + handle2: Some(handle_id), + }, + ); + } + Some(state) => { + state.handle2 = Some(handle_id); + if state.is_full() { + self.add_entity_to_set(entity); + } + } + } + } +} + +impl AssetBatcher for AssetSetBatcher2 { + fn set_entity_handle(&mut self, entity: Entity, handle_type: TypeId, handle_id: HandleId) { + if handle_type == self.key.handle1_type { + self.set_entity_handle1(entity, handle_id); + } else if handle_type == self.key.handle2_type { + self.set_entity_handle2(entity, handle_id); + } + } + fn get_batch2(&self, key: &BatchKey2) -> Option<&Batch2> { + self.set_batches.get(key) + } + + fn get_batches2(&self) -> std::collections::hash_map::Iter<'_, BatchKey2, Batch2> { + self.set_batches.iter() + } +} + +pub trait AssetBatcher { + fn set_entity_handle(&mut self, entity: Entity, handle_type: TypeId, handle_id: HandleId); + fn get_batch2(&self, key: &BatchKey2) -> Option<&Batch2>; + fn get_batches2(&self) -> std::collections::hash_map::Iter<'_, BatchKey2, Batch2>; +} + +#[derive(Default)] +pub struct AssetBatchers { + asset_batchers: Vec>, + asset_batcher_indices2: HashMap, +} + +impl AssetBatchers { + pub fn set_entity_handle(&mut self, entity: Entity, handle: Handle) + where + T: 'static, + { + let handle_type = TypeId::of::(); + for asset_batcher in self.asset_batchers.iter_mut() { + asset_batcher.set_entity_handle(entity, handle_type, handle.id); + } + } + + pub fn batch_types2(&mut self) + where + T1: 'static, + T2: 'static, + { + let key = AssetSetBatcherKey2 { + handle1_type: TypeId::of::(), + handle2_type: TypeId::of::(), + }; + + self.asset_batchers + .push(Box::new(AssetSetBatcher2::new(key.clone()))); + + self.asset_batcher_indices2 + .insert(key, self.asset_batchers.len() - 1); + } + + pub fn get_batches2(&mut self) -> Option> + where + T1: 'static, + T2: 'static, + { + let key = AssetSetBatcherKey2 { + handle1_type: TypeId::of::(), + handle2_type: TypeId::of::(), + }; + + if let Some(index) = self.asset_batcher_indices2.get(&key) { + Some(self.asset_batchers[*index].get_batches2()) + } else { + None + } + } + + pub fn get_batch2( + &mut self, + handle1: Handle, + handle2: Handle, + ) -> Option<&Batch2> + where + T1: 'static, + T2: 'static, + { + let key = AssetSetBatcherKey2 { + handle1_type: TypeId::of::(), + handle2_type: TypeId::of::(), + }; + + let batch_key = BatchKey2 { + handle1: handle1.id, + handle2: handle2.id, + }; + + if let Some(index) = self.asset_batcher_indices2.get(&key) { + self.asset_batchers[*index].get_batch2(&batch_key) + } else { + None + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use legion::prelude::*; + struct A; + struct B; + struct C; + + #[test] + fn test_batching() { + let mut asset_batchers = AssetBatchers::default(); + asset_batchers.batch_types2::(); + + let mut world = World::new(); + let a1: Handle = Handle::new(1); + let b1: Handle = Handle::new(1); + let c1: Handle = Handle::new(1); + + let a2: Handle = Handle::new(2); + let b2: Handle = Handle::new(2); + + let entities = world.insert((), (0..3).map(|_| ())); + asset_batchers.set_entity_handle(entities[0], a1); + // batch is empty when Handle is missing + assert_eq!(asset_batchers.get_batch2(a1, b1), None); + asset_batchers.set_entity_handle(entities[0], b1); + // entity[0] is added to batch when it has both Handle and Handle + assert_eq!( + asset_batchers.get_batch2(a1, b1).unwrap(), + &Batch2 { + entities: vec![entities[0],], + buffer1: None, + buffer2: None, + } + ); + asset_batchers.set_entity_handle(entities[0], c1); + + asset_batchers.set_entity_handle(entities[1], a1); + asset_batchers.set_entity_handle(entities[1], b1); + + // all entities with Handle and Handle are returned + assert_eq!( + asset_batchers.get_batch2(a1, b1).unwrap(), + &Batch2 { + entities: vec![entities[0], entities[1],], + buffer1: None, + buffer2: None, + } + ); + + // uncreated batches are empty + assert_eq!(asset_batchers.get_batch2(a1, c1), None); + + // batch iteration works + asset_batchers.set_entity_handle(entities[2], a2); + asset_batchers.set_entity_handle(entities[2], b2); + assert_eq!( + asset_batchers + .get_batches2::() + .unwrap() + .collect::>(), + vec![( + &BatchKey2 { + handle1: a1.id, + handle2: b1.id, + }, + &Batch2 { + buffer1: None, + buffer2: None, + entities: vec![entities[0], entities[1]] + } + ),( + &BatchKey2 { + handle1: a2.id, + handle2: b2.id, + }, + &Batch2 { + buffer1: None, + buffer2: None, + entities: vec![entities[2]] + } + )] + ); + } +} diff --git a/src/render/render_resource/mod.rs b/src/render/render_resource/mod.rs index a1f3699ed2..7f819cb245 100644 --- a/src/render/render_resource/mod.rs +++ b/src/render/render_resource/mod.rs @@ -3,9 +3,11 @@ mod render_resource; mod resource_info; pub mod resource_name; mod resource_provider; +mod asset_batcher; pub mod resource_providers; pub use buffer::*; pub use render_resource::*; pub use resource_info::*; pub use resource_provider::*; +pub use asset_batcher::*;