wfassoc/wfassoc-cdylib/src/object_pool.rs

//! When exporting resources for C interface, resource management and ownership are important things.
//! In this dynamic library, we hold all resources' ownership in Rust world,
//! and only expose a token for C code manipulation.
//! 
//! We need to create a container for holding all resources and providing corresponding operations.
//! So we introduce [ObjectPool] in this module for this purpose.
use slotmap::{DefaultKey, Key, KeyData, SlotMap};
use thiserror::Error as TeError;

/// Error occurs when operating with [ObjectPool].
#[derive(Debug, TeError)]
pub enum Error {
    #[error("given token is not presented in object pool")]
    NoSuchToken,
}

/// The token for fetching object in [ObjectPool].
pub type Token = u64;

/// A pool for managing objects with unique tokens.
/// 
/// It is highly suggested to use this pool with [std::sync::RwLock] guard.
pub struct ObjectPool<T> {
    objs: SlotMap<DefaultKey, T>,
}

impl<T> ObjectPool<T> {
    /// Create a new [ObjectPool].
    pub fn new() -> Self {
        Self {
            objs: SlotMap::new(),
        }
    }

    /// Convert [slotmap] crate's [DefaultKey] to our [Token].
    fn key_to_token(key: &DefaultKey) -> Token {
        key.data().as_ffi()
    }

    /// Convert our [Token] to [slotmap] crate's [DefaultKey].
    fn token_to_key(token: Token) -> DefaultKey {
        DefaultKey::from(KeyData::from_ffi(token))
    }

    /// Put given object into the pool and return its token.
    /// 
    /// The ownership of given object is transferred to the pool.
    pub fn allocate(&mut self, value: T) -> Result<Token, Error> {
        let key = self.objs.insert(value);
        Ok(Self::key_to_token(&key))
    }

    /// Free the object in the pool corresponding to the given token.
    pub fn free(&mut self, token: Token) -> Result<(), Error> {
        let _ = self.pop(token)?;
        Ok(())
    }

    /// Remove the object from the pool corresponding to the given token and return it.
    /// 
    /// The ownership of the object is transferred to the caller.
    pub fn pop(&mut self, token: Token) -> Result<T, Error> {
        match self.objs.remove(Self::token_to_key(token)) {
            Some(obj) => Ok(obj),
            None => Err(Error::NoSuchToken),
        }
    }

    /// Clear all objects in the pool.
    pub fn clear(&mut self) -> () {
        self.objs.clear();
    }

    /// Get a reference to the object in the pool corresponding to the given token.
    pub fn get(&self, token: Token) -> Result<&T, Error> {
        self.objs
            .get(Self::token_to_key(token))
            .ok_or(Error::NoSuchToken)
    }

    /// Get a mutable reference to the object in the pool corresponding to the given token.
    pub fn get_mut(&mut self, token: Token) -> Result<&mut T, Error> {
        self.objs
            .get_mut(Self::token_to_key(token))
            .ok_or(Error::NoSuchToken)
    }
}