AppFlowy/shared-lib/lib-infra/src/ref_map.rs

use async_trait::async_trait;
use std::collections::HashMap;
use std::sync::Arc;

#[async_trait]
pub trait RefCountValue {
    async fn did_remove(&self) {}
}

struct RefCountHandler<T> {
    ref_count: usize,
    inner: T,
}

impl<T> RefCountHandler<T> {
    pub fn new(inner: T) -> Self {
        Self { ref_count: 1, inner }
    }

    pub fn increase_ref_count(&mut self) {
        self.ref_count += 1;
    }
}

pub struct RefCountHashMap<T>(HashMap<String, RefCountHandler<T>>);

impl<T> std::default::Default for RefCountHashMap<T> {
    fn default() -> Self {
        Self(HashMap::new())
    }
}

impl<T> RefCountHashMap<T>
where
    T: Clone + Send + Sync + RefCountValue + 'static,
{
    pub fn new() -> Self {
        Self::default()
    }

    pub fn get(&self, key: &str) -> Option<T> {
        self.0.get(key).map(|handler| handler.inner.clone())
    }

    pub fn values(&self) -> Vec<T> {
        self.0.values().map(|value| value.inner.clone()).collect::<Vec<T>>()
    }

    pub fn insert(&mut self, key: String, value: T) {
        if let Some(handler) = self.0.get_mut(&key) {
            handler.increase_ref_count();
        } else {
            let handler = RefCountHandler::new(value);
            self.0.insert(key, handler);
        }
    }

    pub async fn remove(&mut self, key: &str) {
        let mut should_remove = false;
        if let Some(value) = self.0.get_mut(key) {
            if value.ref_count > 0 {
                value.ref_count -= 1;
            }
            should_remove = value.ref_count == 0;
        }

        if should_remove {
            if let Some(handler) = self.0.remove(key) {
                tokio::spawn(async move {
                    handler.inner.did_remove().await;
                });
            }
        }
    }
}

#[async_trait]
impl<T> RefCountValue for Arc<T>
where
    T: RefCountValue + Sync + Send,
{
    async fn did_remove(&self) {
        (**self).did_remove().await
    }
}