Refactor and integrate AssetBatcher. Add AssignedBatchesDrawTarget shim

bevy_derive/src/lib.rs

@@ -4,7 +4,7 @@ use darling::FromMeta;
 use inflector::Inflector;
 use proc_macro::TokenStream;
 use quote::{format_ident, quote};
-use syn::{parse_macro_input, Data, DataStruct, DeriveInput, Field, Fields, Type};
+use syn::{parse_macro_input, Data, DataStruct, DeriveInput, Field, Fields};
 
 #[derive(FromMeta, Debug, Default)]
 struct EntityArchetypeAttributeArgs {
@@ -112,10 +112,6 @@ pub fn derive_uniforms(input: TokenStream) -> TokenStream {
         })
         .map(|(f, _attr)| *f)
         .collect::<Vec<&Field>>();
-    let active_uniform_field_types = active_uniform_fields
-        .iter()
-        .map(|f| &f.ty)
-        .collect::<Vec<&Type>>();
 
     let shader_def_fields = uniform_fields
         .iter()

src/app/app_builder.rs

@@ -122,6 +122,8 @@ impl AppBuilder {
     }
 
     pub fn add_default_resources(mut self) -> Self {
+        let mut asset_batchers = AssetBatchers::default();
+        asset_batchers.batch_types2::<Mesh, StandardMaterial>();
         self.resources.insert(Time::new());
         self.resources.insert(AssetStorage::<Mesh>::new());
         self.resources.insert(AssetStorage::<Texture>::new());
@@ -132,7 +134,7 @@ impl AppBuilder {
             .insert(AssetStorage::<PipelineDescriptor>::new());
         self.resources.insert(ShaderPipelineAssignments::new());
         self.resources.insert(CompiledShaderMap::new());
-        self.resources.insert(AssetBatchers::default());
+        self.resources.insert(asset_batchers);
         self
     }
 
@@ -149,6 +151,7 @@ impl AppBuilder {
         self.setup_render_graph(|builder, pipeline_storage, shader_storage| {
             builder
                 .add_draw_target(MeshesDrawTarget::default())
+                .add_draw_target(AssignedBatchesDrawTarget::default())
                 .add_draw_target(AssignedMeshesDrawTarget::default())
                 .add_draw_target(UiDrawTarget::default())
                 .add_resource_provider(CameraResourceProvider::default())

src/render/draw_target/draw_targets/assigned_batches_draw_target.rs

@@ -0,0 +1,48 @@
+use crate::{
+    asset::Handle,
+    legion::prelude::*,
+    render::{
+        draw_target::DrawTarget,
+        pipeline::PipelineDescriptor,
+        render_resource::{resource_name, AssetBatchers},
+        renderer::{RenderPass, Renderer},
+    },
+};
+
+#[derive(Default)]
+pub struct AssignedBatchesDrawTarget;
+
+impl DrawTarget for AssignedBatchesDrawTarget {
+    fn draw(
+        &self,
+        _world: &World,
+        resources: &Resources,
+        _render_pass: &mut dyn RenderPass,
+        _pipeline_handle: Handle<PipelineDescriptor>,
+    ) {
+        let asset_batches = resources.get::<AssetBatchers>().unwrap();
+        // let renderer = render_pass.get_renderer();
+        // println!("Drawing batches");
+        for batch in asset_batches.get_batches() {
+            // render_resources.get
+            // println!("{:?}", batch);
+        }
+
+        // println!();
+        // println!();
+        // println!();
+    }
+
+    fn setup(
+        &mut self,
+        _world: &World,
+        _resources: &Resources,
+        _renderer: &mut dyn Renderer,
+        _pipeline_handle: Handle<PipelineDescriptor>,
+    ) {
+    }
+
+    fn get_name(&self) -> String {
+        resource_name::draw_target::ASSIGNED_BATCHES.to_string()
+    }
+}

src/render/draw_target/draw_targets/mod.rs

@@ -1,7 +1,9 @@
 mod assigned_meshes_draw_target;
 mod meshes_draw_target;
 mod ui_draw_target;
+mod assigned_batches_draw_target;
 
 pub use assigned_meshes_draw_target::*;
 pub use meshes_draw_target::*;
 pub use ui_draw_target::*;
+pub use assigned_batches_draw_target::*;

src/render/pipeline/pipelines/forward/mod.rs

@@ -71,6 +71,7 @@ impl ForwardPipelineBuilder for RenderGraphBuilder {
                 write_mask: ColorWrite::ALL,
             })
             .add_draw_target(resource_name::draw_target::ASSIGNED_MESHES)
+            .add_draw_target(resource_name::draw_target::ASSIGNED_BATCHES)
             .finish(),
         )
     }

src/render/render_resource/asset_batcher.rs

@@ -5,10 +5,10 @@ use std::{any::TypeId, collections::HashMap, hash::Hash};
 
 // TODO: if/when const generics land, revisit this design
 
-#[derive(Hash, Eq, PartialEq, Debug)]
+#[derive(Hash, Eq, PartialEq, Debug, Ord, PartialOrd)]
 pub struct BatchKey2 {
-    handle1: HandleId,
+    pub handle1: HandleId,
-    handle2: HandleId,
+    pub handle2: HandleId,
 }
 
 #[derive(Hash, Eq, PartialEq, Clone, Debug)]
@@ -28,16 +28,41 @@ impl EntitySetState2 {
     }
 }
 
