ronantakizawa/codereview-bench huggingface dataset:
Repository: tokio-rs/tokio
File: tokio-test/src/stream_mock.rs (Rust)

Code before:
    ) -> std::task::Poll<Option<Self::Item>> {
        // Try polling the sleep future first
        if let Some(ref mut sleep) = self.sleep {
            ready!(Pin::new(sleep).poll(cx));
            // Since we're ready, discard the sleep future
            self.sleep.take();
        }

        match self.next_action() {
            Some(action) => match action {
                Action::Next(item) => Poll::Ready(Some(item)),
                Action::Wait(duration) => {
                    // Set up a sleep future and schedule this future to be polled again for it.
                    self.sleep = Some(Box::pin(sleep_until(Instant::now() + duration)));
                    cx.waker().wake_by_ref();

                    Poll::Pending
                }
            },
            None => Poll::Ready(None),
        }
    }
}

impl<T: Unpin> Drop for StreamMock<T> {
    fn drop(&mut self) {
        let undropped_count = self
            .actions
            .iter()
            .filter(|action| match action {
                Action::Next(_) => true,
                Action::Wait(_) => false,
            })
            .count();

        assert!(
            undropped_count == 0,
            "StreamMock was dropped before all actions were consumed, {} actions were not consumed",
            undropped_count
        );
    }
}


Review comment (suggestion):
You should first verify that we aren't panicking already to avoid a double panic, like we do here: https://github.com/tokio-rs/tokio/blob/master/tokio-test/src/io.rs#L473

Code after:
    ) -> std::task::Poll<Option<Self::Item>> {
        // Try polling the sleep future first
        if let Some(ref mut sleep) = self.sleep {
            ready!(Pin::new(sleep).poll(cx));
            // Since we're ready, discard the sleep future
            self.sleep.take();
        }

        match self.next_action() {
            Some(action) => match action {
                Action::Next(item) => Poll::Ready(Some(item)),
                Action::Wait(duration) => {
                    // Set up a sleep future and schedule this future to be polled again for it.
                    self.sleep = Some(Box::pin(sleep_until(Instant::now() + duration)));
                    cx.waker().wake_by_ref();

                    Poll::Pending
                }
            },
            None => Poll::Ready(None),
        }
    }
}

impl<T: Unpin> Drop for StreamMock<T> {
    fn drop(&mut self) {
        // Avoid double panicking to make debugging easier.
        if std::thread::panicking() {
            return;
        }

        let undropped_count = self
            .actions
            .iter()
            .filter(|action| match action {
                Action::Next(_) => true,
                Action::Wait(_) => false,
            })
            .count();

        assert!(
            undropped_count == 0,
            "StreamMock was dropped before all actions were consumed, {} actions were not consumed",
            undropped_count
        );
    }
}
Repository: tokio-rs/tokio
File: tokio-util/src/io/copy_to_bytes.rs (Rust)

Code before:

    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.project().inner.poll_ready(cx)
    }

    fn start_send(self: Pin<&mut Self>, item: &'a [u8]) -> Result<(), Self::Error> {
        self.project()
            .inner
            .start_send(Bytes::copy_from_slice(item))
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.project().inner.poll_flush(cx)
    }

    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.project().inner.poll_close(cx)
    }
}

impl<S, T> Stream for CopyToBytes<S>
where
    S: Stream<Item = T> + std::marker::Unpin,
{
    type Item = T;
    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let ref_mut = self.get_mut().get_mut();
        Pin::new(ref_mut).poll_next(cx)
    }
}


Review comment (suggestion):
We don't want `Unpin` bounds if they're not necessary. Also, the `T` parameter can be dropped.
```suggestion
impl<S: Stream> Stream for CopyToBytes<S> {
    type Item = S::Item;
    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        self.project().inner.poll_next(cx)
    }
}
```

