perf: add indexing on custom asset id (#16)

Author: ssouthcity

src/assets/indexing.rs

@@ -0,0 +1,56 @@
+use std::{collections::HashMap, marker::PhantomData};
+
+use bevy::prelude::*;
+
+pub trait Indexable {
+    fn index(&self) -> &String;
+}
+
+pub struct AssetIndexPlugin<T>(PhantomData<T>);
+
+impl<T> Default for AssetIndexPlugin<T> {
+    fn default() -> Self {
+        Self(PhantomData)
+    }
+}
+
+impl<T> Plugin for AssetIndexPlugin<T>
+where
+    T: Asset + Indexable,
+{
+    fn build(&self, app: &mut App) {
+        app.register_type::<IndexMap<T>>();
+        app.init_resource::<IndexMap<T>>();
+
+        app.add_systems(Update, populate_index::<T>);
+    }
+}
+
+#[derive(Resource, Reflect, Deref, DerefMut)]
+#[reflect(Resource)]
+pub struct IndexMap<T: Asset>(HashMap<String, AssetId<T>>);
+
+impl<T> Default for IndexMap<T>
+where
+    T: Asset,
+{
+    fn default() -> Self {
+        Self(HashMap::default())
+    }
+}
+
+fn populate_index<T: Asset + Indexable>(
+    mut events: EventReader<AssetEvent<T>>,
+    assets: Res<Assets<T>>,
+    mut index_map: ResMut<IndexMap<T>>,
+) {
+    for event in events.read() {
+        if let AssetEvent::Added { id } = event {
+            let Some(asset) = assets.get(*id) else {
+                continue;
+            };
+
+            index_map.insert(asset.index().to_owned(), *id);
+        }
+    }
+}

src/assets/mod.rs

@@ -1,5 +1,6 @@
 use bevy::prelude::*;
 
+pub mod indexing;
 pub mod loaders;
 pub mod manifest;
 mod tracking;

src/simulation/item/assets.rs

@@ -2,11 +2,18 @@ use bevy::{asset::LoadedFolder, prelude::*};
 use bevy_aseprite_ultra::prelude::*;
 use serde::Deserialize;
 
-use crate::assets::{LoadResource, loaders::toml::TomlAssetPlugin};
+use crate::assets::{
+    LoadResource,
+    indexing::{AssetIndexPlugin, Indexable},
+    loaders::toml::TomlAssetPlugin,
+};
 
 pub fn plugin(app: &mut App) {
     app.register_type::<ItemDef>();
-    app.add_plugins(TomlAssetPlugin::<ItemDef>::extensions(&["item.toml"]));
+    app.add_plugins((
+        TomlAssetPlugin::<ItemDef>::extensions(&["item.toml"]),
+        AssetIndexPlugin::<ItemDef>::default(),
+    ));
 
     app.register_type::<ItemAssets>();
     app.load_resource::<ItemAssets>();
@@ -19,6 +26,12 @@ pub struct ItemDef {
     pub stack_size: u32,
 }
 
+impl Indexable for ItemDef {
+    fn index(&self) -> &String {
+        &self.id
+    }
+}
+
 #[derive(Asset, Clone, Resource, Reflect)]
 #[reflect(Resource)]
 pub struct ItemAssets {

src/simulation/logistics/conveyor_belt.rs

@@ -7,6 +7,7 @@ use bevy::{
 use bevy_aseprite_ultra::prelude::*;
 
 use crate::{
+    assets::indexing::IndexMap,
     simulation::{
         FactorySystems,
         dismantle::QueueDismantle,
@@ -255,6 +256,7 @@ fn receive_items_from_belt(
     mut inputs: Query<&mut InputInventory>,
     mut commands: Commands,
     items: Res<Assets<ItemDef>>,
+    item_index: Res<IndexMap<ItemDef>>,
 ) {
     for (entity, item, progress, item_of) in conveyored_items {
         if progress.0 < 1.0 {
@@ -276,10 +278,9 @@ fn receive_items_from_belt(
             continue;
         };
 
-        let Some(item_def) = items
-            .iter()
-            .map(|(_, item)| item)
-            .find(|item_def| item_def.id == item.0)
+        let Some(item_def) = item_index
+            .get(&item.0)
+            .and_then(|asset_id| items.get(*asset_id))
         else {
             continue;
         };

src/simulation/recipe/assets.rs

@@ -3,11 +3,18 @@ use std::{collections::HashMap, time::Duration};
 use bevy::{asset::LoadedFolder, prelude::*};
 use serde::Deserialize;
 
-use crate::assets::{LoadResource, loaders::toml::TomlAssetPlugin};
+use crate::assets::{
+    LoadResource,
+    indexing::{AssetIndexPlugin, Indexable},
+    loaders::toml::TomlAssetPlugin,
+};
 
 pub fn plugin(app: &mut App) {
     app.register_type::<RecipeDef>();
-    app.add_plugins(TomlAssetPlugin::<RecipeDef>::extensions(&["recipe.toml"]));
+    app.add_plugins((
+        TomlAssetPlugin::<RecipeDef>::extensions(&["recipe.toml"]),
+        AssetIndexPlugin::<RecipeDef>::default(),
+    ));
 
     app.register_type::<RecipeAssets>();
     app.load_resource::<RecipeAssets>();
@@ -27,6 +34,12 @@ pub struct RecipeDef {
     pub tags: Vec<RecipeTags>,
 }
 
+impl Indexable for RecipeDef {
+    fn index(&self) -> &String {
+        &self.id
+    }
+}
+
 #[derive(Debug, Reflect, Deserialize, PartialEq, Eq)]
 #[serde(tag = "kind", content = "value", rename_all = "snake_case")]
 pub enum RecipeTags {

src/simulation/recipe/process.rs

@@ -1,11 +1,14 @@
 use bevy::prelude::*;
 
-use crate::simulation::{
-    FactorySystems,
-    item::{ItemDef, Stack},
-    logistics::{InputInventory, OutputInventory},
-    machine::power::Powered,
-    recipe::RecipeDef,
+use crate::{
+    assets::indexing::IndexMap,
+    simulation::{
+        FactorySystems,
+        item::{ItemDef, Stack},
+        logistics::{InputInventory, OutputInventory},
+        machine::power::Powered,
+        recipe::RecipeDef,
+    },
 };
 
 use super::{SelectedRecipe, progress::on_progress_state_add};
@@ -44,6 +47,7 @@ impl ProcessState {
 fn consume_input(
     query: Query<(&mut ProcessState, &mut InputInventory, &SelectedRecipe), With<Powered>>,
     recipes: Res<Assets<RecipeDef>>,
+    recipe_index: Res<IndexMap<RecipeDef>>,
 ) {
     for (mut state, mut inventory, selected_recipe) in query {
         if !matches!(*state, ProcessState::InsufficientInput) {
@@ -54,10 +58,9 @@ fn consume_input(
             continue;
         };
 
-        let recipe = recipes
-            .iter()
-            .map(|(_, recipe_def)| recipe_def)
-            .find(|recipe_def| recipe_def.id == *recipe_id)
+        let recipe = recipe_index
+            .get(recipe_id)
+            .and_then(|asset_id| recipes.get(*asset_id))
             .expect("Selected recipe refers to non-existent recipe");
 
         if inventory.consume_input(recipe).is_err() {
@@ -73,7 +76,9 @@ fn consume_input(
 fn progress_work(
     query: Query<(&mut ProcessState, &SelectedRecipe), With<Powered>>,
     recipes: Res<Assets<RecipeDef>>,
+    recipe_index: Res<IndexMap<RecipeDef>>,
     items: Res<Assets<ItemDef>>,
+    item_index: Res<IndexMap<ItemDef>>,
     time: Res<Time>,
 ) {
     for (mut state, selected_recipe) in query {
@@ -89,20 +94,18 @@ fn progress_work(
             continue;
         };
 
-        let recipe = recipes
-            .iter()
-            .map(|(_, recipe_def)| recipe_def)
-            .find(|recipe_def| recipe_def.id == *recipe_id)
+        let recipe = recipe_index
+            .get(recipe_id)
+            .and_then(|asset_id| recipes.get(*asset_id))
             .expect("Selected recipe refers to non-existent id");
 
         let output: Vec<Stack> = recipe
             .output
             .iter()
             .map(|(item_id, quantity)| {
-                let item_def = items
-                    .iter()
-                    .map(|(_, item_def)| item_def)
-                    .find(|item_def| item_def.id == *item_id)
+                let item_def = item_index
+                    .get(item_id)
+                    .and_then(|asset_id| items.get(*asset_id))
                     .expect("Recipe refers to non-existent output item");
 
                 Stack {

src/simulation/recipe/select.rs

@@ -1,9 +1,12 @@
 use bevy::prelude::*;
 
-use crate::simulation::{
-    item::{Inventory, ItemDef, Stack},
-    logistics::{InputInventory, OutputInventory},
-    recipe::{ProcessState, assets::RecipeDef},
+use crate::{
+    assets::indexing::IndexMap,
+    simulation::{
+        item::{Inventory, ItemDef, Stack},
+        logistics::{InputInventory, OutputInventory},
+        recipe::{ProcessState, assets::RecipeDef},
+    },
 };
 
 pub fn plugin(app: &mut App) {
@@ -24,26 +27,23 @@ pub struct SelectRecipe(pub String);
 fn on_select_recipe(
     trigger: Trigger<SelectRecipe>,
     recipes: Res<Assets<RecipeDef>>,
+    recipe_index: Res<IndexMap<RecipeDef>>,
     items: Res<Assets<ItemDef>>,
+    item_index: Res<IndexMap<ItemDef>>,
     mut commands: Commands,
 ) {
     let event = trigger.event();
 
-    let Some(recipe_def) = recipes
-        .iter()
-        .map(|(_, recipe)| recipe)
-        .find(|recipe| recipe.id == event.0)
-    else {
-        warn!("Attempted to select invalid recipe");
-        return;
-    };
+    let recipe_def = recipe_index
+        .get(&event.0)
+        .and_then(|asset_id| recipes.get(*asset_id))
+        .expect("Attempted to select invalid recipe");
 
     let mut input_inventory = Inventory::default();
-    for id in recipe_def.input.keys() {
-        let item_def = items
-            .iter()
-            .map(|(_, item)| item)
-            .find(|item| item.id == *id)
+    for item_id in recipe_def.input.keys() {
+        let item_def = item_index
+            .get(item_id)
+            .and_then(|asset_id| items.get(*asset_id))
             .expect("Recipe refers to non-existent item");
 
         let slot = Stack {
@@ -56,11 +56,10 @@ fn on_select_recipe(
     }
 
     let mut output_inventory = Inventory::default();
-    for id in recipe_def.output.keys() {
-        let item_def = items
-            .iter()
-            .map(|(_, item)| item)
-            .find(|item| item.id == *id)
+    for item_id in recipe_def.output.keys() {
+        let item_def = item_index
+            .get(item_id)
+            .and_then(|asset_id| items.get(*asset_id))
             .expect("Recipe refers to non-existent item");
 
         let slot = Stack {

src/simulation/world/deposit.rs

@@ -6,6 +6,7 @@ use serde::Deserialize;
 use crate::{
     assets::{
         LoadResource,
+        indexing::IndexMap,
         manifest::{Id, Manifest, ManifestPlugin},
     },
     screens::Screen,
@@ -108,23 +109,23 @@ fn on_mine_deposit(
     deposits: Query<&DepositRecipe>,
     mut inventory: Single<&mut PlayerInventory>,
     recipes: Res<Assets<RecipeDef>>,
+    recipe_index: Res<IndexMap<RecipeDef>>,
     items: Res<Assets<ItemDef>>,
+    item_index: Res<IndexMap<ItemDef>>,
 ) {
     let Ok(deposit_recipe) = deposits.get(trigger.target()) else {
         return;
     };
 
-    let recipe_def = recipes
-        .iter()
-        .map(|(_, recipe_def)| recipe_def)
-        .find(|recipe_def| recipe_def.id == deposit_recipe.0)
+    let recipe_def = recipe_index
+        .get(&deposit_recipe.0)
+        .and_then(|asset_id| recipes.get(*asset_id))
         .expect("Deposit refers to non-existent recipe");
 
     for (item_id, quantity) in recipe_def.output.iter() {
-        let item_def = items
-            .iter()
-            .map(|(_, item_def)| item_def)
-            .find(|item_def| item_def.id == *item_id)
+        let item_def = item_index
+            .get(item_id)
+            .and_then(|asset_id| items.get(*asset_id))
             .expect("Recipe refers to invalid item id");
 
         let mut stack = Stack {

src/ui/inspect/info.rs

@@ -4,6 +4,7 @@ use bevy::{
 };
 
 use crate::{
+    assets::indexing::IndexMap,
     simulation::recipe::{RecipeDef, SelectedRecipe},
     ui::{
         inspect::{InspectedEntity, InspectionMenuState},
@@ -24,6 +25,7 @@ pub fn open_recipe_menu(
     selected_recipes: Query<&SelectedRecipe>,
     // icons: Res<ItemAssets>,
     recipes: Res<Assets<RecipeDef>>,
+    recipe_index: Res<IndexMap<RecipeDef>>,
 ) {
     let Ok(selected_recipe) = selected_recipes.get(inspected_entity.0) else {
         return;
@@ -33,10 +35,9 @@ pub fn open_recipe_menu(
         return;
     };
 
-    let Some(recipe) = recipes
-        .iter()
-        .map(|(_, recipe_def)| recipe_def)
-        .find(|recipe_def| recipe_def.id == *recipe_id)
+    let Some(recipe) = recipe_index
+        .get(recipe_id)
+        .and_then(|asset_id| recipes.get(*asset_id))
     else {
         return;
     };