use std::any::{type_name, Any, TypeId}; use std::borrow::Cow; use std::cell::RefCell; use std::convert::Infallible; use std::fmt::{Debug, Display}; use std::hash::Hash; use std::marker::PhantomData; use std::panic::{catch_unwind, RefUnwindSafe, UnwindSafe}; use std::path::{Path, PathBuf}; use std::sync::mpsc::{ channel as mpsc_channel, Receiver as MpscReceiver, Sender as MpscSender, }; use std::thread::available_parallelism; use rayon::{ThreadPool, ThreadPoolBuilder}; use crate::ecs::pair::ChildOf; use crate::ecs::phase::{Phase, PRE_UPDATE as PRE_UPDATE_PHASE}; use crate::ecs::sole::Single; use crate::ecs::{declare_entity, pair, Sole}; const BACKUP_IMPORT_WORK_THREAD_CNT: usize = 2; declare_entity! { pub HANDLE_ASSETS_PHASE: (Phase, pair!(ChildOf, { *PRE_UPDATE_PHASE })); } pub trait Asset: Send + Sync + 'static { fn dependencies(&self) -> impl Iterator>; } #[derive(Debug)] pub struct Dependency<'a> { label: Label<'a>, asset_ty_id: AssetTypeId, } impl<'a> Dependency<'a> { #[must_use] pub fn new(label: Label<'a>) -> Self { Self { label, asset_ty_id: AssetTypeId::of::(), } } } /// Asset label. #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] pub struct Label<'a> { pub path: Cow<'a, Path>, pub name: Option>, } impl Label<'_> { #[must_use] fn to_owned(&self) -> LabelOwned { LabelOwned { path: self.path.clone().into_owned(), name: self.name.clone().map(Cow::into_owned), } } } impl<'a> From<&'a Path> for Label<'a> { fn from(path: &'a Path) -> Self { Self { path: path.into(), name: None } } } impl From for Label<'_> { fn from(path: PathBuf) -> Self { Self { path: path.into(), name: None } } } impl Display for Label<'_> { fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(formatter, "{}", self.path.display())?; if let Some(name) = &self.name { formatter.write_str("::")?; formatter.write_str(name)?; } Ok(()) } } #[derive(Debug, Sole)] pub struct Assets { store: store::Store, metadata_lut: RefCell>, importers: hashbrown::HashMap, import_work_thread_pool: ThreadPool, import_work_msg_receiver: MpscReceiver, import_work_msg_sender: MpscSender, curr_tick_events: Vec, } impl Assets { #[must_use] #[tracing::instrument(skip_all)] pub(crate) fn with_capacity(capacity: usize) -> Self { let (import_work_msg_sender, import_work_msg_receiver) = mpsc_channel::(); let import_work_thread_cnt = calc_import_work_thread_cnt(); Self { store: store::Store::with_capacity(capacity), metadata_lut: RefCell::new(hashbrown::HashMap::with_capacity(capacity)), importers: hashbrown::HashMap::new(), import_work_thread_pool: ThreadPoolBuilder::new() .thread_name(|thread_index| format!("asset_import_worker_{thread_index}")) .num_threads(import_work_thread_cnt) .build() .expect("Unable to create thread pool"), import_work_msg_receiver, import_work_msg_sender, curr_tick_events: Vec::with_capacity(capacity), } } pub fn set_importer( &mut self, func: impl Fn(&mut Submitter<'_>, &Path, Option<&AssetSettings>) -> Result<(), Err>, ) where AssetT: Asset, AssetSettings: UnwindSafe + RefUnwindSafe + 'static, Err: std::error::Error + Send + Sync + 'static, { self.importers .insert(AssetTypeId::of::(), WrappedImporterFn::new(func)); } #[tracing::instrument( skip_all, fields(asset_label, asset_type = %type_name::()) )] pub fn get(&self, handle: &Handle) -> Option<&AssetT> where AssetT: Asset, { let id = handle.id.into_valid()?; let stored_asset = self.store.get(id.index)?; tracing::Span::current() .record("asset_label", tracing::field::display(&stored_asset.label)); let AssetState::Ready { data: asset_data } = &stored_asset.state else { return None; }; let Some(asset) = asset_data.downcast_ref::() else { tracing::error!("Wrong asset type"); return None; }; Some(asset) } #[tracing::instrument( skip_all, fields(asset_label, asset_type = %type_name::()) )] pub fn get_mut(&mut self, handle: &Handle) -> Option<&mut AssetT> where AssetT: Asset, { let id = handle.id.into_valid()?; let stored_asset = self.store.get_mut(id.index)?; tracing::Span::current() .record("asset_label", tracing::field::display(&stored_asset.label)); let AssetState::Ready { data: asset_data } = &mut stored_asset.state else { return None; }; let Some(asset) = asset_data.downcast_mut::() else { tracing::error!("Wrong asset type"); return None; }; Some(asset) } #[tracing::instrument(skip_all, fields(label, asset_type = %type_name::()))] pub fn get_handle_to_loaded<'label, AssetT>( &self, label: impl Into>, ) -> Option> where AssetT: Asset, { let label = label.into(); tracing::Span::current().record("label", tracing::field::display(&label)); let (stored_asset, id) = self.get_asset_by_label(&label)?; let Some(StoredAsset { state: AssetState::Ready { data: asset_data }, .. }) = stored_asset else { return None; }; if asset_data.downcast_ref::().is_none() { tracing::error!("Wrong asset type"); return None; } Some(Handle::from_id(Id::from_valid(id))) } pub fn is_loaded_and_has_type(&self, handle: &Handle) -> bool where AssetT: Asset, { let Some(id) = handle.id.into_valid() else { return false; }; let Some(StoredAsset { state: AssetState::Ready { data: asset_data }, .. }) = self.store.get(id.index) else { return false; }; asset_data.downcast_ref::().is_some() } pub fn get_label(&self, handle: &Handle) -> Option> where AssetT: Asset, { self.get_label_by_id(handle.id) } pub fn get_label_by_id(&self, id: Id) -> Option> { let id = id.into_valid()?; let stored_asset = self.store.get(id.index)?; Some(stored_asset.label.clone().into_label()) } #[tracing::instrument(skip_all, fields(label, asset_type = %type_name::()))] pub fn load<'i, AssetT>(&self, label: impl Into>) -> Handle where AssetT: Asset, { let label = label.into(); tracing::Span::current().record("label", tracing::field::display(&label)); let id = self.load_asset::(AssetTypeId::of::(), label, None); Handle::from_id(Id::from_valid(id)) } #[tracing::instrument(skip_all, fields(label, asset_type = %type_name::()))] pub fn load_with_settings<'i, AssetT, AssetSettings>( &self, label: impl Into> + Debug, asset_settings: AssetSettings, ) -> Handle where AssetT: Asset, AssetSettings: UnwindSafe + RefUnwindSafe + Send + Sync + Debug + 'static, { let label = label.into(); tracing::Span::current().record("label", tracing::field::display(&label)); let id = self.load_asset(AssetTypeId::of::(), label, Some(asset_settings)); Handle::from_id(Id::from_valid(id)) } pub fn store_with_name<'name, AssetT>( &mut self, name: impl Into>, asset: AssetT, ) -> Handle where AssetT: Asset, { self.store_with_label( Label { path: Path::new("").into(), name: Some(name.into()), }, asset, ) } pub fn store_with_name_with<'name, AssetT>( &mut self, name: impl Into>, func: impl FnOnce(&mut Self) -> AssetT, ) -> Handle where AssetT: Asset, { let asset = func(self); self.store_with_label( Label { path: Path::new("").into(), name: Some(name.into()), }, asset, ) } #[tracing::instrument(skip_all, fields(label, asset_type = %type_name::()))] pub fn store_with_label<'i, AssetT>( &mut self, label: impl Into>, asset: AssetT, ) -> Handle where AssetT: Asset, { let label = label.into(); tracing::Span::current().record("label", tracing::field::display(&label)); let id = match self.get_asset_by_label(&label) { Some(( Some(StoredAsset { state: AssetState::Ready { .. } | AssetState::DependenciesPending { .. }, .. }), id, )) => { tracing::error!("Asset already exists"); return Handle::from_id(Id::from_valid(id)); } Some((Some(StoredAsset { state: AssetState::Pending, .. }) | None, id)) => id, None => IdValid::new(&self.store, &label), }; let mut missing_dependencies = 0; for dependency in asset.dependencies() { match self.get_asset_by_label_mut(&dependency.label) { Some((Some(StoredAsset { state: AssetState::Ready { .. }, .. }), _)) => {} Some(( Some(StoredAsset { state: AssetState::Pending | AssetState::DependenciesPending { .. }, waiting_dependants, .. }), _, )) => { tracing::trace!( dependency = %dependency.label, "Asset has missing dependency" ); missing_dependencies += 1; waiting_dependants.push(id); } Some((None, _)) | None => { tracing::trace!( dependency = %dependency.label, "Asset has missing dependency" ); missing_dependencies += 1; let dependency_id = self.load_asset::( dependency.asset_ty_id, dependency.label.clone(), None, ); self.store .get_or_insert_with(dependency_id.index, || StoredAsset { state: AssetState::Pending, label: dependency.label.to_owned(), waiting_dependants: Vec::new(), }) .waiting_dependants .push(id); } } } tracing::debug!("Storing asset"); let stored_asset = self.store.get_or_insert_with(id.index, || StoredAsset { state: AssetState::Pending, label: label.to_owned(), waiting_dependants: Vec::new(), }); let label = label.to_owned(); if missing_dependencies != 0 { stored_asset.state = AssetState::DependenciesPending { data: Box::new(asset), remaining_dependencies: missing_dependencies, }; } else { stored_asset.state = AssetState::Ready { data: Box::new(asset) }; self.curr_tick_events.push(Event::Stored( Id::from_valid(id), label.clone().into_label(), )); let waiting_dependants = std::mem::take(&mut stored_asset.waiting_dependants); self.recurse_update_dependants(waiting_dependants); } let metadata_lut = self.metadata_lut.get_mut(); let parent_id = label.name.is_some().then(|| { let parent_label = Label { path: label.path.as_path().into(), name: None, }; metadata_lut .entry(parent_label.to_owned()) .or_insert_with(|| Metadata { id: IdValid::new(&self.store, &parent_label), parent_id: None, }) .id }); metadata_lut.insert(label.clone(), Metadata { id, parent_id }); Handle::from_id(Id::from_valid(id)) } fn recurse_update_dependants(&mut self, initial_dependants: Vec) { let mut all_dependants = initial_dependants; loop { let dependants = std::mem::take(&mut all_dependants); if dependants.is_empty() { break; } for dependant_id in dependants { let Some(dependant_asset) = self.store.get_mut(dependant_id.index) else { continue; }; { let AssetState::DependenciesPending { remaining_dependencies, .. } = &mut dependant_asset.state else { continue; }; *remaining_dependencies = remaining_dependencies.saturating_sub(1); if *remaining_dependencies > 0 { continue; } } let AssetState::DependenciesPending { data: dependant_data, .. } = std::mem::replace(&mut dependant_asset.state, AssetState::Pending) else { unreachable!(); }; dependant_asset.state = AssetState::Ready { data: dependant_data }; self.curr_tick_events.push(Event::Stored( Id::from_valid(dependant_id), dependant_asset.label.clone().into_label(), )); all_dependants.append(&mut dependant_asset.waiting_dependants); } } } fn load_asset( &self, asset_ty_id: AssetTypeId, label: Label<'_>, asset_settings: Option, ) -> IdValid where AssetSettings: UnwindSafe + RefUnwindSafe + Send + Sync + Debug + 'static, { let id = match self.get_asset_by_label(&label) { Some((Some(_), id)) => { return id; } Some((None, id)) => id, None => IdValid::new(&self.store, &label), }; let mut metadata_lut = self.metadata_lut.borrow_mut(); let mut parent_exists = false; let parent_id = label.name.is_some().then(|| { let parent_label = Label { path: label.path.as_ref().into(), name: None, }; parent_exists = metadata_lut.contains_key(&parent_label); metadata_lut .entry(parent_label.to_owned()) .or_insert_with(|| Metadata { id: IdValid::new(&self.store, &parent_label), parent_id: None, }) .id }); metadata_lut.insert(label.to_owned(), Metadata { id, parent_id }); if parent_exists { return id; } self.add_import_work::(asset_ty_id, &label, asset_settings); id } fn add_import_work( &self, asset_ty_id: AssetTypeId, label: &Label<'_>, asset_settings: Option, ) where AssetSettings: UnwindSafe + RefUnwindSafe + Any + Send + Sync, { let Some(importer) = self.importers.get(&asset_ty_id).cloned() else { tracing::error!("No importer exists for asset"); return; }; let import_work_msg_sender = self.import_work_msg_sender.clone(); let asset_path = label.path.to_path_buf(); self.import_work_thread_pool.spawn(move || { let result = catch_unwind(|| { if let Err(err) = importer.call( import_work_msg_sender, asset_path.as_path(), asset_settings .as_ref() .map(|asset_settings| asset_settings as &(dyn Any + Send + Sync)), ) { tracing::error!( "Failed to import asset {}: {:#}", asset_path.display(), crate::Error::new(err) ); } }); if let Err(panic_msg) = result { let panic_msg = if let Some(msg) = panic_msg.downcast_ref::<&str>() { msg } else if let Some(msg) = panic_msg.downcast_ref::() { msg.as_str() } else { "(unknown panic payload type)" }; tracing::error!( "Failed to import asset {}: Thread panicked: {panic_msg}", asset_path.display(), ); } }); } fn get_asset_by_label( &self, label: &Label<'_>, ) -> Option<(Option<&StoredAsset>, IdValid)> { let metadata_lut = self.metadata_lut.borrow(); let id = metadata_lut.get(label)?.id; Some((self.store.get(id.index), id)) } fn get_asset_by_label_mut( &mut self, label: &Label<'_>, ) -> Option<(Option<&mut StoredAsset>, IdValid)> { let metadata_lut = self.metadata_lut.get_mut(); let id = metadata_lut.get(label)?.id; Some((self.store.get_mut(id.index), id)) } } impl Default for Assets { fn default() -> Self { Self::with_capacity(0) } } pub struct Submitter<'path> { import_work_msg_sender: MpscSender, asset_path: &'path Path, } impl Submitter<'_> { pub fn submit_store(&self, asset: AssetT) where AssetT: Asset, { let label = LabelOwned { path: self.asset_path.into(), name: None, }; let _ = self.import_work_msg_sender.send(ImportWorkMessage::Store { do_store: Box::new(move |assets, label| { assets.store_with_label::(label.as_label(), asset); }), label, }); } pub fn submit_store_named(&self, name: impl AsRef, asset: AssetT) where AssetT: Asset, { let label = LabelOwned { path: self.asset_path.into(), name: Some(name.as_ref().into()), }; let _ = self.import_work_msg_sender.send(ImportWorkMessage::Store { do_store: Box::new(move |assets, label| { assets.store_with_label::(label.as_label(), asset); }), label, }); } } /// Asset handle. #[derive(Debug)] pub struct Handle { id: Id, _pd: PhantomData, } impl Handle { #[must_use] pub fn invalid() -> Self { Self::from_id(Id::invalid()) } #[must_use] pub fn from_id(id: Id) -> Self { Self { id, _pd: PhantomData } } #[must_use] pub fn id(&self) -> Id { self.id } } impl Clone for Handle { fn clone(&self) -> Self { Self { id: self.id, _pd: PhantomData } } } /// Asset ID. #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] pub struct Id { inner: IdInner, } impl Id { #[must_use] pub fn invalid() -> Self { Self { inner: IdInner::Invalid } } fn into_valid(self) -> Option { match self.inner { IdInner::Valid(id) => Some(id), IdInner::Invalid => None, } } fn from_valid(id: IdValid) -> Self { Self { inner: IdInner::Valid(id) } } } #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] enum IdInner { Valid(IdValid), Invalid, } #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] struct IdValid { index: store::ReservedIndex, } impl IdValid { fn new(store: &store::Store, label: &Label<'_>) -> Self { tracing::trace!("Generating valid ID for {label}"); let index = store.reserve_index(); Self { index } } } #[derive(Debug, Default, Sole)] pub struct Events { last_tick: Vec, } impl Events { pub fn iter_last_tick(&self) -> impl Iterator { self.last_tick.iter() } } /// Asset event. #[derive(Debug)] pub enum Event { /// Asset stored. Stored(Id, Label<'static>), } #[derive(Debug, thiserror::Error)] enum ImporterError { #[error("Settings has a incorrect type")] IncorrectAssetSettingsType(PathBuf), #[error(transparent)] Other(Box), } #[derive(Debug, Clone)] struct WrappedImporterFn { wrapper_func: fn( MpscSender, &Path, Option<&(dyn Any + Send + Sync)>, ) -> Result<(), ImporterError>, } impl WrappedImporterFn { fn new(inner_func_param: InnerFunc) -> Self where InnerFunc: Fn(&mut Submitter<'_>, &Path, Option<&AssetSettings>) -> Result<(), Err>, AssetSettings: UnwindSafe + RefUnwindSafe + 'static, Err: std::error::Error + Send + Sync + 'static, { assert_eq!(size_of::(), 0); let wrapper_func = |import_work_msg_sender: MpscSender, asset_path: &Path, asset_settings: Option<&(dyn Any + Send + Sync)>| { let inner_func = unsafe { std::mem::zeroed::() }; let asset_settings = asset_settings .map(|asset_settings| { asset_settings .downcast_ref::() .ok_or_else(|| { ImporterError::IncorrectAssetSettingsType( asset_path.to_path_buf(), ) }) }) .transpose()?; inner_func( &mut Submitter { import_work_msg_sender, asset_path }, asset_path, asset_settings, ) .map_err(|err| ImporterError::Other(Box::new(err)))?; Ok(()) }; std::mem::forget(inner_func_param); Self { wrapper_func } } fn call( &self, import_work_msg_sender: MpscSender, asset_path: &Path, asset_settings: Option<&(dyn Any + Send + Sync)>, ) -> Result<(), ImporterError> { (self.wrapper_func)(import_work_msg_sender, asset_path, asset_settings) } } #[derive(Debug, Default)] pub(crate) struct Extension { pub assets: Assets, } impl crate::ecs::extension::Extension for Extension { fn collect(self, mut collector: crate::ecs::extension::Collector<'_>) { let _ = collector.add_sole(self.assets); let _ = collector.add_sole(Events::default()); collector.spawn_declared_entity(&HANDLE_ASSETS_PHASE); collector.add_system(*HANDLE_ASSETS_PHASE, add_received_assets); } } fn add_received_assets(mut assets: Single, mut events: Single) { let Ok(assets) = assets.get_mut() else { unreachable!(); }; let Ok(events) = events.get_mut() else { unreachable!(); }; std::mem::swap(&mut events.last_tick, &mut assets.curr_tick_events); assets.curr_tick_events.clear(); while let Ok(import_work_msg) = assets.import_work_msg_receiver.try_recv() { match import_work_msg { ImportWorkMessage::Store { do_store, label } => { do_store(assets, label); } } } } enum ImportWorkMessage { Store { do_store: Box, label: LabelOwned, }, } #[derive(Debug)] struct Metadata { id: IdValid, parent_id: Option, } #[derive(Debug)] enum AssetState { Ready { data: Box, }, Pending, DependenciesPending { data: Box, remaining_dependencies: usize, }, } #[derive(Debug)] struct StoredAsset { state: AssetState, label: LabelOwned, waiting_dependants: Vec, } #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] struct AssetTypeId { inner: TypeId, } impl AssetTypeId { fn of() -> Self { Self { inner: TypeId::of::() } } } #[derive(Debug, Clone, Eq)] struct LabelOwned { path: PathBuf, name: Option, } impl LabelOwned { #[must_use] fn as_label(&self) -> Label<'_> { Label { path: self.path.as_path().into(), name: self.name.as_deref().map(Into::into), } } fn into_label(self) -> Label<'static> { Label { path: self.path.into(), name: self.name.map(Into::into), } } } impl Display for LabelOwned { fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { as Display>::fmt(&self.as_label(), formatter) } } impl PartialEq for LabelOwned { fn eq(&self, other: &Self) -> bool { self.as_label() == other.as_label() } } impl Hash for LabelOwned { fn hash(&self, state: &mut HasherT) { as Hash>::hash(&self.as_label(), state); } } impl hashbrown::Equivalent for Label<'_> { fn equivalent(&self, key: &LabelOwned) -> bool { self == &key.as_label() } } fn calc_import_work_thread_cnt() -> usize { available_parallelism() .map(|avail_parallelism| avail_parallelism.get().div_ceil(2)) .inspect_err(|err| { tracing::warn!( concat!( "Failed to get available parallelism for calculating import work ", "thread count. {} threads will be used. Error: {}" ), BACKUP_IMPORT_WORK_THREAD_CNT, err ); }) .unwrap_or(BACKUP_IMPORT_WORK_THREAD_CNT) } mod store { use std::num::NonZero; use std::sync::atomic::{AtomicUsize, Ordering}; use super::StoredAsset; #[derive(Debug)] pub struct Store { store: Vec>, next_index: AtomicUsize, } impl Store { pub fn with_capacity(capacity: usize) -> Self { Self { store: Vec::with_capacity(capacity), next_index: AtomicUsize::new(1), } } pub fn get(&self, index: ReservedIndex) -> Option<&StoredAsset> { self.store.get(index.0.get())?.as_ref() } pub fn get_mut(&mut self, index: ReservedIndex) -> Option<&mut StoredAsset> { self.store.get_mut(index.0.get())?.as_mut() } pub fn get_or_insert_with( &mut self, index: ReservedIndex, func: impl FnOnce() -> StoredAsset, ) -> &mut StoredAsset { let index = index.0.get(); let curr_len = self.store.len(); if curr_len <= index { let missing = (index + 1) - curr_len; self.store .extend(std::iter::repeat_with(|| None).take(missing)); } self.store[index].get_or_insert_with(func) } pub fn reserve_index(&self) -> ReservedIndex { let index = self.next_index.fetch_add(1, Ordering::Relaxed); let index = NonZero::::new(index) .expect("Asset index wrapped around to zero. Too many assets"); ReservedIndex(index) } } #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] pub struct ReservedIndex(NonZero); }