Code after:

    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.project().inner.poll_ready(cx)
    }

    fn start_send(self: Pin<&mut Self>, item: &'a [u8]) -> Result<(), Self::Error> {
        self.project()
            .inner
            .start_send(Bytes::copy_from_slice(item))
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.project().inner.poll_flush(cx)
    }

    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.project().inner.poll_close(cx)
    }
}

impl<S: Stream> Stream for CopyToBytes<S> {
    type Item = S::Item;
    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        self.project().inner.poll_next(cx)
    }
}
Repository: tokio-rs/tokio
File: tokio/tests/fs_file.rs (Rust)

Code before:
    let tempfile = tempfile();

    let mut file = File::create(tempfile.path()).await.unwrap();

    let ret = file
        .write_vectored(&[IoSlice::new(HELLO), IoSlice::new(HELLO)])
        .await
        .unwrap();
    assert_eq!(ret, HELLO.bytes().count() * 2);
    file.flush().await.unwrap();

    let file = std::fs::read(tempfile.path()).unwrap();
    assert_eq!(file, [HELLO, HELLO].concat());
}

#[tokio::test]
async fn write_vectored_and_shutdown() {
    let tempfile = tempfile();

    let mut file = File::create(tempfile.path()).await.unwrap();

    let ret = file
        .write_vectored(&[IoSlice::new(HELLO), IoSlice::new(HELLO)])
        .await
        .unwrap();
    assert_eq!(ret, HELLO.bytes().count() * 2);
    file.shutdown().await.unwrap();

    let file = std::fs::read(tempfile.path()).unwrap();
    assert_eq!(file, [HELLO, HELLO].concat());
}

#[tokio::test]
async fn rewind_seek_position() {
    let tempfile = tempfile();

    let mut file = File::create(tempfile.path()).await.unwrap();

    file.seek(SeekFrom::Current(10)).await.unwrap();

    file.rewind().await.unwrap();

    assert_eq!(file.stream_position().await.unwrap(), 0);
}

