better batch iteration (no allocations + abstracted out)

src/render/render_resource/asset_batcher.rs

@@ -265,49 +265,45 @@ impl AssetBatchers {
             .flatten()
     }
 
-    pub fn get_batcher_indices<T>(&self) -> impl Iterator<Item = &usize>
+    pub fn get_handle_batches<T>(&self) -> Option<impl Iterator<Item = &Batch>>
     where
         T: 'static,
     {
         let handle_type = TypeId::of::<T>();
-        self.handle_batchers.get(&handle_type).unwrap().iter()
+        if let Some(batcher_indices) = self.handle_batchers.get(&handle_type) {
+            Some(
+                // NOTE: it would be great to use batcher_indices.iter().map(|i| self.asset_batchers[*i].get_batches()) here
+                // but unfortunately the lifetimes don't work out for some reason
+                self.asset_batchers
+                    .iter()
+                    .enumerate()
+                    .filter(move |(index, _a)| batcher_indices.contains(index))
+                    .map(|(_index, a)| a.get_batches())
+                    .flatten(),
+            )
+        } else {
+            None
+        }
     }
 
-    pub fn get_batches_from_batcher(&self, index: usize) -> impl Iterator<Item = &Batch> {
+    pub fn get_handle_batches_mut<T>(&mut self) -> Option<impl Iterator<Item = &mut Batch>>
-        self.asset_batchers[index].get_batches()
+    where
+        T: 'static,
+    {
+        let handle_type = TypeId::of::<T>();
+        if let Some(batcher_indices) = self.handle_batchers.get(&handle_type) {
+            Some(
+                self.asset_batchers
+                    .iter_mut()
+                    .enumerate()
+                    .filter(move |(index, _a)| batcher_indices.contains(index))
+                    .map(|(_index, a)| a.get_batches_mut())
+                    .flatten(),
+            )
+        } else {
+            None
         }
-
-    pub fn get_batches_from_batcher_mut(
-        &mut self,
-        index: usize,
-    ) -> impl Iterator<Item = &mut Batch> {
-        self.asset_batchers[index].get_batches_mut()
     }
-
-    // pub fn get_handle_batches<T>(&self) -> Option<impl Iterator<Item = &Batch>>
-    // where
-    //     T: 'static,
-    // {
-    //     let handle_type = TypeId::of::<T>();
-    //     if let Some(batcher_indices) = self.handle_batchers.get(&handle_type) {
-    //         Some(
-    //             self.asset_batchers
-    //                 .iter()
-    //                 .enumerate()
-    //                 .filter(|(index, a)| {
-    //                     let handle_type = TypeId::of::<T>();
-    //                     self.handle_batchers
-    //                         .get(&handle_type)
-    //                         .unwrap()
-    //                         .contains(index)
-    //                 })
-    //                 .map(|(index, a)| a.get_batches())
-    //                 .flatten(),
-    //         )
-    //     } else {
-    //         None
-    //     }
-    // }
 }
 
 #[cfg(test)]

src/render/render_resource/resource_providers/uniform_resource_provider.rs

@@ -546,19 +546,15 @@ where
         _world: &mut World,
         resources: &Resources,
     ) {
+        // update batch resources. this needs to run in "finish_update" because batches aren't finalized across
+        // all members of the batch until "UniformResourceProvider.update" has run for all members of the batch 
         if let Some(asset_storage) = resources.get::<AssetStorage<T>>() {
-            let handle_type = std::any::TypeId::of::<T>();
             let mut asset_batchers = resources.get_mut::<AssetBatchers>().unwrap();
             let mut render_resource_assignments_provider = resources
                 .get_mut::<RenderResourceAssignmentsProvider>()
                 .unwrap();
-            // TODO: work out lifetime issues here so allocation isn't necessary
+            let handle_type = std::any::TypeId::of::<T>();
-            for index in asset_batchers
+            for batch in asset_batchers.get_handle_batches_mut::<T>().unwrap() {
-                .get_batcher_indices::<T>()
-                .map(|i| *i)
-                .collect::<Vec<usize>>()
-            {
-                for batch in asset_batchers.get_batches_from_batcher_mut(index) {
                 let handle: Handle<T> = batch
                     .handles
                     .iter()
@@ -583,4 +579,3 @@ where
         }
     }
 }
-}