-#[derive(Hash, PartialEq, Debug)]
+#[derive(PartialEq, Eq, Debug, Default)]
-pub struct Batch2 {
+pub struct Batch {
-    pub entities: Vec<Entity>,
+    pub entity_indices: HashMap<Entity, usize>,
-    pub buffer1: Option<RenderResource>,
+    pub current_index: usize,
-    pub buffer2: Option<RenderResource>,
+    pub render_resource_assignments: RenderResourceAssignments,
+}
+
+impl Batch {
+    pub fn add_entity(&mut self, entity: Entity) {
+        if let None = self.entity_indices.get(&entity) {
+            self.entity_indices.insert(entity, self.current_index);
+            self.current_index += 1;
+        }
+    }
+}
+
+// TODO: consider merging this with entity_uniform_resource
+#[derive(Eq, PartialEq, Debug, Default)]
+pub struct RenderResourceAssignments {
+    render_resources: HashMap<String, RenderResource>,
+}
+
+impl RenderResourceAssignments {
+    pub fn get_render_resource(&self, name: &str) -> Option<RenderResource> {
+        self.render_resources.get(name).cloned()
+    }
+
+    pub fn set_render_resource(&mut self, name: &str, resource: RenderResource) {
+        self.render_resources.insert(name.to_string(), resource);
+    }
 }
 
 pub struct AssetSetBatcher2 {
     key: AssetSetBatcherKey2,
-    set_batches: HashMap<BatchKey2, Batch2>,
+    set_batches: HashMap<BatchKey2, Batch>,
     entity_set_states: HashMap<Entity, EntitySetState2>,
 }
 
@@ -59,18 +84,13 @@ impl AssetSetBatcher2 {
         };
 
         match self.set_batches.get_mut(&key) {
-            Some(instance_set) => {
+            Some(batch) => {
-                instance_set.entities.push(entity);
+                batch.add_entity(entity);
             }
             None => {
-                self.set_batches.insert(
+                let mut batch = Batch::default();
-                    key,
+                batch.add_entity(entity);
-                    Batch2 {
+                self.set_batches.insert(key, batch);
-                        entities: vec![entity],
-                        buffer1: None,
-                        buffer2: None,
-                    },
-                );
             }
         }
     }
@@ -126,19 +146,25 @@ impl AssetBatcher for AssetSetBatcher2 {
             self.set_entity_handle2(entity, handle_id);
         }
     }
-    fn get_batch2(&self, key: &BatchKey2) -> Option<&Batch2> {
+    fn get_batch2(&self, key: &BatchKey2) -> Option<&Batch> {
         self.set_batches.get(key)
     }
 
-    fn get_batches2(&self) -> std::collections::hash_map::Iter<'_, BatchKey2, Batch2> {
+    fn get_batches2(&self) -> std::collections::hash_map::Iter<'_, BatchKey2, Batch> {
         self.set_batches.iter()
     }
