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<'_, u8> = (&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        let written = writer.write(self)?;
64        *self = &self[written..];
65        Ok(written)
66    }
67}
68
69#[cfg(feature = "alloc")]
70impl IoBufSpec for VecDeque<u8> {
71    fn write_to<W: Write + ?Sized>(&mut self, writer: &mut W) -> Result<usize> {
72        let (front, _back) = self.as_slices();
73        let written = writer.write(front)?;
74        self.drain(..written);
75        Ok(written)
76    }
77}
78
79impl<I: ?Sized> IoBufSpec for BufReader<I>
80where
81    Self: Read + IoBuf,
82{
83    fn write_to<W: Write + ?Sized>(&mut self, writer: &mut W) -> Result<usize> {
84        // Hack: this relies on `impl Read for BufReader` always calling fill_buf
85        // if the buffer is empty, even for empty slices.
86        // It can't be called directly here since specialization prevents us
87        // from adding I: Read
88        self.read(&mut [])?;
89
90        let buf = self.buffer();
91        if buf.is_empty() {
92            return Ok(0);
93        }
94
95        let written = writer.write(buf)?;
96        self.consume(written);
97        Ok(written)
98    }
99}
100
101trait IoBufMutSpec {
102    fn read_from<R: Read + ?Sized>(&mut self, reader: &mut R) -> Result<usize>;
103}
104
105impl<W: Write + IoBufMut + ?Sized> IoBufMutSpec for W {
106    default fn read_from<R: Read + ?Sized>(&mut self, reader: &mut R) -> Result<usize> {
107        stack_buffer_transfer(reader, self, self.remaining_mut())
108    }
109}
110
111impl IoBufMutSpec for &mut [u8] {
112    fn read_from<R: Read + ?Sized>(&mut self, reader: &mut R) -> Result<usize> {
113        let read = reader.read(self)?;
114        *self = core::mem::take(self).split_at_mut(read).1;
115        Ok(read)
116    }
117}
118
119macro_rules! read_from_vec_impl {
120    ($buf:ident, $reader:ident) => {{
121        let mut read_buf: BorrowedBuf<'_, u8> = $buf.spare_capacity_mut().into();
122        let result = $reader.read_buf(read_buf.unfilled());
123        let bytes_read = read_buf.len();
124        unsafe {
125            $buf.set_len($buf.len() + bytes_read);
126        }
127        result.map(|()| bytes_read)
128    }};
129}
130
131#[cfg(feature = "alloc")]
132impl IoBufMutSpec for Vec<u8> {
133    #[inline]
134    fn read_from<R: Read + ?Sized>(&mut self, reader: &mut R) -> Result<usize> {
135        read_from_vec_impl!(self, reader)
136    }
137}
138
139impl IoBufMutSpec for BorrowedCursor<'_, u8> {
140    fn read_from<R: Read + ?Sized>(&mut self, reader: &mut R) -> Result<usize> {
141        let before = self.written();
142        reader.read_buf(self.reborrow())?;
143        Ok(self.written() - before)
144    }
145}
146
147impl<I: Write + ?Sized> IoBufMutSpec for BufWriter<I>
148where
149    Self: IoBufMut,
150{
151    fn read_from<R: Read + ?Sized>(&mut self, reader: &mut R) -> Result<usize> {
152        let buf = self.buffer_mut();
153        read_from_vec_impl!(buf, reader)
154    }
155}
ax_io/iobuf/ext.rs

ax_io/iobuf/
ext.rs
