Repository: tokio-rs/tokio
File: tokio/src/util/sharded_list.rs (Rust)

Code before:
    lists: Box<[Mutex<LinkedList<L, T>>]>,
    count: AtomicUsize,
    shard_mask: usize,
}

/// Defines the id of a node in a linked list, different ids cause inner list nodes
/// to be scattered in different mutexs.
///
/// # Safety
///
/// Implementations must guarantee that `Target` types are pinned in memory. In
/// other words, when a node is inserted, the value will not be moved as long as
/// it is stored in the list.
pub(crate) unsafe trait ShardedListItem: Link {
    // The returned id is used to pick which list this item should go into.
    unsafe fn get_shared_id(target: NonNull<Self::Target>) -> usize;
}

impl<L, T> ShardedList<L, T> {
    /// Creates a new and empty sharded linked list under specified size
    pub(crate) fn new(sharded_size: usize) -> Self {
        // Find power-of-two sizes best matching arguments
        let mut size = 1;
        while size < sharded_size {
            size <<= 1;
        }
        let shard_mask = size - 1;

        let mut lists = Vec::with_capacity(size as usize);
        for _ in 0..size {
            lists.push(Mutex::new(LinkedList::<L, T>::new()))
        }
        Self {
            lists: lists.into_boxed_slice(),
            count: AtomicUsize::new(0),
            shard_mask,
        }
    }
}

/// Used to get the lock of shard
pub(crate) struct ShardGuard<'a, L, T> {
    lock: MutexGuard<'a, LinkedList<L, T>>,
    count: &'a AtomicUsize,
    id: usize,
}

impl<L: ShardedListItem> ShardedList<L, L::Target> {
    #[allow(dead_code)]
    /// Push a value to this shard.
    pub(crate) fn push(&self, val: L::Handle) {

Review comment (suggestion):
I'm not a big fan of this. It means that the size is not always accurate, and the loop will run infinitely if `sharded_size > 2^63`.
```suggestion
        assert!(sharded_size.is_power_of_two());
```

Code after:
    count: AtomicUsize,
    shard_mask: usize,
}

/// Determines which linked list an item should be stored in.
///
/// # Safety
///
/// Implementations must guarantee that the id of an item does not change from
/// call to call.
pub(crate) unsafe trait ShardedListItem: Link {
    /// # Safety
    /// The provided pointer must point at a valid list item.
    unsafe fn get_shard_id(target: NonNull<Self::Target>) -> usize;
}

impl<L, T> ShardedList<L, T> {
    /// Creates a new and empty sharded linked list with the specified size.
    pub(crate) fn new(sharded_size: usize) -> Self {
        assert!(sharded_size.is_power_of_two());

        let shard_mask = sharded_size - 1;
        let mut lists = Vec::with_capacity(sharded_size);
        for _ in 0..sharded_size {
            lists.push(Mutex::new(LinkedList::<L, T>::new()))
        }
        Self {
            lists: lists.into_boxed_slice(),
            count: AtomicUsize::new(0),
            shard_mask,
        }
    }
}

/// Used to get the lock of shard.
pub(crate) struct ShardGuard<'a, L, T> {
    lock: MutexGuard<'a, LinkedList<L, T>>,
    count: &'a AtomicUsize,
    id: usize,
}

impl<L: ShardedListItem> ShardedList<L, L::Target> {
    /// Removes the last element from a list specified by shard_id and returns it, or None if it is
    /// empty.
    pub(crate) fn pop_back(&self, shard_id: usize) -> Option<L::Handle> {
        let mut lock = self.shard_inner(shard_id);
        let node = lock.pop_back();
        if node.is_some() {
            self.count.fetch_sub(1, Ordering::Relaxed);
        }
        node