freya_components/

cache.rs

use std::{
    any::Any,
    cell::RefCell,
    collections::HashMap,
    hash::{
        DefaultHasher,
        Hash,
        Hasher,
    },
    rc::Rc,
    time::Duration,
};

use freya_core::{
    integration::FxHashSet,
    prelude::*,
};

/// Defines the duration for which an Asset will remain cached after it's user has stopped using it.
/// The default is 1h (3600s).
#[derive(Hash, PartialEq, Eq, Clone, Copy)]
pub enum AssetAge {
    /// Asset will be cached for the specified duration
    Duration(Duration),
    /// Asset will be cached until app is closed
    Unspecified,
}

impl Default for AssetAge {
    fn default() -> Self {
        Self::Duration(Duration::from_secs(3600)) // 1h
    }
}

impl AssetAge {
    /// Asset will be cleaned as soon as it stops being used.
    pub fn zero() -> Self {
        Self::Duration(Duration::ZERO)
    }
}

impl From<Duration> for AssetAge {
    fn from(value: Duration) -> Self {
        Self::Duration(value)
    }
}

/// Configuration for a given Asset.
#[derive(Hash, PartialEq, Eq, Clone)]
pub struct AssetConfiguration {
    /// Asset age.
    pub age: AssetAge,
    /// The ID of the asset.
    pub id: u64,
}

impl AssetConfiguration {
    pub fn new(id: impl Hash, age: AssetAge) -> Self {
        let mut state = DefaultHasher::default();
        id.hash(&mut state);
        let id = state.finish();
        Self { id, age }
    }
}

enum AssetUsers {
    Listeners(Rc<RefCell<FxHashSet<ReactiveContext>>>),
    ClearTask(TaskHandle),
}

#[derive(Clone)]
pub enum Asset {
    /// Asset is cached.
    Cached(Rc<dyn Any>),
    /// Asset is currently being fetched.
    Loading,
    /// Asset has yet to be fetched.
    Pending,
    /// Failed to fetch asset.
    Error(String),
}

impl Asset {
    /// Try to get asset.
    pub fn try_get(&self) -> Option<&Rc<dyn Any>> {
        match self {
            Self::Cached(asset) => Some(asset),
            _ => None,
        }
    }
}

struct AssetState {
    users: AssetUsers,
    asset: Asset,
}

#[derive(Clone, Copy, PartialEq)]
pub struct AssetCacher {
    registry: State<HashMap<AssetConfiguration, AssetState>>,
}

impl AssetCacher {
    pub fn create() -> Self {
        Self {
            registry: State::create(HashMap::new()),
        }
    }

    pub fn try_get() -> Option<Self> {
        try_consume_root_context()
    }

    #[track_caller]
    pub fn get() -> Self {
        consume_root_context()
    }

    /// Attempt to resolve a [Asset] given a [AssetConfiguration].
    pub fn read_asset(&self, asset_config: &AssetConfiguration) -> Option<Asset> {
        self.registry
            .peek()
            .get(asset_config)
            .map(|a| a.asset.clone())
    }

    /// Subscribes to a [Asset] given a [AssetConfiguration].
    pub fn subscribe_asset(&self, asset_config: &AssetConfiguration) -> Option<Asset> {
        self.listen(ReactiveContext::current(), asset_config.clone());
        self.registry
            .peek()
            .get(asset_config)
            .map(|a| a.asset.clone())
    }

    /// Update an [Asset] given a [AssetConfiguration].
    ///
    /// Does nothing if the entry is not present.
    pub fn update_asset(&mut self, asset_config: AssetConfiguration, new_asset: Asset) {
        let mut registry = self.registry.write();

        let Some(asset) = registry.get_mut(&asset_config) else {
            return;
        };

        asset.asset = new_asset;

        // Reruns those listening components
        if let AssetUsers::Listeners(listeners) = &asset.users {
            for sub in listeners.borrow().iter() {
                sub.notify();
            }
        }
    }

    /// Try to clean an asset with no more listeners given a [AssetConfiguration].
    pub fn try_clean(&mut self, asset_config: &AssetConfiguration) {
        let mut registry = self.registry;

        let spawn_clear_task = {
            let mut registry = registry.write();

            let entry = registry.get_mut(asset_config);
            if let Some(asset_state) = entry {
                match &mut asset_state.users {
                    AssetUsers::Listeners(listeners) => {
                        // Only spawn a clear-task if there are no more listeners using this asset
                        listeners.borrow().is_empty()
                    }
                    AssetUsers::ClearTask(task) => {
                        // This case should never happen but... we leave it here anyway.
                        task.cancel();
                        true
                    }
                }
            } else {
                false
            }
        };

        if spawn_clear_task {
            // Only clear the asset if a duration was specified
            if let AssetAge::Duration(duration) = asset_config.age {
                let clear_task = spawn_in_window({
                    let asset_config = asset_config.clone();
                    async move {
                        timer(duration).await;
                        registry.write().remove(&asset_config);
                    }
                });

                // Registry the clear-task
                let mut registry = registry.write();
                if let Some(entry) = registry.get_mut(asset_config) {
                    entry.users = AssetUsers::ClearTask(clear_task);
                } else {
                    #[cfg(debug_assertions)]
                    tracing::info!(
                        "Failed to spawn clear task to remove cache of {}",
                        asset_config.id
                    )
                }
            }
        }
    }

    pub(crate) fn listen(&self, mut rc: ReactiveContext, asset_config: AssetConfiguration) {
        let mut registry = self.registry.write_unchecked();

        let asset = registry.entry(asset_config).or_insert_with(|| AssetState {
            asset: Asset::Pending,
            users: AssetUsers::Listeners(Rc::default()),
        });

        match &mut asset.users {
            AssetUsers::Listeners(users) => {
                rc.subscribe(users);
            }
            AssetUsers::ClearTask(clear_task) => {
                clear_task.cancel();
                let listeners = Rc::default();
                rc.subscribe(&listeners);
                asset.users = AssetUsers::Listeners(listeners);
            }
        }
    }

    /// Read the number of successfully cached assets.
    pub fn cached_size(&self) -> usize {
        self.registry
            .read()
            .values()
            .filter(|asset_state| matches!(&asset_state.asset, Asset::Cached(_)))
            .count()
    }

    /// Read the size of the cache registry.
    pub fn size(&self) -> usize {
        self.registry.read().len()
    }
}

/// Start listening to an asset given a [AssetConfiguration].
pub fn use_asset(asset_config: &AssetConfiguration) {
    let mut asset_cacher = use_hook(AssetCacher::get);

    let mut current_config = use_state::<Option<AssetConfiguration>>(|| None);

    use_drop(move || {
        let Some(asset_config) = current_config.peek().clone() else {
            return;
        };
        // Try to clean in the next async tick, when this scope will already be dropped
        spawn_in_window(async move {
            asset_cacher.try_clean(&asset_config);
        });
    });

    {
        let mut current_config = current_config.write();
        if current_config.as_ref() != Some(asset_config) {
            if let Some(prev) = &*current_config
                && prev != asset_config
            {
                // Try to clean the previous asset
                asset_cacher.try_clean(prev);
            }
            current_config.replace(asset_config.clone());
        }
        asset_cacher.listen(ReactiveContext::current(), asset_config.clone());
    }
}