ax_io/iobuf/
ext.rs

1#[cfg(feature = "alloc")]
2use alloc::{collections::vec_deque::VecDeque, vec::Vec};
3use core::{
4    io::{BorrowedBuf, BorrowedCursor},
5    mem::MaybeUninit,
6};
7
8use crate::{BufReader, BufWriter, DEFAULT_BUF_SIZE, IoBuf, IoBufMut, Read, Result, Write};
9
10/// Extension methods for [`IoBuf`].
11pub trait IoBufExt: Read + IoBuf {
12    /// Reads some bytes from this buffer and writes them into `writer`.
13    #[inline]
14    fn write_to<W: Write + ?Sized>(&mut self, writer: &mut W) -> Result<usize> {
15        IoBufSpec::write_to(self, writer)
16    }
17}
18
19impl<T: Read + IoBuf + ?Sized> IoBufExt for T {}
20
21/// Extension methods for [`IoBufMut`].
22pub trait IoBufMutExt: Write + IoBufMut {
23    /// Reads some bytes from `reader` and writes them into this buffer.
24    #[inline]
25    fn read_from<R: Read + ?Sized>(&mut self, reader: &mut R) -> Result<usize> {
26        IoBufMutSpec::read_from(self, reader)
27    }
28}
29
30impl<T: Write + IoBufMut + ?Sized> IoBufMutExt for T {}
31
32fn stack_buffer_transfer<R, W>(reader: &mut R, writer: &mut W, size_limit: usize) -> Result<usize>
33where
34    R: Read + ?Sized,
35    W: Write + ?Sized,
36{
37    let mut read_buf = [MaybeUninit::uninit(); DEFAULT_BUF_SIZE];
38
39    let limit = read_buf.len().min(size_limit);
40    let mut buf: BorrowedBuf<'_> = (&mut read_buf[..limit]).into();
41
42    reader.read_buf(buf.unfilled())?;
43
44    if buf.len() == 0 {
45        return Ok(0);
46    }
47
48    writer.write(buf.filled())
49}
50
51trait IoBufSpec {
52    fn write_to<W: Write + ?Sized>(&mut self, writer: &mut W) -> Result<usize>;
53}
54
55impl<R: Read + IoBuf + ?Sized> IoBufSpec for R {
56    default fn write_to<W: Write + ?Sized>(&mut self, writer: &mut W) -> Result<usize> {
57        stack_buffer_transfer(self, writer, self.remaining())
58    }
59}
60
61impl IoBufSpec for &[u8] {
62    fn write_to<W: Write + ?Sized>(&mut self, writer: &mut W) -> Result<usize> {
63        writer.write(self)
64    }
65}
66
67#[cfg(feature = "alloc")]
68impl IoBufSpec for VecDeque<u8> {
69    fn write_to<W: Write + ?Sized>(&mut self, writer: &mut W) -> Result<usize> {
70        let (front, _back) = self.as_slices();
71        let written = writer.write(front)?;
72        self.drain(..written);
73        Ok(written)
74    }
75}
76
77impl<I: ?Sized> IoBufSpec for BufReader<I>
78where
79    Self: Read + IoBuf,
80{
81    fn write_to<W: Write + ?Sized>(&mut self, writer: &mut W) -> Result<usize> {
82        // Hack: this relies on `impl Read for BufReader` always calling fill_buf
83        // if the buffer is empty, even for empty slices.
84        // It can't be called directly here since specialization prevents us
85        // from adding I: Read
86        self.read(&mut [])?;
87
88        let buf = self.buffer();
89        if buf.is_empty() {
90            return Ok(0);
91        }
92
93        let written = writer.write(buf)?;
94        self.consume(written);
95        Ok(written)
96    }
97}
98
99trait IoBufMutSpec {
100    fn read_from<R: Read + ?Sized>(&mut self, reader: &mut R) -> Result<usize>;
101}
102
103impl<W: Write + IoBufMut + ?Sized> IoBufMutSpec for W {
104    default fn read_from<R: Read + ?Sized>(&mut self, reader: &mut R) -> Result<usize> {
105        stack_buffer_transfer(reader, self, self.remaining_mut())
106    }
107}
108
109impl IoBufMutSpec for &mut [u8] {
110    fn read_from<R: Read + ?Sized>(&mut self, reader: &mut R) -> Result<usize> {
111        reader.read(self)
112    }
113}
114
115macro_rules! read_from_vec_impl {
116    ($buf:ident, $reader:ident) => {{
117        let mut read_buf: BorrowedBuf<'_> = $buf.spare_capacity_mut().into();
118        let result = $reader.read_buf(read_buf.unfilled());
119        let bytes_read = read_buf.len();
120        unsafe {
121            $buf.set_len($buf.len() + bytes_read);
122        }
123        result.map(|()| bytes_read)
124    }};
125}
126
127#[cfg(feature = "alloc")]
128impl IoBufMutSpec for Vec<u8> {
129    #[inline]
130    fn read_from<R: Read + ?Sized>(&mut self, reader: &mut R) -> Result<usize> {
131        read_from_vec_impl!(self, reader)
132    }
133}
134
135impl IoBufMutSpec for BorrowedCursor<'_> {
136    fn read_from<R: Read + ?Sized>(&mut self, reader: &mut R) -> Result<usize> {
137        reader.read_buf(self.reborrow())?;
138        Ok(self.written())
139    }
140}
141
142impl<I: Write + ?Sized> IoBufMutSpec for BufWriter<I>
143where
144    Self: IoBufMut,
145{
146    fn read_from<R: Read + ?Sized>(&mut self, reader: &mut R) -> Result<usize> {
147        let buf = self.buffer_mut();
148        read_from_vec_impl!(buf, reader)
149    }
150}
ax_io/iobuf/ext.rs

ax_io/iobuf/
ext.rs