+
+    fn get_batches<'a>(&'a self) -> Box<dyn Iterator<Item = &'a Batch> + 'a> {
+        Box::new(self.set_batches.values())
+    }
 }
 
 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_batch2(&self, key: &BatchKey2) -> Option<&Batch>;
-    fn get_batches2(&self) -> std::collections::hash_map::Iter<'_, BatchKey2, Batch2>;
+    // TODO: add pipeline handle here
+    fn get_batches2(&self) -> std::collections::hash_map::Iter<'_, BatchKey2, Batch>;
+    fn get_batches<'a>(&'a self) -> Box<dyn Iterator<Item = &Batch> + 'a>;
 }
 
 #[derive(Default)]
@@ -175,7 +201,9 @@ impl AssetBatchers {
             .insert(key, self.asset_batchers.len() - 1);
     }
 
-    pub fn get_batches2<T1, T2>(&mut self) -> Option<std::collections::hash_map::Iter<'_, BatchKey2, Batch2>>
+    pub fn get_batches2<T1, T2>(
+        &self,
+    ) -> Option<std::collections::hash_map::Iter<'_, BatchKey2, Batch>>
     where
         T1: 'static,
         T2: 'static,
@@ -192,11 +220,7 @@ impl AssetBatchers {
         }
     }
 
-    pub fn get_batch2<T1, T2>(
+    pub fn get_batch2<T1, T2>(&self, handle1: Handle<T1>, handle2: Handle<T2>) -> Option<&Batch>
-        &mut self,
-        handle1: Handle<T1>,
-        handle2: Handle<T2>,
-    ) -> Option<&Batch2>
     where
         T1: 'static,
         T2: 'static,
@@ -217,6 +241,15 @@ impl AssetBatchers {
             None
         }
     }
+
+    pub fn get_batches<'a>(&'a self) -> Box<dyn Iterator<Item = &Batch> + 'a> {
+        Box::new(
+            self.asset_batchers
+                .iter()
+                .map(|a| a.get_batches())
+                .flatten(),
+        )
+    }
 }
 
 #[cfg(test)]
