lyra-engine/lyra-resource/src/resource.rs

use std::{any::{Any, TypeId}, marker::PhantomData, ops::{Deref, DerefMut}, sync::{Arc, Condvar, Mutex, RwLock, RwLockReadGuard, RwLockWriteGuard}};

use lyra_ecs::Component;
use crate::{loader::LoaderError, lyra_engine, DependencyState};
use uuid::Uuid;

use crate::ResourceStorage;

pub fn optionally_add_to_dep<R: ResourceData>(deps: &mut Vec<UntypedResHandle>, handle: &Option<ResHandle<R>>) {
    if let Some(h) = handle {
        deps.push(h.untyped_clone());
    }
}

/// A trait that that each resource type should implement.
pub trait ResourceData: Send + Sync + Any + 'static {
    fn as_any(&self) -> &dyn Any;
    fn as_any_mut(&mut self) -> &mut dyn Any;

    /// Collect the dependencies of the Resource.
    fn dependencies(&self) -> Vec<UntypedResHandle>;

    /// Recursively collect the dependencies of the Resource.
    /// 
    /// If a dependency has a child dependency, it will not show up in this list until its
    /// parent is loaded.
    fn recur_dependencies(&self) -> Vec<UntypedResHandle> {
        let deps = self.dependencies();
        let mut all_deps = deps.clone();

        for dep in deps.into_iter() {
            let dep = dep.read();
            match &dep.state {
                ResourceState::Ready(data) => {
                    let mut deps_dep = data.dependencies();
                    all_deps.append(&mut deps_dep);
                },
                _ => {}
            }
        }

        all_deps
    }
}

//impl<T: Send + Sync + Reflect> ResourceData for T { }

pub enum ResourceState {
    Loading,
    Error(Arc<LoaderError>),
    Ready(Box<dyn ResourceData>),
}

impl ResourceState {
    /// Returns a boolean indicating if the state of still loading
    pub fn is_loading(&self) -> bool {
        matches!(self, ResourceState::Loading)
    }

    pub fn is_error(&self) -> bool {
        matches!(self, ResourceState::Error(_))
    }

    pub fn is_ready(&self) -> bool {
        matches!(self, ResourceState::Ready(_))
    }
}

pub struct ResourceDataRef<'a, T> {
    guard: std::sync::RwLockReadGuard<'a, UntypedResource>,
    _marker: PhantomData<T>,
}

impl<'a, T: 'static> std::ops::Deref for ResourceDataRef<'a, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        match &self.guard.state {
            ResourceState::Ready(d) => {
                // for some reason, if I didn't use `.as_ref`, the downcast would fail.
                let d = d.as_ref().as_any();
                d.downcast_ref::<T>().unwrap()
            },
            _ => unreachable!() // ResHandler::data_ref shouldn't allow this to run
        }
    }
}

pub struct UntypedResource {
    pub(crate) version: usize,
    pub(crate) state: ResourceState,
    uuid: Uuid,
    path: Option<String>,
    pub(crate) is_watched: bool,
    /// can be used to wait for the resource to load.
    pub(crate) condvar: Arc<(Mutex<bool>, Condvar)>,
}

#[derive(Clone, Component)]
pub struct UntypedResHandle{
    pub(crate) res: Arc<RwLock<UntypedResource>>,
    #[allow(dead_code)]
    tyid: TypeId,
}

impl UntypedResHandle {
    pub fn new(res: UntypedResource, tyid: TypeId) -> Self {
        Self {
            res: Arc::new(RwLock::new(res)),
            tyid
        }
    }

    pub fn read(&self) -> RwLockReadGuard<UntypedResource> {
        self.res.read().unwrap()
    }

    pub fn write(&self) -> RwLockWriteGuard<UntypedResource> {
        self.res.write().unwrap()
    }

    /// Returns a boolean indicating if this resource's path is being watched.
    pub fn is_watched(&self) -> bool {
        let d = self.read();
        d.is_watched
    }

    /// Returns a boolean indicating if this resource is loaded
    pub fn is_loaded(&self) -> bool {
        let d = self.read();
        matches!(d.state, ResourceState::Ready(_))
    }

    /// Returns the uuid of the resource.
    pub fn uuid(&self) -> Uuid {
        let d = self.read();
        d.uuid
    }

    pub fn path(&self) -> Option<String> {
        let d = self.read();
        d.path.clone()
    }

    /// Retrieves the current version of the resource. This gets incremented when the resource
    /// is reloaded.
    pub fn version(&self) -> usize {
        let d = self.read();
        d.version
    }

