distill_loader/
handle.rs

1use std::{
2    collections::{HashMap, HashSet},
3    fmt::Debug,
4    hash::Hash,
5    marker::PhantomData,
6    sync::{
7        atomic::{AtomicU64, Ordering},
8        Arc, Mutex, RwLock,
9    },
10};
11
12use crossbeam_channel::{unbounded, Receiver, Sender};
13use futures_core::future::{BoxFuture, Future};
14use serde::{
15    de::{self, Deserialize, Visitor},
16    ser::{self, Serialize, Serializer},
17};
18
19use crate::{
20    storage::{LoadStatus, LoaderInfoProvider},
21    AssetRef, AssetUuid, LoadHandle, Loader,
22};
23
24/// Operations on an asset reference.
25#[derive(Debug)]
26pub enum RefOp {
27    Decrease(LoadHandle),
28    Increase(LoadHandle),
29    IncreaseUuid(AssetUuid),
30}
31
32pub fn process_ref_ops(loader: &Loader, rx: &Receiver<RefOp>) {
33    loop {
34        match rx.try_recv() {
35            Err(_) => break,
36            Ok(RefOp::Decrease(handle)) => loader.remove_ref(handle),
37            Ok(RefOp::Increase(handle)) => {
38                loader.add_ref_handle(handle);
39            }
40            Ok(RefOp::IncreaseUuid(uuid)) => {
41                loader.add_ref(uuid);
42            }
43        }
44    }
45}
46
47/// Keeps track of whether a handle ref is a strong, weak or "internal" ref
48#[derive(Debug)]
49pub enum HandleRefType {
50    /// Strong references decrement the count on drop
51    Strong(Sender<RefOp>),
52    /// Weak references do nothing on drop.
53    Weak(Sender<RefOp>),
54    /// Internal references do nothing on drop, but turn into Strong references on clone.
55    /// Should only be used for references stored in loaded assets to avoid self-referencing
56    Internal(Sender<RefOp>),
57    /// Implementation detail, used when changing state in this enum
58    None,
59}
60
61struct HandleRef {
62    id: LoadHandle,
63    ref_type: HandleRefType,
64}
65impl PartialEq for HandleRef {
66    fn eq(&self, other: &Self) -> bool {
67        self.id.eq(&other.id)
68    }
69}
70impl Hash for HandleRef {
71    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
72        self.id.hash(state)
73    }
74}
75impl Eq for HandleRef {}
76impl Debug for HandleRef {
77    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
78        self.id.fmt(f)
79    }
80}
81
82impl Drop for HandleRef {
83    fn drop(&mut self) {
84        use HandleRefType::*;
85        self.ref_type = match std::mem::replace(&mut self.ref_type, None) {
86            Strong(sender) => {
87                let _ = sender.send(RefOp::Decrease(self.id));
88                Weak(sender)
89            }
90            r => r,
91        };
92    }
93}
94
95impl Clone for HandleRef {
96    fn clone(&self) -> Self {
97        use HandleRefType::*;
98        Self {
99            id: self.id,
100            ref_type: match &self.ref_type {
101                Internal(sender) | Strong(sender) => {
102                    let _ = sender.send(RefOp::Increase(self.id));
103                    Strong(sender.clone())
104                }
105                Weak(sender) => Weak(sender.clone()),
106                None => panic!("unexpected ref type in clone()"),
107            },
108        }
109    }
110}
111
112impl AssetHandle for HandleRef {
113    fn load_handle(&self) -> LoadHandle {
114        self.id
115    }
116}
117
118/// Handle to an asset.
119#[derive(Eq)]
120pub struct Handle<T: ?Sized> {
121    handle_ref: HandleRef,
122    marker: PhantomData<T>,
123}
124
125impl<T: ?Sized> PartialEq for Handle<T> {
126    fn eq(&self, other: &Self) -> bool {
127        self.handle_ref == other.handle_ref
128    }
129}
130
131impl<T: ?Sized> Clone for Handle<T> {
132    fn clone(&self) -> Self {
133        Self {
134            handle_ref: self.handle_ref.clone(),
135            marker: PhantomData,
136        }
137    }
138}
139
140impl<T: ?Sized> Hash for Handle<T> {
141    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
142        self.handle_ref.hash(state);
143    }
144}
145
146impl<T: ?Sized> Debug for Handle<T> {
147    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
148        f.debug_struct("Handle")
149            .field("handle_ref", &self.handle_ref)
150            .finish()
151    }
152}
153
154impl<T: ?Sized> From<GenericHandle> for Handle<T> {
155    fn from(handle: GenericHandle) -> Self {
156        Self {
157            handle_ref: handle.handle_ref,
158            marker: PhantomData,
159        }
160    }
161}
162
163impl<T> Handle<T> {
164    /// Creates a new handle with `HandleRefType::Strong`
165    pub fn new(chan: Sender<RefOp>, handle: LoadHandle) -> Self {
166        Self {
167            handle_ref: HandleRef {
168                id: handle,
169                ref_type: HandleRefType::Strong(chan),
170            },
171            marker: PhantomData,
172        }
173    }
174
175    /// Creates a new handle with `HandleRefType::Internal`
176    pub(crate) fn new_internal(chan: Sender<RefOp>, handle: LoadHandle) -> Self {
177        Self {
178            handle_ref: HandleRef {
179                id: handle,
180                ref_type: HandleRefType::Internal(chan),
181            },
182            marker: PhantomData,
183        }
184    }
185
186    pub fn asset<'a>(&self, storage: &'a impl TypedAssetStorage<T>) -> Option<&'a T> {
187        AssetHandle::asset(self, storage)
188    }
189}
190
191impl<T> AssetHandle for Handle<T> {
192    fn load_handle(&self) -> LoadHandle {
193        self.handle_ref.load_handle()
194    }
195}
196
197/// Handle to an asset whose type is unknown during loading.
198///
199/// This is returned by `Loader::load_asset_generic` for assets loaded by UUID.
200#[derive(Debug, Clone, PartialEq, Eq, Hash)]
201pub struct GenericHandle {
202    handle_ref: HandleRef,
203}
204
205impl GenericHandle {
206    /// Creates a new handle with `HandleRefType::Strong`
207    pub fn new(chan: Sender<RefOp>, handle: LoadHandle) -> Self {
208        Self {
209            handle_ref: HandleRef {
210                id: handle,
211                ref_type: HandleRefType::Strong(chan),
212            },
213        }
214    }
215
216    /// Creates a new handle with `HandleRefType::Internal`
217    pub(crate) fn new_internal(chan: Sender<RefOp>, handle: LoadHandle) -> Self {
218        Self {
219            handle_ref: HandleRef {
220                id: handle,
221                ref_type: HandleRefType::Internal(chan),
222            },
223        }
224    }
225}
226
227impl AssetHandle for GenericHandle {
228    fn load_handle(&self) -> LoadHandle {
229        self.handle_ref.load_handle()
230    }
231}
232
233impl<T: ?Sized> From<Handle<T>> for GenericHandle {
234    fn from(handle: Handle<T>) -> Self {
235        Self {
236            handle_ref: handle.handle_ref,
237        }
238    }
239}
240
241/// Handle to an asset that does not prevent the asset from being unloaded.
242///
243/// Weak handles are primarily used when you want to use something that is already loaded.
244///
245/// For example, a strong handle to an asset may be guaranteed to exist elsewhere in the program,
246/// and so you can simply get and use a weak handle to that asset in other parts of your code. This
247/// removes reference counting overhead, but also ensures that the system which uses the weak handle
248/// is not in control of when to unload the asset.
249#[derive(Clone, Eq, Hash, PartialEq, Debug)]
250pub struct WeakHandle {
251    id: LoadHandle,
252}
253
254impl WeakHandle {
255    pub fn new(handle: LoadHandle) -> Self {
256        WeakHandle { id: handle }
257    }
258}
259
260impl AssetHandle for WeakHandle {
261    fn load_handle(&self) -> LoadHandle {
262        self.id
263    }
264}
265
266tokio::task_local! {
267    static LOADER: &'static dyn LoaderInfoProvider;
268    static REFOP_SENDER: Sender<RefOp>;
269}
270
271/// Used to make some limited Loader interactions available to `serde` Serialize/Deserialize
272/// implementations by using thread-local storage. Required to support Serialize/Deserialize of Handle.
273pub struct SerdeContext;
274impl SerdeContext {
275    pub fn with_active<R>(f: impl FnOnce(&dyn LoaderInfoProvider, &Sender<RefOp>) -> R) -> R {
276        LOADER.with(|l| REFOP_SENDER.with(|r| f(*l, &r)))
277    }
278
279    pub async fn with<F>(loader: &dyn LoaderInfoProvider, sender: Sender<RefOp>, f: F) -> F::Output
280    where
281        F: Future,
282    {
283        // The loader lifetime needs to be transmuted to 'static to be able to be stored in task_local.
284        // This is safe since SerdeContext's lifetime cannot be shorter than the opened scope, and the loader
285        // must live at least as long.
286        let loader = unsafe {
287            std::mem::transmute::<&dyn LoaderInfoProvider, &'static dyn LoaderInfoProvider>(loader)
288        };
289
290        LOADER.scope(loader, REFOP_SENDER.scope(sender, f)).await
291    }
292}
293
294/// This context can be used to maintain AssetUuid references through a serialize/deserialize cycle
295/// even if the LoadHandles produced are invalid. This is useful when a loader is not
296/// present, such as when processing in the Distill Daemon.
297struct DummySerdeContext {
298    maps: RwLock<DummySerdeContextMaps>,
299    current: Mutex<DummySerdeContextCurrent>,
300    ref_sender: Sender<RefOp>,
301    handle_gen: AtomicU64,
302}
303
304struct DummySerdeContextMaps {
305    uuid_to_load: HashMap<AssetRef, LoadHandle>,
306    load_to_uuid: HashMap<LoadHandle, AssetRef>,
307}
308
309struct DummySerdeContextCurrent {
310    current_serde_dependencies: HashSet<AssetRef>,
311    current_serde_asset: Option<AssetUuid>,
312}
313
314impl DummySerdeContext {
315    pub fn new() -> Self {
316        let (tx, _) = unbounded();
317        Self {
318            maps: RwLock::new(DummySerdeContextMaps {
319                uuid_to_load: HashMap::default(),
320                load_to_uuid: HashMap::default(),
321            }),
322            current: Mutex::new(DummySerdeContextCurrent {
323                current_serde_dependencies: HashSet::new(),
324                current_serde_asset: None,
325            }),
326            ref_sender: tx,
327            handle_gen: AtomicU64::new(1),
328        }
329    }
330}
331
332impl LoaderInfoProvider for DummySerdeContext {
333    fn get_load_handle(&self, asset_ref: &AssetRef) -> Option<LoadHandle> {
334        let mut maps = self.maps.write().unwrap();
335        let maps = &mut *maps;
336        let uuid_to_load = &mut maps.uuid_to_load;
337        let load_to_uuid = &mut maps.load_to_uuid;
338
339        let entry = uuid_to_load.entry(asset_ref.clone());
340        let handle = entry.or_insert_with(|| {
341            let new_id = self.handle_gen.fetch_add(1, Ordering::Relaxed);
342            let handle = LoadHandle(new_id);
343            load_to_uuid.insert(handle, asset_ref.clone());
344            handle
345        });
346
347        Some(*handle)
348    }
349
350    fn get_asset_id(&self, load: LoadHandle) -> Option<AssetUuid> {
351        let maps = self.maps.read().unwrap();
352        let maybe_asset = maps.load_to_uuid.get(&load).cloned();
353        if let Some(asset_ref) = maybe_asset.as_ref() {
354            let mut current = self.current.lock().unwrap();
355            if let Some(ref current_serde_id) = current.current_serde_asset {
356                if AssetRef::Uuid(*current_serde_id) != *asset_ref
357                    && *asset_ref != AssetRef::Uuid(AssetUuid::default())
358                {
359                    current.current_serde_dependencies.insert(asset_ref.clone());
360                }
361            }
362        }
363        if let Some(AssetRef::Uuid(uuid)) = maybe_asset {
364            Some(uuid)
365        } else {
366            None
367        }
368    }
369}
370struct DummySerdeContextHandle {
371    dummy: Arc<DummySerdeContext>,
372}
373impl<'a> distill_core::importer_context::ImporterContextHandle for DummySerdeContextHandle {
374    fn scope<'s>(&'s self, fut: BoxFuture<'s, ()>) -> BoxFuture<'s, ()> {
375        let sender = self.dummy.ref_sender.clone();
376        let loader = &*self.dummy;
377        Box::pin(SerdeContext::with(loader, sender, fut))
378    }
379
380    fn resolve_ref(&mut self, asset_ref: &AssetRef, asset: AssetUuid) {
381        let new_ref = AssetRef::Uuid(asset);
382        let mut maps = self.dummy.maps.write().unwrap();
383        if let Some(handle) = maps.uuid_to_load.get(asset_ref) {
384            let handle = *handle;
385            maps.load_to_uuid.insert(handle, new_ref.clone());
386            maps.uuid_to_load.insert(new_ref, handle);
387        }
388    }
389
390    /// Begin gathering dependencies for an asset
391    fn begin_serialize_asset(&mut self, asset: AssetUuid) {
392        let mut current = self.dummy.current.lock().unwrap();
393        if current.current_serde_asset.is_some() {
394            panic!("begin_serialize_asset when current_serde_asset is already set");
395        }
396        current.current_serde_asset = Some(asset);
397    }
398
399    /// Finish gathering dependencies for an asset
400    fn end_serialize_asset(&mut self, _asset: AssetUuid) -> HashSet<AssetRef> {
401        let mut current = self.dummy.current.lock().unwrap();
402        if current.current_serde_asset.is_none() {
403            panic!("end_serialize_asset when current_serde_asset is not set");
404        }
405        current.current_serde_asset = None;
406        std::mem::replace(&mut current.current_serde_dependencies, HashSet::new())
407    }
408}
409
410/// Register this context with AssetDaemon to add serde support for Handle.
411pub struct HandleSerdeContextProvider;
412impl distill_core::importer_context::ImporterContext for HandleSerdeContextProvider {
413    fn handle(&self) -> Box<dyn distill_core::importer_context::ImporterContextHandle> {
414        let dummy = Arc::new(DummySerdeContext::new());
415        Box::new(DummySerdeContextHandle { dummy })
416    }
417}
418
419fn serialize_handle<S>(load: LoadHandle, serializer: S) -> Result<S::Ok, S::Error>
420where
421    S: Serializer,
422{
423    SerdeContext::with_active(|loader, _| {
424        use ser::SerializeSeq;
425        let uuid: AssetUuid = loader.get_asset_id(load).unwrap_or_default();
426        let mut seq = serializer.serialize_seq(Some(uuid.0.len()))?;
427        for element in &uuid.0 {
428            seq.serialize_element(element)?;
429        }
430        seq.end()
431    })
432}
433impl<T> Serialize for Handle<T> {
434    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
435    where
436        S: Serializer,
437    {
438        serialize_handle(self.handle_ref.id, serializer)
439    }
440}
441impl Serialize for GenericHandle {
442    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
443    where
444        S: Serializer,
445    {
446        serialize_handle(self.handle_ref.id, serializer)
447    }
448}
449
450fn get_handle_ref(asset_ref: AssetRef) -> (LoadHandle, Sender<RefOp>) {
451    SerdeContext::with_active(|loader, sender| {
452        let handle = if asset_ref == AssetRef::Uuid(AssetUuid::default()) {
453            LoadHandle(0)
454        } else {
455            loader
456                .get_load_handle(&asset_ref)
457                .unwrap_or_else(|| panic!("Handle for AssetUuid {:?} was not present when deserializing a Handle. This indicates missing dependency metadata, and can be caused by dependency cycles.", asset_ref))
458        };
459        (handle, sender.clone())
460    })
461}
462
463impl<'de, T> Deserialize<'de> for Handle<T> {
464    fn deserialize<D>(deserializer: D) -> Result<Handle<T>, D::Error>
465    where
466        D: de::Deserializer<'de>,
467    {
468        let asset_ref = if deserializer.is_human_readable() {
469            deserializer.deserialize_any(AssetRefVisitor)?
470        } else {
471            deserializer.deserialize_seq(AssetRefVisitor)?
472        };
473        let (handle, sender) = get_handle_ref(asset_ref);
474        Ok(Handle::new_internal(sender, handle))
475    }
476}
477
478impl<'de> Deserialize<'de> for GenericHandle {
479    fn deserialize<D>(deserializer: D) -> Result<GenericHandle, D::Error>
480    where
481        D: de::Deserializer<'de>,
482    {
483        let asset_ref = if deserializer.is_human_readable() {
484            deserializer.deserialize_any(AssetRefVisitor)?
485        } else {
486            deserializer.deserialize_seq(AssetRefVisitor)?
487        };
488        let (handle, sender) = get_handle_ref(asset_ref);
489        Ok(GenericHandle::new_internal(sender, handle))
490    }
491}
492
493struct AssetRefVisitor;
494
495impl<'de> Visitor<'de> for AssetRefVisitor {
496    type Value = AssetRef;
497
498    fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
499        formatter.write_str("an array of 16 u8")
500    }
501
502    fn visit_newtype_struct<D>(self, deserializer: D) -> Result<Self::Value, D::Error>
503    where
504        D: de::Deserializer<'de>,
505    {
506        deserializer.deserialize_seq(self)
507    }
508
509    fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
510    where
511        A: de::SeqAccess<'de>,
512    {
513        use de::Error;
514        let mut uuid: [u8; 16] = Default::default();
515        for (i, uuid_byte) in uuid.iter_mut().enumerate() {
516            if let Some(byte) = seq.next_element::<u8>()? {
517                *uuid_byte = byte;
518            } else {
519                return Err(A::Error::custom(format!(
520                    "expected byte at element {} when deserializing handle",
521                    i
522                )));
523            }
524        }
525        if seq.next_element::<u8>()?.is_some() {
526            return Err(A::Error::custom(
527                "too many elements when deserializing handle",
528            ));
529        }
530        Ok(AssetRef::Uuid(AssetUuid(uuid)))
531    }
532
533    fn visit_str<E>(self, v: &str) -> Result<Self::Value, E>
534    where
535        E: de::Error,
536    {
537        use std::str::FromStr;
538        match std::path::PathBuf::from_str(v) {
539            Ok(path) => {
540                if let Ok(uuid) = uuid::Uuid::parse_str(&path.to_string_lossy()) {
541                    Ok(AssetRef::Uuid(AssetUuid(*uuid.as_bytes())))
542                } else {
543                    Ok(AssetRef::Path(path))
544                }
545            }
546            Err(err) => Err(E::custom(format!(
547                "failed to parse Handle string: {:?}",
548                err
549            ))),
550        }
551    }
552
553    fn visit_bytes<E>(self, v: &[u8]) -> Result<Self::Value, E>
554    where
555        E: de::Error,
556    {
557        if v.len() != 16 {
558            Err(E::custom(format!(
559                "byte array len == {}, expected {}",
560                v.len(),
561                16
562            )))
563        } else {
564            let mut a = <[u8; 16]>::default();
565            a.copy_from_slice(v);
566            Ok(AssetRef::Uuid(AssetUuid(a)))
567        }
568    }
569}
570
571/// Implementors of [`crate::storage::AssetStorage`] can implement this trait to enable convenience
572/// functions on the common [`AssetHandle`] trait, which is implemented by all handle types.
573pub trait TypedAssetStorage<A> {
574    /// Returns the asset for the given handle, or `None` if has not completed loading.
575    ///
576    /// # Parameters
577    ///
578    /// * `handle`: Handle of the asset.
579    ///
580    /// # Type Parameters
581    ///
582    /// * `T`: Asset handle type.
583    fn get<T: AssetHandle>(&self, handle: &T) -> Option<&A>;
584
585    /// Returns the version of a loaded asset, or `None` if has not completed loading.
586    ///
587    /// # Parameters
588    ///
589    /// * `handle`: Handle of the asset.
590    ///
591    /// # Type Parameters
592    ///
593    /// * `T`: Asset handle type.
594    fn get_version<T: AssetHandle>(&self, handle: &T) -> Option<u32>;
595
596    /// Returns the loaded asset and its version, or `None` if has not completed loading.
597    ///
598    /// # Parameters
599    ///
600    /// * `handle`: Handle of the asset.
601    ///
602    /// # Type Parameters
603    ///
604    /// * `T`: Asset handle type.
605    fn get_asset_with_version<T: AssetHandle>(&self, handle: &T) -> Option<(&A, u32)>;
606}
607
608/// The contract of an asset handle.
609///
610/// There are two types of asset handles:
611///
612/// * **Typed -- `Handle<T>`:** When the asset's type is known when loading.
613/// * **Generic -- `GenericHandle`:** When only the asset's UUID is known when loading.
614pub trait AssetHandle {
615    /// Returns the load status of the asset.
616    ///
617    /// # Parameters
618    ///
619    /// * `loader`: Loader that is loading the asset.
620    ///
621    /// # Type Parameters
622    ///
623    /// * `L`: Asset loader type.
624    fn load_status(&self, loader: &Loader) -> LoadStatus {
625        loader.get_load_status(self.load_handle())
626    }
627
628    /// Returns an immutable reference to the asset if it is committed.
629    ///
630    /// # Parameters
631    ///
632    /// * `storage`: Asset storage.
633    fn asset<'a, T, S: TypedAssetStorage<T>>(&self, storage: &'a S) -> Option<&'a T>
634    where
635        Self: Sized,
636    {
637        storage.get(self)
638    }
639
640    /// Returns the version of the asset if it is committed.
641    ///
642    /// # Parameters
643    ///
644    /// * `storage`: Asset storage.
645    fn asset_version<T, S: TypedAssetStorage<T>>(&self, storage: &S) -> Option<u32>
646    where
647        Self: Sized,
648    {
649        storage.get_version(self)
650    }
651
652    /// Returns the asset with the given version if it is committed.
653    ///
654    /// # Parameters
655    ///
656    /// * `storage`: Asset storage.
657    fn asset_with_version<'a, T, S: TypedAssetStorage<T>>(
658        &self,
659        storage: &'a S,
660    ) -> Option<(&'a T, u32)>
661    where
662        Self: Sized,
663    {
664        storage.get_asset_with_version(self)
665    }
666
667    /// Downgrades this handle into a `WeakHandle`.
668    ///
669    /// Be aware that if there are no longer any strong handles to the asset, then the underlying
670    /// asset may be freed at any time.
671    fn downgrade(&self) -> WeakHandle {
672        WeakHandle::new(self.load_handle())
673    }
674
675    /// Returns the `LoadHandle` of this asset handle.
676    fn load_handle(&self) -> LoadHandle;
677}
distill_loader/handle.rs

distill_loader/
handle.rs
