txmap 3.0.1

A concurrent transactional hash map for Rust with fine-grained locking, internal mutability and composable transactions
Documentation
use crate::{
    lock_guards::LockGuards,
    lock_policies::lock_policy::LockPolicy,
    new_types::{BitMask, ShardCount, ShardIndex},
    shard::Shard,
};
use hashbrown::HashTable;
use intmap::IntMap;

pub(crate) struct Custodian<K, V, L>
where
    L: LockPolicy,
{
    pub(crate) shard_count: ShardCount,
    pub(crate) shards: Vec<L::Lock<Shard<K, V>>>,
}

impl<K, V, L> Custodian<K, V, L>
where
    L: LockPolicy,
{
    pub fn new(shard_count: ShardCount, capacity: usize) -> Self {
        let mut shards = Vec::with_capacity(shard_count.0 as usize);
        let capacity_per_shard = capacity.div_ceil(shard_count.0 as usize);
        for _ in 0..shard_count.0 {
            shards.push(L::new(HashTable::with_capacity(capacity_per_shard)));
        }
        Self {
            shard_count,
            shards,
        }
    }
    pub fn all_read_guards(&self) -> IntMap<u8, L::ReadGuard<'_, Shard<K, V>>> {
        self.lock_guards(self.all_bitmask(), BitMask::ZERO).read
    }
    pub fn all_write_guards(&self) -> IntMap<u8, L::WriteGuard<'_, Shard<K, V>>> {
        self.lock_guards(BitMask::ZERO, self.all_bitmask()).write
    }
    fn all_bitmask(&self) -> BitMask {
        let bitmask = if self.shard_count.0 == 128 {
            !0u128
        } else {
            (1 << self.shard_count.0) - 1
        };
        BitMask(bitmask)
    }
    pub fn lock_guards(&self, read: BitMask, write: BitMask) -> LockGuards<'_, K, V, L> {
        let mut read_guards = IntMap::new();
        let mut write_guards = IntMap::new();
        for i in 0..self.shard_count.0 {
            let bitmask = ShardIndex(i).bitmask();
            let shard_lock = &self.shards[i as usize];
            if (write & bitmask) != BitMask::ZERO {
                let write_guard = L::write(shard_lock);
                write_guards.insert(i, write_guard);
            } else if (read & bitmask) != BitMask::ZERO {
                let read_guard = L::read(shard_lock);
                read_guards.insert(i, read_guard);
            };
        }
        LockGuards {
            read: read_guards,
            write: write_guards,
            write_bitmask: write,
        }
    }
    pub fn write_guards(&self, write: BitMask) -> IntMap<u8, L::WriteGuard<'_, Shard<K, V>>> {
        let mut write_guards = IntMap::new();
        for i in 0..self.shard_count.0 {
            let bitmask = ShardIndex(i).bitmask();
            if (write & bitmask) != BitMask::ZERO {
                let shard_lock = &self.shards[i as usize];
                let write_guard = L::write(shard_lock);
                write_guards.insert(i, write_guard);
            };
        }
        write_guards
    }
    pub fn read_guard_at(&self, shard_index: ShardIndex) -> L::ReadGuard<'_, Shard<K, V>> {
        let shard_lock = &self.shards[shard_index.0 as usize];
        L::read(shard_lock)
    }
    pub fn write_guard_at(&self, shard_index: ShardIndex) -> L::WriteGuard<'_, Shard<K, V>> {
        let shard_lock = &self.shards[shard_index.0 as usize];
        L::write(shard_lock)
    }
}