    /// Wait for the resource to be loaded, not including its dependencies
    /// (see[`UntypedResHandle::wait_recurse_dependencies_load`]).
    /// 
    /// This blocks the thread without consuming CPU time; its backed by a
    /// [`Condvar`](std::sync::Condvar).
    pub fn wait_for_load(&self) {
        let d = self.read();
        
        if matches!(d.state, ResourceState::Ready(_)) {
            return;
        }

        let cv = d.condvar.clone();
        drop(d);

        let l = cv.0.lock().unwrap();
        let _unused = cv.1.wait(l).unwrap();
    }

    /// Wait for the entire resource, including its dependencies to be loaded.
    /// 
    /// This blocks the thread without consuming CPU time; its backed by a
    /// [`Condvar`](std::sync::Condvar).
    pub fn wait_recurse_dependencies_load(&self) {
        self.wait_for_load();

        let res = self.read();
        match &res.state {
            ResourceState::Ready(data) => {
                // `recur_dependencies` wont return resources that are not loaded in yet
                // if we did not check if the resource was finished loading, we could miss
                // waiting for some resources and finish early.
                while self.recurse_dependency_state().is_loading() {
                    for dep in data.recur_dependencies() {
                        dep.wait_for_load();
                    }
                }
            },
            _ => unreachable!() // the self.wait_for_load ensures that the state is ready
        }
    }

    /// Recursively get the state of the dependencies.
    /// 
    /// This doesn't return any resource data, it can be used to check if the resource and its
    /// dependencies are loaded.
    pub fn recurse_dependency_state(&self) -> DependencyState {
        DependencyState::from_res_recurse(self)
    }

    pub fn as_typed<T: ResourceData>(&self) -> ResHandle<T> {
        ResHandle {
            handle: self.clone(),
            _marker: PhantomData::<T>,
        }
    }
}

/// A handle to a resource. 
/// 
/// # Note
/// This struct has an inner [`RwLock`] to the resource data, so most methods may be blocking.
/// However, the only times it will be blocking is if another thread is reloading the resource
/// and has a write lock on the data. This means that most of the time, it is not blocking.
#[derive(Component)]
pub struct ResHandle<T: ResourceData> {
    pub(crate) handle: UntypedResHandle,
    _marker: PhantomData<T>,
}

impl<T: ResourceData> Clone for ResHandle<T> {
    fn clone(&self) -> Self {
        Self {
            handle: self.handle.clone(),
            _marker: PhantomData::<T>
        }
    }
}

impl<T: ResourceData> Deref for ResHandle<T> {
    type Target = UntypedResHandle;

    fn deref(&self) -> &Self::Target {
        &self.handle
    }
}

impl<T: ResourceData> DerefMut for ResHandle<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.handle
    }
}

impl<T: ResourceData> ResHandle<T> {
    pub fn new_loading(path: Option<&str>) -> Self {
        let res_version = UntypedResource {
            version: 0,
            path: path.map(str::to_string),
            state: ResourceState::Loading,
            uuid: Uuid::new_v4(),
            is_watched: false,
            condvar: Arc::new((Mutex::new(false), Condvar::new())),
        };

        Self {
            handle: UntypedResHandle::new(res_version, TypeId::of::<T>()),
            _marker: PhantomData::<T>,
        }
    }

    /// Create the resource with data, its assumed the state is `Ready`
    pub fn new_ready(path: Option<&str>, data: T) -> Self {
        let han = Self::new_loading(path);
        han.set_state(ResourceState::Ready(Box::new(data)));
        han
    }

    /// Retrieve an untyped clone of the handle
    pub fn untyped_clone(&self) -> UntypedResHandle {
        self.handle.clone()
    }

    /// Get a reference to the data in the resource
    pub fn data_ref<'a>(&'a self) -> Option<ResourceDataRef<'a, T>> {
        if self.is_loaded() {
            let d = self.handle.read();
            Some(ResourceDataRef {
                guard: d,
                _marker: PhantomData::<T>
            })
        } else {
            None
        }
    }
}

impl<T: ResourceData> ResourceStorage for ResHandle<T> {
    fn as_any(&self) -> &dyn Any {
        self
    }

    fn as_any_mut(&mut self) -> &mut dyn Any {
        self
    }

    fn as_arc_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync> {
        self.clone()
    }

    fn as_box_any(self: Box<Self>) -> Box<dyn Any + Send + Sync> {
        self
    }

    fn set_watched(&self, watched: bool) {
        let mut d = self.handle.write();
        d.is_watched = watched;
    }

    fn version(&self) -> usize {
        self.handle.version()
    }

    fn uuid(&self) -> Uuid {
        self.handle.uuid()
    }

