1use core::fmt;
2use core::mem::MaybeUninit;
3use 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)]
22pub struct UninitSlice([MaybeUninit<u8>]);
23
24impl 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]
36 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]
55 pub fn uninit(slice: &mut [MaybeUninit<u8>]) -> &mut UninitSlice {
56 unsafe { &mut *(slice as *mut [MaybeUninit<u8>] as *mut UninitSlice) }
57 }
58
59 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]
82 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]
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]
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]
161 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 uninitalized
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]
187 pub unsafe fn as_uninit_slice_mut<'a>(&'a mut self) -> &'a mut [MaybeUninit<u8>] {
188 &mut *(self as *mut _ as *mut [MaybeUninit<u8>])
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]
205 pub fn len(&self) -> usize {
206 self.0.len()
207 }
208}
209
210impl fmt::Debug for UninitSlice {
211 fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
212 fmt.debug_struct("UninitSlice[...]").finish()
213 }
214}
215
216impl<'a> From<&'a mut [u8]> for &'a mut UninitSlice {
217 fn from(slice: &'a mut [u8]) -> Self {
218 UninitSlice::new(slice)
219 }
220}
221
222impl<'a> From<&'a mut [MaybeUninit<u8>]> for &'a mut UninitSlice {
223 fn from(slice: &'a mut [MaybeUninit<u8>]) -> Self {
224 UninitSlice::uninit(slice)
225 }
226}
227
228macro_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
250impl_index!(
251 Range<usize>,
252 RangeFrom<usize>,
253 RangeFull,
254 RangeInclusive<usize>,
255 RangeTo<usize>,
256 RangeToInclusive<usize>
257);
258