@@ -246,13 +279,11 @@ mod tests {
         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<A> and Handle<B>
+        let mut expected_batch = Batch::default();
+        expected_batch.add_entity(entities[0]);
         assert_eq!(
             asset_batchers.get_batch2(a1, b1).unwrap(),
-            &Batch2 {
+            &expected_batch
-                entities: vec![entities[0],],
-                buffer1: None,
-                buffer2: None,
-            }
         );
         asset_batchers.set_entity_handle(entities[0], c1);
 
@@ -260,13 +291,12 @@ mod tests {
         asset_batchers.set_entity_handle(entities[1], b1);
 
         // all entities with Handle<A> and Handle<B> are returned
+        let mut expected_batch = Batch::default();
+        expected_batch.add_entity(entities[0]);
+        expected_batch.add_entity(entities[1]);
         assert_eq!(
             asset_batchers.get_batch2(a1, b1).unwrap(),
-            &Batch2 {
+            &expected_batch
-                entities: vec![entities[0], entities[1],],
-                buffer1: None,
-                buffer2: None,
-            }
         );
 
         // uncreated batches are empty
@@ -275,32 +305,41 @@ mod tests {
         // batch iteration works
         asset_batchers.set_entity_handle(entities[2], a2);
         asset_batchers.set_entity_handle(entities[2], b2);
-        assert_eq!(
+
-            asset_batchers
+        let mut batches = asset_batchers
             .get_batches2::<A, B>()
             .unwrap()
-                .collect::<Vec<(&BatchKey2, &Batch2)>>(),
+            .collect::<Vec<(&BatchKey2, &Batch)>>();
-            vec![(
+
-                &BatchKey2 {
+        batches.sort_by(|a, b| a.0.cmp(b.0));
+        let mut expected_batch1 = Batch::default();
+        expected_batch1.add_entity(entities[0]);
+        expected_batch1.add_entity(entities[1]);
+        let mut expected_batch2 = Batch::default();
+        expected_batch2.add_entity(entities[2]);
+        let mut expected_batches = vec![
+            (
+                BatchKey2 {
                     handle1: a1.id,
                     handle2: b1.id,
                 },
-                &Batch2 {
+                expected_batch1
-                    buffer1: None,
+            ),
-                    buffer2: None,
+            (
-                    entities: vec![entities[0], entities[1]]
+                BatchKey2 {
-                }
-            ),(
-                &BatchKey2 {
                     handle1: a2.id,
                     handle2: b2.id,
                 },
-                &Batch2 {
+                expected_batch2
-                    buffer1: None,
+            ),
-                    buffer2: None,
+        ];
-                    entities: vec![entities[2]]
+        expected_batches.sort_by(|a, b| a.0.cmp(&b.0));
-                }
+        assert_eq!(
-            )]
+            batches,
+            expected_batches
+                .iter()
+                .map(|(a, b)| (a, b))
+                .collect::<Vec<(&BatchKey2, &Batch)>>()
         );
     }
 }

src/render/render_resource/resource_name.rs

@@ -18,6 +18,7 @@ pub mod buffer {
 pub mod draw_target {
     pub const MESHES: &str = "Meshes";
     pub const ASSIGNED_MESHES: &str = "AssignedMeshes";
+    pub const ASSIGNED_BATCHES: &str = "AssignedBatches";
     pub const UI: &str = "Ui";
 }
 

src/render/render_resource/resource_providers/mesh_resource_provider.rs

@@ -2,7 +2,7 @@ use crate::{
     asset::{AssetStorage, Handle},
     prelude::Renderable,
     render::{
-        render_resource::{BufferUsage, ResourceProvider},
+        render_resource::{BufferUsage, ResourceProvider, AssetBatchers},
         renderer::Renderer,
         mesh::Mesh,
     },
@@ -41,7 +41,9 @@ impl MeshResourceProvider {
 impl ResourceProvider for MeshResourceProvider {
     fn update(&mut self, renderer: &mut dyn Renderer, world: &mut World, resources: &Resources) {
         let mesh_storage = resources.get_mut::<AssetStorage<Mesh>>().unwrap();
-        for (mesh_handle, _renderable) in self.mesh_query.iter(world) {
+        let mut asset_batchers = resources.get_mut::<AssetBatchers>().unwrap();
+        for (entity, (mesh_handle, _renderable)) in self.mesh_query.iter_entities(world) {
+            asset_batchers.set_entity_handle(entity, *mesh_handle);
             if let None = renderer
                 .get_render_resources()
                 .get_mesh_vertices_resource(*mesh_handle)

src/render/render_resource/resource_providers/uniform_resource_provider.rs

@@ -2,7 +2,7 @@ use crate::{
     asset::{AssetStorage, Handle},
     render::{
         pipeline::BindType,
-        render_resource::{BufferUsage, RenderResource, ResourceProvider},
+        render_resource::{BufferUsage, RenderResource, ResourceProvider, AssetBatchers},
         renderer::Renderer,
         shader::{AsUniforms, DynamicUniformBufferInfo, UniformInfoIter},
         texture::{SamplerDescriptor, Texture, TextureDescriptor},
@@ -67,9 +67,11 @@ where
         resources: &Resources,
     ) {
         let handle_query = self.handle_query.take().unwrap();
+        let mut asset_batchers = resources.get_mut::<AssetBatchers>().unwrap();
         // TODO: only update handle values when Asset value has changed
         if let Some(asset_storage) = resources.get::<AssetStorage<T>>() {
             for (entity, (handle, _renderable)) in handle_query.iter_entities(world) {
+                asset_batchers.set_entity_handle(entity, *handle);
                 if let Some(uniforms) = asset_storage.get(&handle) {
                     self.setup_entity_uniform_resources(
                         entity,