    fn path(&self) -> Option<String> {
        self.handle.path()
    }

    fn is_watched(&self) -> bool {
        self.handle.is_watched()
    }

    fn is_loaded(&self) -> bool {
        self.handle.is_loaded()
    }

    fn set_state(&self, new: ResourceState) {
        let mut d = self.handle.write();
        d.state = new;
    }
}

#[cfg(test)]
mod tests {
    use std::{path::PathBuf, str::FromStr, sync::Arc};

    use async_std::task;
    use instant::Duration;
    use rand::Rng;

    use crate::{loader::ResourceLoader, ResHandle, ResourceData, ResourceManager};

    #[allow(dead_code)]
    struct SimpleDepend {
        file_name: String,
        ext: String,
    }

    impl ResourceData for SimpleDepend {
        fn as_any(&self) -> &dyn std::any::Any {
            self
        }

        fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
            self
        }

        fn dependencies(&self) -> Vec<crate::UntypedResHandle> {
            vec![]
        }
    }

    #[derive(Default)]
    struct SlowSimpleDependLoader;

    impl ResourceLoader for SlowSimpleDependLoader {
        fn extensions(&self) -> &[&str] {
            &["txt", "buf"]
        }

        fn mime_types(&self) -> &[&str] {
            &[]
        }

        fn load(&self, _: crate::ResourceManager, path: &str) -> crate::loader::PinedBoxLoaderFuture {
            let path = path.to_string();
            Box::pin(async move {
                let path = PathBuf::from_str(&path).unwrap();

                let file_name = path.file_name()
                    .and_then(|os| os.to_str())
                    .unwrap();
                let path_ext = path.extension()
                    .and_then(|os| os.to_str())
                    .unwrap();

                let res = rand::thread_rng().gen_range(500..1000);

                task::sleep(Duration::from_millis(res)).await;

                let simple = SimpleDepend {
                    file_name: file_name.to_string(),
                    ext: path_ext.to_string(),
                };

                Ok(Box::new(simple) as Box<dyn ResourceData>)
            })
        }

        fn load_bytes(&self, _: crate::ResourceManager, _: Vec<u8>, _: usize, _: usize) -> crate::loader::PinedBoxLoaderFuture {
            unreachable!()
        }

        fn create_erased_handle(&self) -> std::sync::Arc<dyn crate::ResourceStorage> {
            Arc::from(ResHandle::<SimpleDepend>::new_loading(None))
        }
    }


    struct SimpleResource {
        depend_a: ResHandle<SimpleDepend>,
    }

    impl ResourceData for SimpleResource {
        fn as_any(&self) -> &dyn std::any::Any {
            self
        }

        fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
            self
        }

        fn dependencies(&self) -> Vec<crate::UntypedResHandle> {
            vec![self.depend_a.untyped_clone()]
        }
    }

    #[derive(Default)]
    struct SlowSimpleResourceLoader;

    impl ResourceLoader for SlowSimpleResourceLoader {
        fn extensions(&self) -> &[&str] {
            &["res", "large"]
        }

        fn mime_types(&self) -> &[&str] {
            &[]
        }

        fn load(&self, res_man: crate::ResourceManager, _: &str) -> crate::loader::PinedBoxLoaderFuture {
            Box::pin(async move {
                let res = rand::thread_rng().gen_range(500..1000);

                task::sleep(Duration::from_millis(res)).await;

                // load dummy dependency that will take a bit
                let depend_path = "depend.txt";
                let depend_han = res_man.request::<SimpleDepend>(depend_path).unwrap();

                let simple = SimpleResource {
                    depend_a: depend_han,
                };

                Ok(Box::new(simple) as Box<dyn ResourceData>)
            })
        }

        fn load_bytes(&self, _: crate::ResourceManager, _: Vec<u8>, _: usize, _: usize) -> crate::loader::PinedBoxLoaderFuture {
            unreachable!()
        }

        fn create_erased_handle(&self) -> std::sync::Arc<dyn crate::ResourceStorage> {
            Arc::from(ResHandle::<SimpleResource>::new_loading(None))
        }
    }

    #[test]
    fn recursive() {
        let man = ResourceManager::new();
        man.register_loader::<SlowSimpleDependLoader>();
        man.register_loader::<SlowSimpleResourceLoader>();

        let res = man.request::<SimpleResource>("massive_asset.res").unwrap();

        let state = res.recurse_dependency_state();
        assert!(state.is_loading());

        // this will take a bit
        res.wait_recurse_dependencies_load();

        let state = res.recurse_dependency_state();
        assert!(!state.is_loading());
    }
}