#[tokio::test]
async fn coop() {
    let mut tempfile = tempfile();
    tempfile.write_all(HELLO).unwrap();

    let mut task = task::spawn(async {

Review comment (suggestion):
This is in the test twice.
```suggestion
    assert_eq!(ret, HELLO.len() * 2);
```

Code after:
    let tempfile = tempfile();

    let mut file = File::create(tempfile.path()).await.unwrap();

    let ret = file
        .write_vectored(&[IoSlice::new(HELLO), IoSlice::new(HELLO)])
        .await
        .unwrap();
    assert_eq!(ret, HELLO.len() * 2);
    file.flush().await.unwrap();

    let file = std::fs::read(tempfile.path()).unwrap();
    assert_eq!(file, [HELLO, HELLO].concat());
}

#[tokio::test]
async fn write_vectored_and_shutdown() {
    let tempfile = tempfile();

    let mut file = File::create(tempfile.path()).await.unwrap();

    let ret = file
        .write_vectored(&[IoSlice::new(HELLO), IoSlice::new(HELLO)])
        .await
        .unwrap();
    assert_eq!(ret, HELLO.len() * 2);
    file.shutdown().await.unwrap();

    let file = std::fs::read(tempfile.path()).unwrap();
    assert_eq!(file, [HELLO, HELLO].concat());
}

#[tokio::test]
async fn rewind_seek_position() {
    let tempfile = tempfile();

    let mut file = File::create(tempfile.path()).await.unwrap();

    file.seek(SeekFrom::Current(10)).await.unwrap();

    file.rewind().await.unwrap();

    assert_eq!(file.stream_position().await.unwrap(), 0);
}

#[tokio::test]
async fn coop() {
    let mut tempfile = tempfile();
    tempfile.write_all(HELLO).unwrap();

    let mut task = task::spawn(async {
Repository: tokio-rs/tokio
File: tokio/src/io/blocking.rs (Rust)

Code before:
        res
    }

    pub(crate) fn write_to<T: Write>(&mut self, wr: &mut T) -> io::Result<()> {
        assert_eq!(self.pos, 0);

        // `write_all` already ignores interrupts
        let res = wr.write_all(&self.buf);
        self.buf.clear();
        res
    }
}

cfg_fs! {
    impl Buf {
        pub(crate) fn discard_read(&mut self) -> i64 {
            let ret = -(self.bytes().len() as i64);
            self.pos = 0;
            self.buf.truncate(0);
            ret
        }

        pub(crate) fn copy_from_bufs(&mut self, bufs: &[io::IoSlice<'_>]) -> usize {
            assert!(self.is_empty());

            let n = bufs.iter().map(|b| b.len()).sum::<usize>().min(MAX_BUF);

            let mut rem = n;
            for buf in bufs {
                if rem == 0 {
                    break
                }

                let len = buf.len().min(rem);
                self.buf.extend_from_slice(&buf[..len]);
                rem -= len;
            }

            n
        }
    }
}


Review comment (performance):
This iterates all of the buffers every time, even if we only write a few of them. If the buffers are very long and this is called in a loop, that gives quadratic performance.

We should be able to embed this logic inside the for loop instead to avoid that.

Code after:
        res
    }

    pub(crate) fn write_to<T: Write>(&mut self, wr: &mut T) -> io::Result<()> {
        assert_eq!(self.pos, 0);

        // `write_all` already ignores interrupts
        let res = wr.write_all(&self.buf);
        self.buf.clear();
        res
    }
}

cfg_fs! {
    impl Buf {
        pub(crate) fn discard_read(&mut self) -> i64 {
            let ret = -(self.bytes().len() as i64);
            self.pos = 0;
            self.buf.truncate(0);
            ret
        }

        pub(crate) fn copy_from_bufs(&mut self, bufs: &[io::IoSlice<'_>]) -> usize {
            assert!(self.is_empty());

            let mut rem = MAX_BUF;
            for buf in bufs {
                if rem == 0 {
                    break
                }

                let len = buf.len().min(rem);
                self.buf.extend_from_slice(&buf[..len]);
                rem -= len;
            }

            MAX_BUF - rem
        }
    }
}
Repository: tokio-rs/tokio
File: tokio/src/runtime/task/id.rs (Rust)

Code before:
/// This function is similar to  [`task::id()`](crate::runtime::task::id()), except
/// that it returns `None` rather than panicking if called outside of a task
/// context.
///
/// **Note**: This is an [unstable API][unstable]. The public API of this type
/// may break in 1.x releases. See [the documentation on unstable
/// features][unstable] for details.
///
/// [task ID]: crate::task::Id
/// [unstable]: crate#unstable-features
#[cfg_attr(not(tokio_unstable), allow(unreachable_pub))]
#[track_caller]
pub fn try_id() -> Option<Id> {
    context::current_task_id()
}

impl fmt::Display for Id {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl Id {
    pub(crate) fn next() -> Self {
        #[cfg(loom)]
        return Self(context::next_task_id().unwrap());

        #[cfg(not(loom))]
        {
            use crate::loom::sync::atomic::{Ordering::Relaxed, StaticAtomicU64};

            static NEXT_ID: StaticAtomicU64 = StaticAtomicU64::new(1);

            Self(NEXT_ID.fetch_add(1, Relaxed))
        }
    }

    pub(crate) fn as_u64(&self) -> u64 {
        self.0
    }
}


Review comment (suggestion):
I think what you can do is use loom's lazy_static wrapper, which you can read about at https://docs.rs/loom/latest/loom/macro.lazy_static.html

This avoids a problem where tasks created in different threads get the same task id, since thread locals have a separate counter per thread.

By using the lazy_static wrapper, you should be able to put a `loom::sync::atomic::AtomicU64` in a global.

Code after:
/// This function is similar to  [`task::id()`](crate::runtime::task::id()), except
/// that it returns `None` rather than panicking if called outside of a task
/// context.
///
/// **Note**: This is an [unstable API][unstable]. The public API of this type
/// may break in 1.x releases. See [the documentation on unstable
/// features][unstable] for details.
///
/// [task ID]: crate::task::Id
/// [unstable]: crate#unstable-features
#[cfg_attr(not(tokio_unstable), allow(unreachable_pub))]
#[track_caller]
pub fn try_id() -> Option<Id> {
    context::current_task_id()
}

impl fmt::Display for Id {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.0.fmt(f)
    }
}

impl Id {
    pub(crate) fn next() -> Self {
        use crate::loom::sync::atomic::Ordering::Relaxed;
        use crate::loom::sync::atomic::StaticAtomicU64;

        #[cfg(all(test, loom))]
        {
            crate::loom::lazy_static! {
                static ref NEXT_ID: StaticAtomicU64 = StaticAtomicU64::new(1);
            }
            Self(NEXT_ID.fetch_add(1, Relaxed))
        }

        #[cfg(not(all(test, loom)))]
        {
            static NEXT_ID: StaticAtomicU64 = StaticAtomicU64::new(1);
            Self(NEXT_ID.fetch_add(1, Relaxed))
        }
    }

    pub(crate) fn as_u64(&self) -> u64 {
        self.0
    }
}
Repository: tokio-rs/tokio
File: tokio/src/util/sharded_list.rs (Rust)

Code before:
#![cfg_attr(not(feature = "full"), allow(dead_code))]

use std::{ptr::NonNull, sync::atomic::Ordering};

use crate::loom::sync::{Mutex, MutexGuard};
use std::sync::atomic::AtomicUsize;

use super::linked_list::{Link, LinkedList};

/// An intrusive  linked list supporting high concurrent updates.
///
/// It currently relies on `LinkedList`, so it is the caller's
/// responsibility to ensure the list is empty before dropping it.
///
/// Note: Due to its inner sharded design, the order of node cannot be guaranteed.
pub(crate) struct ShardedList<L, T> {
    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

Review comment (suggestion):
We only use `{}` for things in the same module.
```suggestion
use std::ptr::NonNull;
use std::sync::atomic::Ordering;
```

Code after:
use std::ptr::NonNull;
use std::sync::atomic::Ordering;

use crate::loom::sync::{Mutex, MutexGuard};
use std::sync::atomic::AtomicUsize;

use super::linked_list::{Link, LinkedList};

/// An intrusive linked list supporting highly concurrent updates.
///
/// It currently relies on `LinkedList`, so it is the caller's
/// responsibility to ensure the list is empty before dropping it.
///
/// Note: Due to its inner sharded design, the order of nodes cannot be guaranteed.
pub(crate) struct ShardedList<L, T> {
    lists: Box<[Mutex<LinkedList<L, T>>]>,
    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
Repository: tokio-rs/tokio
File: tokio/src/util/sharded_list.rs (Rust)

Code before:
#![cfg_attr(not(feature = "full"), allow(dead_code))]

use std::{ptr::NonNull, sync::atomic::Ordering};

use crate::loom::sync::{Mutex, MutexGuard};
use std::sync::atomic::AtomicUsize;

use super::linked_list::{Link, LinkedList};

/// An intrusive  linked list supporting high concurrent updates.
///
/// It currently relies on `LinkedList`, so it is the caller's
/// responsibility to ensure the list is empty before dropping it.
///
/// Note: Due to its inner sharded design, the order of node cannot be guaranteed.
pub(crate) struct ShardedList<L, T> {
    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> {

Review comment (suggestion):
```suggestion
/// An intrusive linked list supporting highly concurrent updates.
```

Code after:
use std::ptr::NonNull;
use std::sync::atomic::Ordering;

use crate::loom::sync::{Mutex, MutexGuard};
use std::sync::atomic::AtomicUsize;

use super::linked_list::{Link, LinkedList};

/// An intrusive linked list supporting highly concurrent updates.
///
/// It currently relies on `LinkedList`, so it is the caller's
/// responsibility to ensure the list is empty before dropping it.
///
/// Note: Due to its inner sharded design, the order of nodes cannot be guaranteed.
pub(crate) struct ShardedList<L, T> {
    lists: Box<[Mutex<LinkedList<L, T>>]>,
    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 {
Repository: tokio-rs/tokio
File: tokio/src/util/sharded_list.rs (Rust)

Code before:
#![cfg_attr(not(feature = "full"), allow(dead_code))]

use std::{ptr::NonNull, sync::atomic::Ordering};

use crate::loom::sync::{Mutex, MutexGuard};
use std::sync::atomic::AtomicUsize;

use super::linked_list::{Link, LinkedList};

/// An intrusive  linked list supporting high concurrent updates.
///
/// It currently relies on `LinkedList`, so it is the caller's
/// responsibility to ensure the list is empty before dropping it.
///
/// Note: Due to its inner sharded design, the order of node cannot be guaranteed.
pub(crate) struct ShardedList<L, T> {
    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 {

Review comment (suggestion):
```suggestion
/// Note: Due to its inner sharded design, the order of nodes cannot be guaranteed.
```

Code after:
use std::ptr::NonNull;
use std::sync::atomic::Ordering;

use crate::loom::sync::{Mutex, MutexGuard};
use std::sync::atomic::AtomicUsize;

use super::linked_list::{Link, LinkedList};

/// An intrusive linked list supporting highly concurrent updates.
///
/// It currently relies on `LinkedList`, so it is the caller's
/// responsibility to ensure the list is empty before dropping it.
///
/// Note: Due to its inner sharded design, the order of nodes cannot be guaranteed.
pub(crate) struct ShardedList<L, T> {
    lists: Box<[Mutex<LinkedList<L, T>>]>,
    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 {
Repository: tokio-rs/tokio
File: tokio/src/util/sharded_list.rs (Rust)

Code before:
    pub(crate) fn shard_size(&self) -> usize {
        self.shard_mask + 1
    }

    #[inline]
    fn shard_inner(&self, id: usize) -> MutexGuard<'_, LinkedList<L, <L as Link>::Target>> {
        // Safety: this modulo operation ensures it is safe here.
        unsafe { self.lists.get_unchecked(id & self.shard_mask).lock() }
    }
}

impl<'a, L: ShardedListItem> ShardGuard<'a, L, L::Target> {
    /// Push a value to this shard.
    pub(crate) fn push(mut self, val: L::Handle) {
        let id = unsafe { L::get_shared_id(L::as_raw(&val)) };
        assert_eq!(id, self.id);
        self.lock.push_front(val);
        self.count.fetch_add(1, Ordering::Relaxed);
    }
}

cfg_taskdump! {
    impl<L: ShardedListItem> ShardedList<L, L::Target> {
        pub(crate) fn for_each<F>(&self, f: &mut F)
        where
        F: FnMut(&L::Handle),
        {
            let mut guards = vec![];
            for list in self.lists.iter(){
                guards.push(list.lock());
            }
            for g in &mut guards{
                g.for_each(f);
            }
        }
    }
}


Review comment (suggestion):
```suggestion
            F: FnMut(&L::Handle),
```

Code after:
impl<'a, L: ShardedListItem> ShardGuard<'a, L, L::Target> {
    /// Push a value to this shard.
    pub(crate) fn push(mut self, val: L::Handle) {
        let id = unsafe { L::get_shard_id(L::as_raw(&val)) };
        assert_eq!(id, self.id);
        self.lock.push_front(val);
        self.count.fetch_add(1, Ordering::Relaxed);
    }
}

cfg_taskdump! {
    impl<L: ShardedListItem> ShardedList<L, L::Target> {
        pub(crate) fn for_each<F>(&self, mut f: F)
        where
            F: FnMut(&L::Handle),
        {
            let mut guards = Vec::with_capacity(self.lists.len());
            for list in self.lists.iter() {
                guards.push(list.lock());
            }
            for g in &mut guards {
                g.for_each(&mut f);
            }
        }
    }
}