1 use core::fmt;
2 use core::mem::MaybeUninit;
3 use core::ops::{
4     Index, IndexMut, Range, RangeFrom, RangeFull, RangeInclusive, RangeTo, RangeToInclusive,
5 };
6 
7 /// Uninitialized byte slice.
8 ///
9 /// Returned by `BufMut::chunk_mut()`, the referenced byte slice may be
10 /// uninitialized. The wrapper provides safe access without introducing
11 /// undefined behavior.
12 ///
13 /// The safety invariants of this wrapper are:
14 ///
15 ///  1. Reading from an `UninitSlice` is undefined behavior.
16 ///  2. Writing uninitialized bytes to an `UninitSlice` is undefined behavior.
17 ///
18 /// The difference between `&mut UninitSlice` and `&mut [MaybeUninit<u8>]` is
19 /// that it is possible in safe code to write uninitialized bytes to an
20 /// `&mut [MaybeUninit<u8>]`, which this type prohibits.
21 #[repr(transparent)]
22 pub struct UninitSlice([MaybeUninit<u8>]);
23 
24 impl UninitSlice {
25     /// Creates a `&mut UninitSlice` wrapping a slice of initialised memory.
26     ///
27     /// # Examples
28     ///
29     /// ```
30     /// use bytes::buf::UninitSlice;
31     ///
32     /// let mut buffer = [0u8; 64];
33     /// let slice = UninitSlice::new(&mut buffer[..]);
34     /// ```
35     #[inline]
new(slice: &mut [u8]) -> &mut UninitSlice36     pub fn new(slice: &mut [u8]) -> &mut UninitSlice {
37         unsafe { &mut *(slice as *mut [u8] as *mut [MaybeUninit<u8>] as *mut UninitSlice) }
38     }
39 
40     /// Creates a `&mut UninitSlice` wrapping a slice of uninitialised memory.
41     ///
42     /// # Examples
43     ///
44     /// ```
45     /// use bytes::buf::UninitSlice;
46     /// use core::mem::MaybeUninit;
47     ///
48     /// let mut buffer = [MaybeUninit::uninit(); 64];
49     /// let slice = UninitSlice::uninit(&mut buffer[..]);
50     ///
51     /// let mut vec = Vec::with_capacity(1024);
52     /// let spare: &mut UninitSlice = vec.spare_capacity_mut().into();
53     /// ```
54     #[inline]
uninit(slice: &mut [MaybeUninit<u8>]) -> &mut UninitSlice55     pub fn uninit(slice: &mut [MaybeUninit<u8>]) -> &mut UninitSlice {
56         unsafe { &mut *(slice as *mut [MaybeUninit<u8>] as *mut UninitSlice) }
57     }
58 
uninit_ref(slice: &[MaybeUninit<u8>]) -> &UninitSlice59     fn uninit_ref(slice: &[MaybeUninit<u8>]) -> &UninitSlice {
60         unsafe { &*(slice as *const [MaybeUninit<u8>] as *const UninitSlice) }
61     }
62 
63     /// Create a `&mut UninitSlice` from a pointer and a length.
64     ///
65     /// # Safety
66     ///
67     /// The caller must ensure that `ptr` references a valid memory region owned
68     /// by the caller representing a byte slice for the duration of `'a`.
69     ///
70     /// # Examples
71     ///
72     /// ```
73     /// use bytes::buf::UninitSlice;
74     ///
75     /// let bytes = b"hello world".to_vec();
76     /// let ptr = bytes.as_ptr() as *mut _;
77     /// let len = bytes.len();
78     ///
79     /// let slice = unsafe { UninitSlice::from_raw_parts_mut(ptr, len) };
80     /// ```
81     #[inline]
from_raw_parts_mut<'a>(ptr: *mut u8, len: usize) -> &'a mut UninitSlice82     pub unsafe fn from_raw_parts_mut<'a>(ptr: *mut u8, len: usize) -> &'a mut UninitSlice {
83         let maybe_init: &mut [MaybeUninit<u8>] =
84             core::slice::from_raw_parts_mut(ptr as *mut _, len);
85         Self::uninit(maybe_init)
86     }
87 
88     /// Write a single byte at the specified offset.
89     ///
90     /// # Panics
91     ///
92     /// The function panics if `index` is out of bounds.
93     ///
94     /// # Examples
95     ///
96     /// ```
97     /// use bytes::buf::UninitSlice;
98     ///
99     /// let mut data = [b'f', b'o', b'o'];
100     /// let slice = unsafe { UninitSlice::from_raw_parts_mut(data.as_mut_ptr(), 3) };
101     ///
102     /// slice.write_byte(0, b'b');
103     ///
104     /// assert_eq!(b"boo", &data[..]);
105     /// ```
106     #[inline]
write_byte(&mut self, index: usize, byte: u8)107     pub fn write_byte(&mut self, index: usize, byte: u8) {
108         assert!(index < self.len());
109 
110         unsafe { self[index..].as_mut_ptr().write(byte) }
111     }
112 
113     /// Copies bytes from `src` into `self`.
114     ///
115     /// The length of `src` must be the same as `self`.
116     ///
117     /// # Panics
118     ///
119     /// The function panics if `src` has a different length than `self`.
120     ///
121     /// # Examples
122     ///
123     /// ```
124     /// use bytes::buf::UninitSlice;
125     ///
126     /// let mut data = [b'f', b'o', b'o'];
127     /// let slice = unsafe { UninitSlice::from_raw_parts_mut(data.as_mut_ptr(), 3) };
128     ///
129     /// slice.copy_from_slice(b"bar");
130     ///
131     /// assert_eq!(b"bar", &data[..]);
132     /// ```
133     #[inline]
copy_from_slice(&mut self, src: &[u8])134     pub fn copy_from_slice(&mut self, src: &[u8]) {
135         use core::ptr;
136 
137         assert_eq!(self.len(), src.len());
138 
139         unsafe {
140             ptr::copy_nonoverlapping(src.as_ptr(), self.as_mut_ptr(), self.len());
141         }
142     }
143 
144     /// Return a raw pointer to the slice's buffer.
145     ///
146     /// # Safety
147     ///
148     /// The caller **must not** read from the referenced memory and **must not**
149     /// write **uninitialized** bytes to the slice either.
150     ///
151     /// # Examples
152     ///
153     /// ```
154     /// use bytes::BufMut;
155     ///
156     /// let mut data = [0, 1, 2];
157     /// let mut slice = &mut data[..];
158     /// let ptr = BufMut::chunk_mut(&mut slice).as_mut_ptr();
159     /// ```
160     #[inline]
as_mut_ptr(&mut self) -> *mut u8161     pub fn as_mut_ptr(&mut self) -> *mut u8 {
162         self.0.as_mut_ptr() as *mut _
163     }
164 
165     /// Return a `&mut [MaybeUninit<u8>]` to this slice's buffer.
166     ///
167     /// # Safety
168     ///
169     /// The caller **must not** read from the referenced memory and **must not** write
170     /// **uninitialized** bytes to the slice either. This is because `BufMut` implementation
171     /// that created the `UninitSlice` knows which parts are initialized. Writing uninitialized
172     /// bytes to the slice may cause the `BufMut` to read those bytes and trigger undefined
173     /// behavior.
174     ///
175     /// # Examples
176     ///
177     /// ```
178     /// use bytes::BufMut;
179     ///
180     /// let mut data = [0, 1, 2];
181     /// let mut slice = &mut data[..];
182     /// unsafe {
183     ///     let uninit_slice = BufMut::chunk_mut(&mut slice).as_uninit_slice_mut();
184     /// };
185     /// ```
186     #[inline]
as_uninit_slice_mut(&mut self) -> &mut [MaybeUninit<u8>]187     pub unsafe fn as_uninit_slice_mut(&mut self) -> &mut [MaybeUninit<u8>] {
188         &mut self.0
189     }
190 
191     /// Returns the number of bytes in the slice.
192     ///
193     /// # Examples
194     ///
195     /// ```
196     /// use bytes::BufMut;
197     ///
198     /// let mut data = [0, 1, 2];
199     /// let mut slice = &mut data[..];
200     /// let len = BufMut::chunk_mut(&mut slice).len();
201     ///
202     /// assert_eq!(len, 3);
203     /// ```
204     #[inline]
len(&self) -> usize205     pub fn len(&self) -> usize {
206         self.0.len()
207     }
208 }
209 
210 impl fmt::Debug for UninitSlice {
fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result211     fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
212         fmt.debug_struct("UninitSlice[...]").finish()
213     }
214 }
215 
216 impl<'a> From<&'a mut [u8]> for &'a mut UninitSlice {
from(slice: &'a mut [u8]) -> Self217     fn from(slice: &'a mut [u8]) -> Self {
218         UninitSlice::new(slice)
219     }
220 }
221 
222 impl<'a> From<&'a mut [MaybeUninit<u8>]> for &'a mut UninitSlice {
from(slice: &'a mut [MaybeUninit<u8>]) -> Self223     fn from(slice: &'a mut [MaybeUninit<u8>]) -> Self {
224         UninitSlice::uninit(slice)
225     }
226 }
227 
228 macro_rules! impl_index {
229     ($($t:ty),*) => {
230         $(
231             impl Index<$t> for UninitSlice {
232                 type Output = UninitSlice;
233 
234                 #[inline]
235                 fn index(&self, index: $t) -> &UninitSlice {
236                     UninitSlice::uninit_ref(&self.0[index])
237                 }
238             }
239 
240             impl IndexMut<$t> for UninitSlice {
241                 #[inline]
242                 fn index_mut(&mut self, index: $t) -> &mut UninitSlice {
243                     UninitSlice::uninit(&mut self.0[index])
244                 }
245             }
246         )*
247     };
248 }
249 
250 impl_index!(
251     Range<usize>,
252     RangeFrom<usize>,
253     RangeFull,
254     RangeInclusive<usize>,
255     RangeTo<usize>,
256     RangeToInclusive<usize>
257 );
258