channels_io/buf/
buf_mut.rs

1use core::cmp::min;
2use core::mem;
3use core::pin::Pin;
4use core::task::{Context, Poll};
5
6use crate::buf::{Chain, Limit};
7use crate::error::{IoError, WriteError};
8use crate::util::copy_slice;
9use crate::{AsyncWrite, Write};
10
11/// A mutable and initialized buffer that contains an internal cursor.
12pub trait BufMut {
13	/// Get the number of bytes this buffer can hold.
14	fn remaining_mut(&self) -> usize;
15
16	/// Get the remaining part of the buffer as a mutable slice.
17	///
18	/// This method is allowed to return slices smaller than what [`BufMut::remaining_mut()`]
19	/// describes. This allows non contiguous representation of the buffer in memory.
20	fn chunk_mut(&mut self) -> &mut [u8];
21
22	/// Advance the internal cursor of the buffer by `n` bytes.
23	///
24	/// # Panics
25	///
26	/// If `n` is a value that would cause the cursor to go out of bounds.
27	fn advance_mut(&mut self, n: usize);
28
29	/// Check whether the buffer has any data left in it.
30	fn has_remaining_mut(&self) -> bool {
31		self.remaining_mut() != 0
32	}
33
34	/// Copy data from the buffer into `slice` advancing the buffer accordingly.
35	fn copy_from_slice(&mut self, slice: &[u8]) -> usize {
36		let n = min(slice.len(), self.remaining_mut());
37		let mut i = 0;
38
39		while i < n {
40			let x = copy_slice(&slice[i..], self.chunk_mut());
41			self.advance_mut(x);
42			i += x;
43		}
44
45		n
46	}
47
48	/// Create a "by reference" adapter that takes the current instance of [`BufMut`]
49	/// by mutable reference.
50	fn by_ref(&mut self) -> &mut Self
51	where
52		Self: Sized,
53	{
54		self
55	}
56
57	/// Create an adapter that implements the [`Write`] and [`AsyncWrite`] traits
58	/// on the current instance of [`BufMut`].
59	fn writer(self) -> Writer<Self>
60	where
61		Self: Sized,
62	{
63		Writer::new(self)
64	}
65
66	/// Create a [`Chain`] adapter that chains this buffer with `other`.
67	///
68	/// The returned [`BufMut`] will behave as a non-contiguous buffer made up of
69	/// `self` and `other`.
70	fn chain<T: BufMut>(self, other: T) -> Chain<Self, T>
71	where
72		Self: Sized,
73	{
74		Chain::new(self, other)
75	}
76
77	/// Create a [`Limit`] adapter that limits accesses to the first `n` bytes of
78	/// this buffer.
79	///
80	/// The returned [`BufMut`] will contain only the first `n` bytes of this buffer.
81	fn limit(self, n: usize) -> Limit<Self>
82	where
83		Self: Sized,
84	{
85		Limit::new(self, n)
86	}
87}
88
89macro_rules! forward_impl_buf_mut {
90	($to:ty) => {
91		fn remaining_mut(&self) -> usize {
92			<$to>::remaining_mut(self)
93		}
94
95		fn chunk_mut(&mut self) -> &mut [u8] {
96			<$to>::chunk_mut(self)
97		}
98
99		fn advance_mut(&mut self, n: usize) {
100			<$to>::advance_mut(self, n)
101		}
102
103		fn has_remaining_mut(&self) -> bool {
104			<$to>::has_remaining_mut(self)
105		}
106
107		fn copy_from_slice(&mut self, slice: &[u8]) -> usize {
108			<$to>::copy_from_slice(self, slice)
109		}
110	};
111}
112
113impl<T: BufMut + ?Sized> BufMut for &mut T {
114	forward_impl_buf_mut!(T);
115}
116
117#[cfg(feature = "alloc")]
118impl<T: BufMut + ?Sized> BufMut for alloc::boxed::Box<T> {
119	forward_impl_buf_mut!(T);
120}
121
122impl BufMut for &mut [u8] {
123	fn remaining_mut(&self) -> usize {
124		self.len()
125	}
126
127	fn chunk_mut(&mut self) -> &mut [u8] {
128		self
129	}
130
131	fn advance_mut(&mut self, n: usize) {
132		let tmp = mem::take(self);
133		*self = &mut tmp[n..];
134	}
135}
136
137/// The error returned by IO write calls to a [`Writer`].
138#[derive(Debug, Clone, Copy)]
139pub enum WriterError {
140	/// The writer has no more space left.
141	WriteZero,
142}
143
144impl IoError for WriterError {
145	fn should_retry(&self) -> bool {
146		false
147	}
148}
149
150impl WriteError for WriterError {
151	fn write_zero() -> Self {
152		Self::WriteZero
153	}
154}
155
156/// An adapter for [`BufMut`] that implements both [`Write`] and [`AsyncWrite`].
157#[derive(Debug, Clone, Copy)]
158pub struct Writer<B> {
159	buf: B,
160}
161
162impl<B> Writer<B> {
163	pub(crate) fn new(buf: B) -> Self {
164		Self { buf }
165	}
166}
167
168impl<B: BufMut> Write for Writer<B> {
169	type Error = WriterError;
170
171	fn write_slice(
172		&mut self,
173		buf: &[u8],
174	) -> Result<usize, Self::Error> {
175		if buf.is_empty() {
176			return Ok(0);
177		}
178
179		let n = self.buf.copy_from_slice(buf);
180		if n == 0 {
181			return Err(WriterError::WriteZero);
182		}
183
184		Ok(n)
185	}
186
187	fn flush_once(&mut self) -> Result<(), Self::Error> {
188		Ok(())
189	}
190}
191
192impl<B: BufMut + Unpin> AsyncWrite for Writer<B> {
193	type Error = WriterError;
194
195	fn poll_write_slice(
196		mut self: Pin<&mut Self>,
197		_: &mut Context,
198		buf: &[u8],
199	) -> Poll<Result<usize, Self::Error>> {
200		Poll::Ready(self.write_slice(buf))
201	}
202
203	fn poll_flush_once(
204		self: Pin<&mut Self>,
205		_: &mut Context,
206	) -> Poll<Result<(), Self::Error>> {
207		Poll::Ready(Ok(()))
208	}
209}