Repository: tokio-rs/tokio
File: tokio-util/tests/io_inspect.rs (Rust)

Code before:
// If you want something similar for real code, you'll want to avoid all the
// fun of manipulating `bufs` - ideally, by the time you read this,
// IoSlice::advance_slices will be stable, and you can use that.
async fn write_all_vectored<W: AsyncWrite + Unpin>(
    mut writer: W,
    mut bufs: Vec<Vec<u8>>,
) -> Result<usize, std::io::Error> {
    let mut res = 0;
    while !bufs.is_empty() {
        let mut written = poll_fn(|cx| {
            let bufs: Vec<IoSlice> = bufs.iter().map(|v| IoSlice::new(v)).collect();
            Pin::new(&mut writer).poll_write_vectored(cx, &bufs)
        })
        .await?;
        res += written;
        while written > 0 {
            let buf_len = bufs[0].len();
            if buf_len <= written {
                bufs.remove(0);
                written -= buf_len;
            } else {
                let buf = &mut bufs[0];
                while written > 0 {
                    buf.remove(0);
                    written -= 1;
                }
            }
        }
    }
    Ok(res)
}

#[tokio::test]
async fn write_tee_vectored() {
    let mut altout: Vec<u8> = Vec::new();
    let mut writeout = SmallWriter {
        contents: Vec::new(),
    };
    let original = b"A very long string split up";
    let bufs: Vec<Vec<u8>> = original
        .split(|b| b.is_ascii_whitespace())
        .map(Vec::from)
        .collect();
    assert!(bufs.len() > 1);
    let expected: Vec<u8> = {
        let mut out = Vec::new();
        for item in &bufs {
            out.extend_from_slice(item)
        }
        out
    };

Review comment (suggestion):
```suggestion
                let drain_len = written.min(buf.len());
                buf.drain(..drain_len);
                written -= drain_len;
```

Code after:
// If you want something similar for real code, you'll want to avoid all the
// fun of manipulating `bufs` - ideally, by the time you read this,
// IoSlice::advance_slices will be stable, and you can use that.
async fn write_all_vectored<W: AsyncWrite + Unpin>(
    mut writer: W,
    mut bufs: Vec<Vec<u8>>,
) -> Result<usize, std::io::Error> {
    let mut res = 0;
    while !bufs.is_empty() {
        let mut written = poll_fn(|cx| {
            let bufs: Vec<IoSlice> = bufs.iter().map(|v| IoSlice::new(v)).collect();
            Pin::new(&mut writer).poll_write_vectored(cx, &bufs)
        })
        .await?;
        res += written;
        while written > 0 {
            let buf_len = bufs[0].len();
            if buf_len <= written {
                bufs.remove(0);
                written -= buf_len;
            } else {
                let buf = &mut bufs[0];
                let drain_len = written.min(buf.len());
                buf.drain(..drain_len);
                written -= drain_len;
            }
        }
    }
    Ok(res)
}

#[tokio::test]
async fn write_tee_vectored() {
    let mut altout: Vec<u8> = Vec::new();
    let mut writeout = SmallWriter {
        contents: Vec::new(),
    };
    let original = b"A very long string split up";
    let bufs: Vec<Vec<u8>> = original
        .split(|b| b.is_ascii_whitespace())
        .map(Vec::from)
        .collect();
    assert!(bufs.len() > 1);
    let expected: Vec<u8> = {
        let mut out = Vec::new();
        for item in &bufs {
            out.extend_from_slice(item)
        }
        out
    };
    {