1 use core::iter::FromIterator;
2 use core::mem::{self, ManuallyDrop, MaybeUninit};
3 use core::ops::{Deref, DerefMut};
4 use core::ptr::{self, NonNull};
5 use core::{cmp, fmt, hash, isize, slice, usize};
6 
7 use alloc::{
8     borrow::{Borrow, BorrowMut},
9     boxed::Box,
10     string::String,
11     vec,
12     vec::Vec,
13 };
14 
15 use crate::buf::{IntoIter, UninitSlice};
16 use crate::bytes::Vtable;
17 #[allow(unused)]
18 use crate::loom::sync::atomic::AtomicMut;
19 use crate::loom::sync::atomic::{AtomicPtr, AtomicUsize, Ordering};
20 use crate::{offset_from, Buf, BufMut, Bytes};
21 
22 /// A unique reference to a contiguous slice of memory.
23 ///
24 /// `BytesMut` represents a unique view into a potentially shared memory region.
25 /// Given the uniqueness guarantee, owners of `BytesMut` handles are able to
26 /// mutate the memory.
27 ///
28 /// `BytesMut` can be thought of as containing a `buf: Arc<Vec<u8>>`, an offset
29 /// into `buf`, a slice length, and a guarantee that no other `BytesMut` for the
30 /// same `buf` overlaps with its slice. That guarantee means that a write lock
31 /// is not required.
32 ///
33 /// # Growth
34 ///
35 /// `BytesMut`'s `BufMut` implementation will implicitly grow its buffer as
36 /// necessary. However, explicitly reserving the required space up-front before
37 /// a series of inserts will be more efficient.
38 ///
39 /// # Examples
40 ///
41 /// ```
42 /// use bytes::{BytesMut, BufMut};
43 ///
44 /// let mut buf = BytesMut::with_capacity(64);
45 ///
46 /// buf.put_u8(b'h');
47 /// buf.put_u8(b'e');
48 /// buf.put(&b"llo"[..]);
49 ///
50 /// assert_eq!(&buf[..], b"hello");
51 ///
52 /// // Freeze the buffer so that it can be shared
53 /// let a = buf.freeze();
54 ///
55 /// // This does not allocate, instead `b` points to the same memory.
56 /// let b = a.clone();
57 ///
58 /// assert_eq!(&a[..], b"hello");
59 /// assert_eq!(&b[..], b"hello");
60 /// ```
61 pub struct BytesMut {
62     ptr: NonNull<u8>,
63     len: usize,
64     cap: usize,
65     data: *mut Shared,
66 }
67 
68 // Thread-safe reference-counted container for the shared storage. This mostly
69 // the same as `core::sync::Arc` but without the weak counter. The ref counting
70 // fns are based on the ones found in `std`.
71 //
72 // The main reason to use `Shared` instead of `core::sync::Arc` is that it ends
73 // up making the overall code simpler and easier to reason about. This is due to
74 // some of the logic around setting `Inner::arc` and other ways the `arc` field
75 // is used. Using `Arc` ended up requiring a number of funky transmutes and
76 // other shenanigans to make it work.
77 struct Shared {
78     vec: Vec<u8>,
79     original_capacity_repr: usize,
80     ref_count: AtomicUsize,
81 }
82 
83 // Assert that the alignment of `Shared` is divisible by 2.
84 // This is a necessary invariant since we depend on allocating `Shared` a
85 // shared object to implicitly carry the `KIND_ARC` flag in its pointer.
86 // This flag is set when the LSB is 0.
87 const _: [(); 0 - mem::align_of::<Shared>() % 2] = []; // Assert that the alignment of `Shared` is divisible by 2.
88 
89 // Buffer storage strategy flags.
90 const KIND_ARC: usize = 0b0;
91 const KIND_VEC: usize = 0b1;
92 const KIND_MASK: usize = 0b1;
93 
94 // The max original capacity value. Any `Bytes` allocated with a greater initial
95 // capacity will default to this.
96 const MAX_ORIGINAL_CAPACITY_WIDTH: usize = 17;
97 // The original capacity algorithm will not take effect unless the originally
98 // allocated capacity was at least 1kb in size.
99 const MIN_ORIGINAL_CAPACITY_WIDTH: usize = 10;
100 // The original capacity is stored in powers of 2 starting at 1kb to a max of
101 // 64kb. Representing it as such requires only 3 bits of storage.
102 const ORIGINAL_CAPACITY_MASK: usize = 0b11100;
103 const ORIGINAL_CAPACITY_OFFSET: usize = 2;
104 
105 const VEC_POS_OFFSET: usize = 5;
106 // When the storage is in the `Vec` representation, the pointer can be advanced
107 // at most this value. This is due to the amount of storage available to track
108 // the offset is usize - number of KIND bits and number of ORIGINAL_CAPACITY
109 // bits.
110 const MAX_VEC_POS: usize = usize::MAX >> VEC_POS_OFFSET;
111 const NOT_VEC_POS_MASK: usize = 0b11111;
112 
113 #[cfg(target_pointer_width = "64")]
114 const PTR_WIDTH: usize = 64;
115 #[cfg(target_pointer_width = "32")]
116 const PTR_WIDTH: usize = 32;
117 
118 /*
119  *
120  * ===== BytesMut =====
121  *
122  */
123 
124 impl BytesMut {
125     /// Creates a new `BytesMut` with the specified capacity.
126     ///
127     /// The returned `BytesMut` will be able to hold at least `capacity` bytes
128     /// without reallocating.
129     ///
130     /// It is important to note that this function does not specify the length
131     /// of the returned `BytesMut`, but only the capacity.
132     ///
133     /// # Examples
134     ///
135     /// ```
136     /// use bytes::{BytesMut, BufMut};
137     ///
138     /// let mut bytes = BytesMut::with_capacity(64);
139     ///
140     /// // `bytes` contains no data, even though there is capacity
141     /// assert_eq!(bytes.len(), 0);
142     ///
143     /// bytes.put(&b"hello world"[..]);
144     ///
145     /// assert_eq!(&bytes[..], b"hello world");
146     /// ```
147     #[inline]
with_capacity(capacity: usize) -> BytesMut148     pub fn with_capacity(capacity: usize) -> BytesMut {
149         BytesMut::from_vec(Vec::with_capacity(capacity))
150     }
151 
152     /// Creates a new `BytesMut` with default capacity.
153     ///
154     /// Resulting object has length 0 and unspecified capacity.
155     /// This function does not allocate.
156     ///
157     /// # Examples
158     ///
159     /// ```
160     /// use bytes::{BytesMut, BufMut};
161     ///
162     /// let mut bytes = BytesMut::new();
163     ///
164     /// assert_eq!(0, bytes.len());
165     ///
166     /// bytes.reserve(2);
167     /// bytes.put_slice(b"xy");
168     ///
169     /// assert_eq!(&b"xy"[..], &bytes[..]);
170     /// ```
171     #[inline]
new() -> BytesMut172     pub fn new() -> BytesMut {
173         BytesMut::with_capacity(0)
174     }
175 
176     /// Returns the number of bytes contained in this `BytesMut`.
177     ///
178     /// # Examples
179     ///
180     /// ```
181     /// use bytes::BytesMut;
182     ///
183     /// let b = BytesMut::from(&b"hello"[..]);
184     /// assert_eq!(b.len(), 5);
185     /// ```
186     #[inline]
len(&self) -> usize187     pub fn len(&self) -> usize {
188         self.len
189     }
190 
191     /// Returns true if the `BytesMut` has a length of 0.
192     ///
193     /// # Examples
194     ///
195     /// ```
196     /// use bytes::BytesMut;
197     ///
198     /// let b = BytesMut::with_capacity(64);
199     /// assert!(b.is_empty());
200     /// ```
201     #[inline]
is_empty(&self) -> bool202     pub fn is_empty(&self) -> bool {
203         self.len == 0
204     }
205 
206     /// Returns the number of bytes the `BytesMut` can hold without reallocating.
207     ///
208     /// # Examples
209     ///
210     /// ```
211     /// use bytes::BytesMut;
212     ///
213     /// let b = BytesMut::with_capacity(64);
214     /// assert_eq!(b.capacity(), 64);
215     /// ```
216     #[inline]
capacity(&self) -> usize217     pub fn capacity(&self) -> usize {
218         self.cap
219     }
220 
221     /// Converts `self` into an immutable `Bytes`.
222     ///
223     /// The conversion is zero cost and is used to indicate that the slice
224     /// referenced by the handle will no longer be mutated. Once the conversion
225     /// is done, the handle can be cloned and shared across threads.
226     ///
227     /// # Examples
228     ///
229     /// ```
230     /// use bytes::{BytesMut, BufMut};
231     /// use std::thread;
232     ///
233     /// let mut b = BytesMut::with_capacity(64);
234     /// b.put(&b"hello world"[..]);
235     /// let b1 = b.freeze();
236     /// let b2 = b1.clone();
237     ///
238     /// let th = thread::spawn(move || {
239     ///     assert_eq!(&b1[..], b"hello world");
240     /// });
241     ///
242     /// assert_eq!(&b2[..], b"hello world");
243     /// th.join().unwrap();
244     /// ```
245     #[inline]
freeze(self) -> Bytes246     pub fn freeze(self) -> Bytes {
247         let bytes = ManuallyDrop::new(self);
248         if bytes.kind() == KIND_VEC {
249             // Just re-use `Bytes` internal Vec vtable
250             unsafe {
251                 let off = bytes.get_vec_pos();
252                 let vec = rebuild_vec(bytes.ptr.as_ptr(), bytes.len, bytes.cap, off);
253                 let mut b: Bytes = vec.into();
254                 b.advance(off);
255                 b
256             }
257         } else {
258             debug_assert_eq!(bytes.kind(), KIND_ARC);
259 
260             let ptr = bytes.ptr.as_ptr();
261             let len = bytes.len;
262             let data = AtomicPtr::new(bytes.data.cast());
263             unsafe { Bytes::with_vtable(ptr, len, data, &SHARED_VTABLE) }
264         }
265     }
266 
267     /// Creates a new `BytesMut` containing `len` zeros.
268     ///
269     /// The resulting object has a length of `len` and a capacity greater
270     /// than or equal to `len`. The entire length of the object will be filled
271     /// with zeros.
272     ///
273     /// On some platforms or allocators this function may be faster than
274     /// a manual implementation.
275     ///
276     /// # Examples
277     ///
278     /// ```
279     /// use bytes::BytesMut;
280     ///
281     /// let zeros = BytesMut::zeroed(42);
282     ///
283     /// assert!(zeros.capacity() >= 42);
284     /// assert_eq!(zeros.len(), 42);
285     /// zeros.into_iter().for_each(|x| assert_eq!(x, 0));
286     /// ```
zeroed(len: usize) -> BytesMut287     pub fn zeroed(len: usize) -> BytesMut {
288         BytesMut::from_vec(vec![0; len])
289     }
290 
291     /// Splits the bytes into two at the given index.
292     ///
293     /// Afterwards `self` contains elements `[0, at)`, and the returned
294     /// `BytesMut` contains elements `[at, capacity)`. It's guaranteed that the
295     /// memory does not move, that is, the address of `self` does not change,
296     /// and the address of the returned slice is `at` bytes after that.
297     ///
298     /// This is an `O(1)` operation that just increases the reference count
299     /// and sets a few indices.
300     ///
301     /// # Examples
302     ///
303     /// ```
304     /// use bytes::BytesMut;
305     ///
306     /// let mut a = BytesMut::from(&b"hello world"[..]);
307     /// let mut b = a.split_off(5);
308     ///
309     /// a[0] = b'j';
310     /// b[0] = b'!';
311     ///
312     /// assert_eq!(&a[..], b"jello");
313     /// assert_eq!(&b[..], b"!world");
314     /// ```
315     ///
316     /// # Panics
317     ///
318     /// Panics if `at > capacity`.
319     #[must_use = "consider BytesMut::truncate if you don't need the other half"]
split_off(&mut self, at: usize) -> BytesMut320     pub fn split_off(&mut self, at: usize) -> BytesMut {
321         assert!(
322             at <= self.capacity(),
323             "split_off out of bounds: {:?} <= {:?}",
324             at,
325             self.capacity(),
326         );
327         unsafe {
328             let mut other = self.shallow_clone();
329             // SAFETY: We've checked that `at` <= `self.capacity()` above.
330             other.advance_unchecked(at);
331             self.cap = at;
332             self.len = cmp::min(self.len, at);
333             other
334         }
335     }
336 
337     /// Removes the bytes from the current view, returning them in a new
338     /// `BytesMut` handle.
339     ///
340     /// Afterwards, `self` will be empty, but will retain any additional
341     /// capacity that it had before the operation. This is identical to
342     /// `self.split_to(self.len())`.
343     ///
344     /// This is an `O(1)` operation that just increases the reference count and
345     /// sets a few indices.
346     ///
347     /// # Examples
348     ///
349     /// ```
350     /// use bytes::{BytesMut, BufMut};
351     ///
352     /// let mut buf = BytesMut::with_capacity(1024);
353     /// buf.put(&b"hello world"[..]);
354     ///
355     /// let other = buf.split();
356     ///
357     /// assert!(buf.is_empty());
358     /// assert_eq!(1013, buf.capacity());
359     ///
360     /// assert_eq!(other, b"hello world"[..]);
361     /// ```
362     #[must_use = "consider BytesMut::clear if you don't need the other half"]
split(&mut self) -> BytesMut363     pub fn split(&mut self) -> BytesMut {
364         let len = self.len();
365         self.split_to(len)
366     }
367 
368     /// Splits the buffer into two at the given index.
369     ///
370     /// Afterwards `self` contains elements `[at, len)`, and the returned `BytesMut`
371     /// contains elements `[0, at)`.
372     ///
373     /// This is an `O(1)` operation that just increases the reference count and
374     /// sets a few indices.
375     ///
376     /// # Examples
377     ///
378     /// ```
379     /// use bytes::BytesMut;
380     ///
381     /// let mut a = BytesMut::from(&b"hello world"[..]);
382     /// let mut b = a.split_to(5);
383     ///
384     /// a[0] = b'!';
385     /// b[0] = b'j';
386     ///
387     /// assert_eq!(&a[..], b"!world");
388     /// assert_eq!(&b[..], b"jello");
389     /// ```
390     ///
391     /// # Panics
392     ///
393     /// Panics if `at > len`.
394     #[must_use = "consider BytesMut::advance if you don't need the other half"]
split_to(&mut self, at: usize) -> BytesMut395     pub fn split_to(&mut self, at: usize) -> BytesMut {
396         assert!(
397             at <= self.len(),
398             "split_to out of bounds: {:?} <= {:?}",
399             at,
400             self.len(),
401         );
402 
403         unsafe {
404             let mut other = self.shallow_clone();
405             // SAFETY: We've checked that `at` <= `self.len()` and we know that `self.len()` <=
406             // `self.capacity()`.
407             self.advance_unchecked(at);
408             other.cap = at;
409             other.len = at;
410             other
411         }
412     }
413 
414     /// Shortens the buffer, keeping the first `len` bytes and dropping the
415     /// rest.
416     ///
417     /// If `len` is greater than the buffer's current length, this has no
418     /// effect.
419     ///
420     /// Existing underlying capacity is preserved.
421     ///
422     /// The [split_off](`Self::split_off()`) method can emulate `truncate`, but this causes the
423     /// excess bytes to be returned instead of dropped.
424     ///
425     /// # Examples
426     ///
427     /// ```
428     /// use bytes::BytesMut;
429     ///
430     /// let mut buf = BytesMut::from(&b"hello world"[..]);
431     /// buf.truncate(5);
432     /// assert_eq!(buf, b"hello"[..]);
433     /// ```
truncate(&mut self, len: usize)434     pub fn truncate(&mut self, len: usize) {
435         if len <= self.len() {
436             // SAFETY: Shrinking the buffer cannot expose uninitialized bytes.
437             unsafe { self.set_len(len) };
438         }
439     }
440 
441     /// Clears the buffer, removing all data. Existing capacity is preserved.
442     ///
443     /// # Examples
444     ///
445     /// ```
446     /// use bytes::BytesMut;
447     ///
448     /// let mut buf = BytesMut::from(&b"hello world"[..]);
449     /// buf.clear();
450     /// assert!(buf.is_empty());
451     /// ```
clear(&mut self)452     pub fn clear(&mut self) {
453         // SAFETY: Setting the length to zero cannot expose uninitialized bytes.
454         unsafe { self.set_len(0) };
455     }
456 
457     /// Resizes the buffer so that `len` is equal to `new_len`.
458     ///
459     /// If `new_len` is greater than `len`, the buffer is extended by the
460     /// difference with each additional byte set to `value`. If `new_len` is
461     /// less than `len`, the buffer is simply truncated.
462     ///
463     /// # Examples
464     ///
465     /// ```
466     /// use bytes::BytesMut;
467     ///
468     /// let mut buf = BytesMut::new();
469     ///
470     /// buf.resize(3, 0x1);
471     /// assert_eq!(&buf[..], &[0x1, 0x1, 0x1]);
472     ///
473     /// buf.resize(2, 0x2);
474     /// assert_eq!(&buf[..], &[0x1, 0x1]);
475     ///
476     /// buf.resize(4, 0x3);
477     /// assert_eq!(&buf[..], &[0x1, 0x1, 0x3, 0x3]);
478     /// ```
resize(&mut self, new_len: usize, value: u8)479     pub fn resize(&mut self, new_len: usize, value: u8) {
480         let additional = if let Some(additional) = new_len.checked_sub(self.len()) {
481             additional
482         } else {
483             self.truncate(new_len);
484             return;
485         };
486 
487         if additional == 0 {
488             return;
489         }
490 
491         self.reserve(additional);
492         let dst = self.spare_capacity_mut().as_mut_ptr();
493         // SAFETY: `spare_capacity_mut` returns a valid, properly aligned pointer and we've
494         // reserved enough space to write `additional` bytes.
495         unsafe { ptr::write_bytes(dst, value, additional) };
496 
497         // SAFETY: There are at least `new_len` initialized bytes in the buffer so no
498         // uninitialized bytes are being exposed.
499         unsafe { self.set_len(new_len) };
500     }
501 
502     /// Sets the length of the buffer.
503     ///
504     /// This will explicitly set the size of the buffer without actually
505     /// modifying the data, so it is up to the caller to ensure that the data
506     /// has been initialized.
507     ///
508     /// # Examples
509     ///
510     /// ```
511     /// use bytes::BytesMut;
512     ///
513     /// let mut b = BytesMut::from(&b"hello world"[..]);
514     ///
515     /// unsafe {
516     ///     b.set_len(5);
517     /// }
518     ///
519     /// assert_eq!(&b[..], b"hello");
520     ///
521     /// unsafe {
522     ///     b.set_len(11);
523     /// }
524     ///
525     /// assert_eq!(&b[..], b"hello world");
526     /// ```
527     #[inline]
set_len(&mut self, len: usize)528     pub unsafe fn set_len(&mut self, len: usize) {
529         debug_assert!(len <= self.cap, "set_len out of bounds");
530         self.len = len;
531     }
532 
533     /// Reserves capacity for at least `additional` more bytes to be inserted
534     /// into the given `BytesMut`.
535     ///
536     /// More than `additional` bytes may be reserved in order to avoid frequent
537     /// reallocations. A call to `reserve` may result in an allocation.
538     ///
539     /// Before allocating new buffer space, the function will attempt to reclaim
540     /// space in the existing buffer. If the current handle references a view
541     /// into a larger original buffer, and all other handles referencing part
542     /// of the same original buffer have been dropped, then the current view
543     /// can be copied/shifted to the front of the buffer and the handle can take
544     /// ownership of the full buffer, provided that the full buffer is large
545     /// enough to fit the requested additional capacity.
546     ///
547     /// This optimization will only happen if shifting the data from the current
548     /// view to the front of the buffer is not too expensive in terms of the
549     /// (amortized) time required. The precise condition is subject to change;
550     /// as of now, the length of the data being shifted needs to be at least as
551     /// large as the distance that it's shifted by. If the current view is empty
552     /// and the original buffer is large enough to fit the requested additional
553     /// capacity, then reallocations will never happen.
554     ///
555     /// # Examples
556     ///
557     /// In the following example, a new buffer is allocated.
558     ///
559     /// ```
560     /// use bytes::BytesMut;
561     ///
562     /// let mut buf = BytesMut::from(&b"hello"[..]);
563     /// buf.reserve(64);
564     /// assert!(buf.capacity() >= 69);
565     /// ```
566     ///
567     /// In the following example, the existing buffer is reclaimed.
568     ///
569     /// ```
570     /// use bytes::{BytesMut, BufMut};
571     ///
572     /// let mut buf = BytesMut::with_capacity(128);
573     /// buf.put(&[0; 64][..]);
574     ///
575     /// let ptr = buf.as_ptr();
576     /// let other = buf.split();
577     ///
578     /// assert!(buf.is_empty());
579     /// assert_eq!(buf.capacity(), 64);
580     ///
581     /// drop(other);
582     /// buf.reserve(128);
583     ///
584     /// assert_eq!(buf.capacity(), 128);
585     /// assert_eq!(buf.as_ptr(), ptr);
586     /// ```
587     ///
588     /// # Panics
589     ///
590     /// Panics if the new capacity overflows `usize`.
591     #[inline]
reserve(&mut self, additional: usize)592     pub fn reserve(&mut self, additional: usize) {
593         let len = self.len();
594         let rem = self.capacity() - len;
595 
596         if additional <= rem {
597             // The handle can already store at least `additional` more bytes, so
598             // there is no further work needed to be done.
599             return;
600         }
601 
602         // will always succeed
603         let _ = self.reserve_inner(additional, true);
604     }
605 
606     // In separate function to allow the short-circuits in `reserve` and `try_reclaim` to
607     // be inline-able. Significantly helps performance. Returns false if it did not succeed.
reserve_inner(&mut self, additional: usize, allocate: bool) -> bool608     fn reserve_inner(&mut self, additional: usize, allocate: bool) -> bool {
609         let len = self.len();
610         let kind = self.kind();
611 
612         if kind == KIND_VEC {
613             // If there's enough free space before the start of the buffer, then
614             // just copy the data backwards and reuse the already-allocated
615             // space.
616             //
617             // Otherwise, since backed by a vector, use `Vec::reserve`
618             //
619             // We need to make sure that this optimization does not kill the
620             // amortized runtimes of BytesMut's operations.
621             unsafe {
622                 let off = self.get_vec_pos();
623 
624                 // Only reuse space if we can satisfy the requested additional space.
625                 //
626                 // Also check if the value of `off` suggests that enough bytes
627                 // have been read to account for the overhead of shifting all
628                 // the data (in an amortized analysis).
629                 // Hence the condition `off >= self.len()`.
630                 //
631                 // This condition also already implies that the buffer is going
632                 // to be (at least) half-empty in the end; so we do not break
633                 // the (amortized) runtime with future resizes of the underlying
634                 // `Vec`.
635                 //
636                 // [For more details check issue #524, and PR #525.]
637                 if self.capacity() - self.len() + off >= additional && off >= self.len() {
638                     // There's enough space, and it's not too much overhead:
639                     // reuse the space!
640                     //
641                     // Just move the pointer back to the start after copying
642                     // data back.
643                     let base_ptr = self.ptr.as_ptr().sub(off);
644                     // Since `off >= self.len()`, the two regions don't overlap.
645                     ptr::copy_nonoverlapping(self.ptr.as_ptr(), base_ptr, self.len);
646                     self.ptr = vptr(base_ptr);
647                     self.set_vec_pos(0);
648 
649                     // Length stays constant, but since we moved backwards we
650                     // can gain capacity back.
651                     self.cap += off;
652                 } else {
653                     if !allocate {
654                         return false;
655                     }
656                     // Not enough space, or reusing might be too much overhead:
657                     // allocate more space!
658                     let mut v =
659                         ManuallyDrop::new(rebuild_vec(self.ptr.as_ptr(), self.len, self.cap, off));
660                     v.reserve(additional);
661 
662                     // Update the info
663                     self.ptr = vptr(v.as_mut_ptr().add(off));
664                     self.cap = v.capacity() - off;
665                     debug_assert_eq!(self.len, v.len() - off);
666                 }
667 
668                 return true;
669             }
670         }
671 
672         debug_assert_eq!(kind, KIND_ARC);
673         let shared: *mut Shared = self.data;
674 
675         // Reserving involves abandoning the currently shared buffer and
676         // allocating a new vector with the requested capacity.
677         //
678         // Compute the new capacity
679         let mut new_cap = match len.checked_add(additional) {
680             Some(new_cap) => new_cap,
681             None if !allocate => return false,
682             None => panic!("overflow"),
683         };
684 
685         unsafe {
686             // First, try to reclaim the buffer. This is possible if the current
687             // handle is the only outstanding handle pointing to the buffer.
688             if (*shared).is_unique() {
689                 // This is the only handle to the buffer. It can be reclaimed.
690                 // However, before doing the work of copying data, check to make
691                 // sure that the vector has enough capacity.
692                 let v = &mut (*shared).vec;
693 
694                 let v_capacity = v.capacity();
695                 let ptr = v.as_mut_ptr();
696 
697                 let offset = offset_from(self.ptr.as_ptr(), ptr);
698 
699                 // Compare the condition in the `kind == KIND_VEC` case above
700                 // for more details.
701                 if v_capacity >= new_cap + offset {
702                     self.cap = new_cap;
703                     // no copy is necessary
704                 } else if v_capacity >= new_cap && offset >= len {
705                     // The capacity is sufficient, and copying is not too much
706                     // overhead: reclaim the buffer!
707 
708                     // `offset >= len` means: no overlap
709                     ptr::copy_nonoverlapping(self.ptr.as_ptr(), ptr, len);
710 
711                     self.ptr = vptr(ptr);
712                     self.cap = v.capacity();
713                 } else {
714                     if !allocate {
715                         return false;
716                     }
717                     // calculate offset
718                     let off = (self.ptr.as_ptr() as usize) - (v.as_ptr() as usize);
719 
720                     // new_cap is calculated in terms of `BytesMut`, not the underlying
721                     // `Vec`, so it does not take the offset into account.
722                     //
723                     // Thus we have to manually add it here.
724                     new_cap = new_cap.checked_add(off).expect("overflow");
725 
726                     // The vector capacity is not sufficient. The reserve request is
727                     // asking for more than the initial buffer capacity. Allocate more
728                     // than requested if `new_cap` is not much bigger than the current
729                     // capacity.
730                     //
731                     // There are some situations, using `reserve_exact` that the
732                     // buffer capacity could be below `original_capacity`, so do a
733                     // check.
734                     let double = v.capacity().checked_shl(1).unwrap_or(new_cap);
735 
736                     new_cap = cmp::max(double, new_cap);
737 
738                     // No space - allocate more
739                     //
740                     // The length field of `Shared::vec` is not used by the `BytesMut`;
741                     // instead we use the `len` field in the `BytesMut` itself. However,
742                     // when calling `reserve`, it doesn't guarantee that data stored in
743                     // the unused capacity of the vector is copied over to the new
744                     // allocation, so we need to ensure that we don't have any data we
745                     // care about in the unused capacity before calling `reserve`.
746                     debug_assert!(off + len <= v.capacity());
747                     v.set_len(off + len);
748                     v.reserve(new_cap - v.len());
749 
750                     // Update the info
751                     self.ptr = vptr(v.as_mut_ptr().add(off));
752                     self.cap = v.capacity() - off;
753                 }
754 
755                 return true;
756             }
757         }
758         if !allocate {
759             return false;
760         }
761 
762         let original_capacity_repr = unsafe { (*shared).original_capacity_repr };
763         let original_capacity = original_capacity_from_repr(original_capacity_repr);
764 
765         new_cap = cmp::max(new_cap, original_capacity);
766 
767         // Create a new vector to store the data
768         let mut v = ManuallyDrop::new(Vec::with_capacity(new_cap));
769 
770         // Copy the bytes
771         v.extend_from_slice(self.as_ref());
772 
773         // Release the shared handle. This must be done *after* the bytes are
774         // copied.
775         unsafe { release_shared(shared) };
776 
777         // Update self
778         let data = (original_capacity_repr << ORIGINAL_CAPACITY_OFFSET) | KIND_VEC;
779         self.data = invalid_ptr(data);
780         self.ptr = vptr(v.as_mut_ptr());
781         self.cap = v.capacity();
782         debug_assert_eq!(self.len, v.len());
783         return true;
784     }
785 
786     /// Attempts to cheaply reclaim already allocated capacity for at least `additional` more
787     /// bytes to be inserted into the given `BytesMut` and returns `true` if it succeeded.
788     ///
789     /// `try_reclaim` behaves exactly like `reserve`, except that it never allocates new storage
790     /// and returns a `bool` indicating whether it was successful in doing so:
791     ///
792     /// `try_reclaim` returns false under these conditions:
793     ///  - The spare capacity left is less than `additional` bytes AND
794     ///  - The existing allocation cannot be reclaimed cheaply or it was less than
795     ///    `additional` bytes in size
796     ///
797     /// Reclaiming the allocation cheaply is possible if the `BytesMut` has no outstanding
798     /// references through other `BytesMut`s or `Bytes` which point to the same underlying
799     /// storage.
800     ///
801     /// # Examples
802     ///
803     /// ```
804     /// use bytes::BytesMut;
805     ///
806     /// let mut buf = BytesMut::with_capacity(64);
807     /// assert_eq!(true, buf.try_reclaim(64));
808     /// assert_eq!(64, buf.capacity());
809     ///
810     /// buf.extend_from_slice(b"abcd");
811     /// let mut split = buf.split();
812     /// assert_eq!(60, buf.capacity());
813     /// assert_eq!(4, split.capacity());
814     /// assert_eq!(false, split.try_reclaim(64));
815     /// assert_eq!(false, buf.try_reclaim(64));
816     /// // The split buffer is filled with "abcd"
817     /// assert_eq!(false, split.try_reclaim(4));
818     /// // buf is empty and has capacity for 60 bytes
819     /// assert_eq!(true, buf.try_reclaim(60));
820     ///
821     /// drop(buf);
822     /// assert_eq!(false, split.try_reclaim(64));
823     ///
824     /// split.clear();
825     /// assert_eq!(4, split.capacity());
826     /// assert_eq!(true, split.try_reclaim(64));
827     /// assert_eq!(64, split.capacity());
828     /// ```
829     // I tried splitting out try_reclaim_inner after the short circuits, but it was inlined
830     // regardless with Rust 1.78.0 so probably not worth it
831     #[inline]
832     #[must_use = "consider BytesMut::reserve if you need an infallible reservation"]
try_reclaim(&mut self, additional: usize) -> bool833     pub fn try_reclaim(&mut self, additional: usize) -> bool {
834         let len = self.len();
835         let rem = self.capacity() - len;
836 
837         if additional <= rem {
838             // The handle can already store at least `additional` more bytes, so
839             // there is no further work needed to be done.
840             return true;
841         }
842 
843         self.reserve_inner(additional, false)
844     }
845 
846     /// Appends given bytes to this `BytesMut`.
847     ///
848     /// If this `BytesMut` object does not have enough capacity, it is resized
849     /// first.
850     ///
851     /// # Examples
852     ///
853     /// ```
854     /// use bytes::BytesMut;
855     ///
856     /// let mut buf = BytesMut::with_capacity(0);
857     /// buf.extend_from_slice(b"aaabbb");
858     /// buf.extend_from_slice(b"cccddd");
859     ///
860     /// assert_eq!(b"aaabbbcccddd", &buf[..]);
861     /// ```
862     #[inline]
extend_from_slice(&mut self, extend: &[u8])863     pub fn extend_from_slice(&mut self, extend: &[u8]) {
864         let cnt = extend.len();
865         self.reserve(cnt);
866 
867         unsafe {
868             let dst = self.spare_capacity_mut();
869             // Reserved above
870             debug_assert!(dst.len() >= cnt);
871 
872             ptr::copy_nonoverlapping(extend.as_ptr(), dst.as_mut_ptr().cast(), cnt);
873         }
874 
875         unsafe {
876             self.advance_mut(cnt);
877         }
878     }
879 
880     /// Absorbs a `BytesMut` that was previously split off.
881     ///
882     /// If the two `BytesMut` objects were previously contiguous and not mutated
883     /// in a way that causes re-allocation i.e., if `other` was created by
884     /// calling `split_off` on this `BytesMut`, then this is an `O(1)` operation
885     /// that just decreases a reference count and sets a few indices.
886     /// Otherwise this method degenerates to
887     /// `self.extend_from_slice(other.as_ref())`.
888     ///
889     /// # Examples
890     ///
891     /// ```
892     /// use bytes::BytesMut;
893     ///
894     /// let mut buf = BytesMut::with_capacity(64);
895     /// buf.extend_from_slice(b"aaabbbcccddd");
896     ///
897     /// let split = buf.split_off(6);
898     /// assert_eq!(b"aaabbb", &buf[..]);
899     /// assert_eq!(b"cccddd", &split[..]);
900     ///
901     /// buf.unsplit(split);
902     /// assert_eq!(b"aaabbbcccddd", &buf[..]);
903     /// ```
unsplit(&mut self, other: BytesMut)904     pub fn unsplit(&mut self, other: BytesMut) {
905         if self.is_empty() {
906             *self = other;
907             return;
908         }
909 
910         if let Err(other) = self.try_unsplit(other) {
911             self.extend_from_slice(other.as_ref());
912         }
913     }
914 
915     // private
916 
917     // For now, use a `Vec` to manage the memory for us, but we may want to
918     // change that in the future to some alternate allocator strategy.
919     //
920     // Thus, we don't expose an easy way to construct from a `Vec` since an
921     // internal change could make a simple pattern (`BytesMut::from(vec)`)
922     // suddenly a lot more expensive.
923     #[inline]
from_vec(vec: Vec<u8>) -> BytesMut924     pub(crate) fn from_vec(vec: Vec<u8>) -> BytesMut {
925         let mut vec = ManuallyDrop::new(vec);
926         let ptr = vptr(vec.as_mut_ptr());
927         let len = vec.len();
928         let cap = vec.capacity();
929 
930         let original_capacity_repr = original_capacity_to_repr(cap);
931         let data = (original_capacity_repr << ORIGINAL_CAPACITY_OFFSET) | KIND_VEC;
932 
933         BytesMut {
934             ptr,
935             len,
936             cap,
937             data: invalid_ptr(data),
938         }
939     }
940 
941     #[inline]
as_slice(&self) -> &[u8]942     fn as_slice(&self) -> &[u8] {
943         unsafe { slice::from_raw_parts(self.ptr.as_ptr(), self.len) }
944     }
945 
946     #[inline]
as_slice_mut(&mut self) -> &mut [u8]947     fn as_slice_mut(&mut self) -> &mut [u8] {
948         unsafe { slice::from_raw_parts_mut(self.ptr.as_ptr(), self.len) }
949     }
950 
951     /// Advance the buffer without bounds checking.
952     ///
953     /// # SAFETY
954     ///
955     /// The caller must ensure that `count` <= `self.cap`.
advance_unchecked(&mut self, count: usize)956     pub(crate) unsafe fn advance_unchecked(&mut self, count: usize) {
957         // Setting the start to 0 is a no-op, so return early if this is the
958         // case.
959         if count == 0 {
960             return;
961         }
962 
963         debug_assert!(count <= self.cap, "internal: set_start out of bounds");
964 
965         let kind = self.kind();
966 
967         if kind == KIND_VEC {
968             // Setting the start when in vec representation is a little more
969             // complicated. First, we have to track how far ahead the
970             // "start" of the byte buffer from the beginning of the vec. We
971             // also have to ensure that we don't exceed the maximum shift.
972             let pos = self.get_vec_pos() + count;
973 
974             if pos <= MAX_VEC_POS {
975                 self.set_vec_pos(pos);
976             } else {
977                 // The repr must be upgraded to ARC. This will never happen
978                 // on 64 bit systems and will only happen on 32 bit systems
979                 // when shifting past 134,217,727 bytes. As such, we don't
980                 // worry too much about performance here.
981                 self.promote_to_shared(/*ref_count = */ 1);
982             }
983         }
984 
985         // Updating the start of the view is setting `ptr` to point to the
986         // new start and updating the `len` field to reflect the new length
987         // of the view.
988         self.ptr = vptr(self.ptr.as_ptr().add(count));
989         self.len = self.len.checked_sub(count).unwrap_or(0);
990         self.cap -= count;
991     }
992 
try_unsplit(&mut self, other: BytesMut) -> Result<(), BytesMut>993     fn try_unsplit(&mut self, other: BytesMut) -> Result<(), BytesMut> {
994         if other.capacity() == 0 {
995             return Ok(());
996         }
997 
998         let ptr = unsafe { self.ptr.as_ptr().add(self.len) };
999         if ptr == other.ptr.as_ptr()
1000             && self.kind() == KIND_ARC
1001             && other.kind() == KIND_ARC
1002             && self.data == other.data
1003         {
1004             // Contiguous blocks, just combine directly
1005             self.len += other.len;
1006             self.cap += other.cap;
1007             Ok(())
1008         } else {
1009             Err(other)
1010         }
1011     }
1012 
1013     #[inline]
kind(&self) -> usize1014     fn kind(&self) -> usize {
1015         self.data as usize & KIND_MASK
1016     }
1017 
promote_to_shared(&mut self, ref_cnt: usize)1018     unsafe fn promote_to_shared(&mut self, ref_cnt: usize) {
1019         debug_assert_eq!(self.kind(), KIND_VEC);
1020         debug_assert!(ref_cnt == 1 || ref_cnt == 2);
1021 
1022         let original_capacity_repr =
1023             (self.data as usize & ORIGINAL_CAPACITY_MASK) >> ORIGINAL_CAPACITY_OFFSET;
1024 
1025         // The vec offset cannot be concurrently mutated, so there
1026         // should be no danger reading it.
1027         let off = (self.data as usize) >> VEC_POS_OFFSET;
1028 
1029         // First, allocate a new `Shared` instance containing the
1030         // `Vec` fields. It's important to note that `ptr`, `len`,
1031         // and `cap` cannot be mutated without having `&mut self`.
1032         // This means that these fields will not be concurrently
1033         // updated and since the buffer hasn't been promoted to an
1034         // `Arc`, those three fields still are the components of the
1035         // vector.
1036         let shared = Box::new(Shared {
1037             vec: rebuild_vec(self.ptr.as_ptr(), self.len, self.cap, off),
1038             original_capacity_repr,
1039             ref_count: AtomicUsize::new(ref_cnt),
1040         });
1041 
1042         let shared = Box::into_raw(shared);
1043 
1044         // The pointer should be aligned, so this assert should
1045         // always succeed.
1046         debug_assert_eq!(shared as usize & KIND_MASK, KIND_ARC);
1047 
1048         self.data = shared;
1049     }
1050 
1051     /// Makes an exact shallow clone of `self`.
1052     ///
1053     /// The kind of `self` doesn't matter, but this is unsafe
1054     /// because the clone will have the same offsets. You must
1055     /// be sure the returned value to the user doesn't allow
1056     /// two views into the same range.
1057     #[inline]
shallow_clone(&mut self) -> BytesMut1058     unsafe fn shallow_clone(&mut self) -> BytesMut {
1059         if self.kind() == KIND_ARC {
1060             increment_shared(self.data);
1061             ptr::read(self)
1062         } else {
1063             self.promote_to_shared(/*ref_count = */ 2);
1064             ptr::read(self)
1065         }
1066     }
1067 
1068     #[inline]
get_vec_pos(&self) -> usize1069     unsafe fn get_vec_pos(&self) -> usize {
1070         debug_assert_eq!(self.kind(), KIND_VEC);
1071 
1072         self.data as usize >> VEC_POS_OFFSET
1073     }
1074 
1075     #[inline]
set_vec_pos(&mut self, pos: usize)1076     unsafe fn set_vec_pos(&mut self, pos: usize) {
1077         debug_assert_eq!(self.kind(), KIND_VEC);
1078         debug_assert!(pos <= MAX_VEC_POS);
1079 
1080         self.data = invalid_ptr((pos << VEC_POS_OFFSET) | (self.data as usize & NOT_VEC_POS_MASK));
1081     }
1082 
1083     /// Returns the remaining spare capacity of the buffer as a slice of `MaybeUninit<u8>`.
1084     ///
1085     /// The returned slice can be used to fill the buffer with data (e.g. by
1086     /// reading from a file) before marking the data as initialized using the
1087     /// [`set_len`] method.
1088     ///
1089     /// [`set_len`]: BytesMut::set_len
1090     ///
1091     /// # Examples
1092     ///
1093     /// ```
1094     /// use bytes::BytesMut;
1095     ///
1096     /// // Allocate buffer big enough for 10 bytes.
1097     /// let mut buf = BytesMut::with_capacity(10);
1098     ///
1099     /// // Fill in the first 3 elements.
1100     /// let uninit = buf.spare_capacity_mut();
1101     /// uninit[0].write(0);
1102     /// uninit[1].write(1);
1103     /// uninit[2].write(2);
1104     ///
1105     /// // Mark the first 3 bytes of the buffer as being initialized.
1106     /// unsafe {
1107     ///     buf.set_len(3);
1108     /// }
1109     ///
1110     /// assert_eq!(&buf[..], &[0, 1, 2]);
1111     /// ```
1112     #[inline]
spare_capacity_mut(&mut self) -> &mut [MaybeUninit<u8>]1113     pub fn spare_capacity_mut(&mut self) -> &mut [MaybeUninit<u8>] {
1114         unsafe {
1115             let ptr = self.ptr.as_ptr().add(self.len);
1116             let len = self.cap - self.len;
1117 
1118             slice::from_raw_parts_mut(ptr.cast(), len)
1119         }
1120     }
1121 }
1122 
1123 impl Drop for BytesMut {
drop(&mut self)1124     fn drop(&mut self) {
1125         let kind = self.kind();
1126 
1127         if kind == KIND_VEC {
1128             unsafe {
1129                 let off = self.get_vec_pos();
1130 
1131                 // Vector storage, free the vector
1132                 let _ = rebuild_vec(self.ptr.as_ptr(), self.len, self.cap, off);
1133             }
1134         } else if kind == KIND_ARC {
1135             unsafe { release_shared(self.data) };
1136         }
1137     }
1138 }
1139 
1140 impl Buf for BytesMut {
1141     #[inline]
remaining(&self) -> usize1142     fn remaining(&self) -> usize {
1143         self.len()
1144     }
1145 
1146     #[inline]
chunk(&self) -> &[u8]1147     fn chunk(&self) -> &[u8] {
1148         self.as_slice()
1149     }
1150 
1151     #[inline]
advance(&mut self, cnt: usize)1152     fn advance(&mut self, cnt: usize) {
1153         assert!(
1154             cnt <= self.remaining(),
1155             "cannot advance past `remaining`: {:?} <= {:?}",
1156             cnt,
1157             self.remaining(),
1158         );
1159         unsafe {
1160             // SAFETY: We've checked that `cnt` <= `self.remaining()` and we know that
1161             // `self.remaining()` <= `self.cap`.
1162             self.advance_unchecked(cnt);
1163         }
1164     }
1165 
copy_to_bytes(&mut self, len: usize) -> Bytes1166     fn copy_to_bytes(&mut self, len: usize) -> Bytes {
1167         self.split_to(len).freeze()
1168     }
1169 }
1170 
1171 unsafe impl BufMut for BytesMut {
1172     #[inline]
remaining_mut(&self) -> usize1173     fn remaining_mut(&self) -> usize {
1174         usize::MAX - self.len()
1175     }
1176 
1177     #[inline]
advance_mut(&mut self, cnt: usize)1178     unsafe fn advance_mut(&mut self, cnt: usize) {
1179         let remaining = self.cap - self.len();
1180         if cnt > remaining {
1181             super::panic_advance(cnt, remaining);
1182         }
1183         // Addition won't overflow since it is at most `self.cap`.
1184         self.len = self.len() + cnt;
1185     }
1186 
1187     #[inline]
chunk_mut(&mut self) -> &mut UninitSlice1188     fn chunk_mut(&mut self) -> &mut UninitSlice {
1189         if self.capacity() == self.len() {
1190             self.reserve(64);
1191         }
1192         self.spare_capacity_mut().into()
1193     }
1194 
1195     // Specialize these methods so they can skip checking `remaining_mut`
1196     // and `advance_mut`.
1197 
put<T: Buf>(&mut self, mut src: T) where Self: Sized,1198     fn put<T: Buf>(&mut self, mut src: T)
1199     where
1200         Self: Sized,
1201     {
1202         while src.has_remaining() {
1203             let s = src.chunk();
1204             let l = s.len();
1205             self.extend_from_slice(s);
1206             src.advance(l);
1207         }
1208     }
1209 
put_slice(&mut self, src: &[u8])1210     fn put_slice(&mut self, src: &[u8]) {
1211         self.extend_from_slice(src);
1212     }
1213 
put_bytes(&mut self, val: u8, cnt: usize)1214     fn put_bytes(&mut self, val: u8, cnt: usize) {
1215         self.reserve(cnt);
1216         unsafe {
1217             let dst = self.spare_capacity_mut();
1218             // Reserved above
1219             debug_assert!(dst.len() >= cnt);
1220 
1221             ptr::write_bytes(dst.as_mut_ptr(), val, cnt);
1222 
1223             self.advance_mut(cnt);
1224         }
1225     }
1226 }
1227 
1228 impl AsRef<[u8]> for BytesMut {
1229     #[inline]
as_ref(&self) -> &[u8]1230     fn as_ref(&self) -> &[u8] {
1231         self.as_slice()
1232     }
1233 }
1234 
1235 impl Deref for BytesMut {
1236     type Target = [u8];
1237 
1238     #[inline]
deref(&self) -> &[u8]1239     fn deref(&self) -> &[u8] {
1240         self.as_ref()
1241     }
1242 }
1243 
1244 impl AsMut<[u8]> for BytesMut {
1245     #[inline]
as_mut(&mut self) -> &mut [u8]1246     fn as_mut(&mut self) -> &mut [u8] {
1247         self.as_slice_mut()
1248     }
1249 }
1250 
1251 impl DerefMut for BytesMut {
1252     #[inline]
deref_mut(&mut self) -> &mut [u8]1253     fn deref_mut(&mut self) -> &mut [u8] {
1254         self.as_mut()
1255     }
1256 }
1257 
1258 impl<'a> From<&'a [u8]> for BytesMut {
from(src: &'a [u8]) -> BytesMut1259     fn from(src: &'a [u8]) -> BytesMut {
1260         BytesMut::from_vec(src.to_vec())
1261     }
1262 }
1263 
1264 impl<'a> From<&'a str> for BytesMut {
from(src: &'a str) -> BytesMut1265     fn from(src: &'a str) -> BytesMut {
1266         BytesMut::from(src.as_bytes())
1267     }
1268 }
1269 
1270 impl From<BytesMut> for Bytes {
from(src: BytesMut) -> Bytes1271     fn from(src: BytesMut) -> Bytes {
1272         src.freeze()
1273     }
1274 }
1275 
1276 impl PartialEq for BytesMut {
eq(&self, other: &BytesMut) -> bool1277     fn eq(&self, other: &BytesMut) -> bool {
1278         self.as_slice() == other.as_slice()
1279     }
1280 }
1281 
1282 impl PartialOrd for BytesMut {
partial_cmp(&self, other: &BytesMut) -> Option<cmp::Ordering>1283     fn partial_cmp(&self, other: &BytesMut) -> Option<cmp::Ordering> {
1284         self.as_slice().partial_cmp(other.as_slice())
1285     }
1286 }
1287 
1288 impl Ord for BytesMut {
cmp(&self, other: &BytesMut) -> cmp::Ordering1289     fn cmp(&self, other: &BytesMut) -> cmp::Ordering {
1290         self.as_slice().cmp(other.as_slice())
1291     }
1292 }
1293 
1294 impl Eq for BytesMut {}
1295 
1296 impl Default for BytesMut {
1297     #[inline]
default() -> BytesMut1298     fn default() -> BytesMut {
1299         BytesMut::new()
1300     }
1301 }
1302 
1303 impl hash::Hash for BytesMut {
hash<H>(&self, state: &mut H) where H: hash::Hasher,1304     fn hash<H>(&self, state: &mut H)
1305     where
1306         H: hash::Hasher,
1307     {
1308         let s: &[u8] = self.as_ref();
1309         s.hash(state);
1310     }
1311 }
1312 
1313 impl Borrow<[u8]> for BytesMut {
borrow(&self) -> &[u8]1314     fn borrow(&self) -> &[u8] {
1315         self.as_ref()
1316     }
1317 }
1318 
1319 impl BorrowMut<[u8]> for BytesMut {
borrow_mut(&mut self) -> &mut [u8]1320     fn borrow_mut(&mut self) -> &mut [u8] {
1321         self.as_mut()
1322     }
1323 }
1324 
1325 impl fmt::Write for BytesMut {
1326     #[inline]
write_str(&mut self, s: &str) -> fmt::Result1327     fn write_str(&mut self, s: &str) -> fmt::Result {
1328         if self.remaining_mut() >= s.len() {
1329             self.put_slice(s.as_bytes());
1330             Ok(())
1331         } else {
1332             Err(fmt::Error)
1333         }
1334     }
1335 
1336     #[inline]
write_fmt(&mut self, args: fmt::Arguments<'_>) -> fmt::Result1337     fn write_fmt(&mut self, args: fmt::Arguments<'_>) -> fmt::Result {
1338         fmt::write(self, args)
1339     }
1340 }
1341 
1342 impl Clone for BytesMut {
clone(&self) -> BytesMut1343     fn clone(&self) -> BytesMut {
1344         BytesMut::from(&self[..])
1345     }
1346 }
1347 
1348 impl IntoIterator for BytesMut {
1349     type Item = u8;
1350     type IntoIter = IntoIter<BytesMut>;
1351 
into_iter(self) -> Self::IntoIter1352     fn into_iter(self) -> Self::IntoIter {
1353         IntoIter::new(self)
1354     }
1355 }
1356 
1357 impl<'a> IntoIterator for &'a BytesMut {
1358     type Item = &'a u8;
1359     type IntoIter = core::slice::Iter<'a, u8>;
1360 
into_iter(self) -> Self::IntoIter1361     fn into_iter(self) -> Self::IntoIter {
1362         self.as_ref().iter()
1363     }
1364 }
1365 
1366 impl Extend<u8> for BytesMut {
extend<T>(&mut self, iter: T) where T: IntoIterator<Item = u8>,1367     fn extend<T>(&mut self, iter: T)
1368     where
1369         T: IntoIterator<Item = u8>,
1370     {
1371         let iter = iter.into_iter();
1372 
1373         let (lower, _) = iter.size_hint();
1374         self.reserve(lower);
1375 
1376         // TODO: optimize
1377         // 1. If self.kind() == KIND_VEC, use Vec::extend
1378         for b in iter {
1379             self.put_u8(b);
1380         }
1381     }
1382 }
1383 
1384 impl<'a> Extend<&'a u8> for BytesMut {
extend<T>(&mut self, iter: T) where T: IntoIterator<Item = &'a u8>,1385     fn extend<T>(&mut self, iter: T)
1386     where
1387         T: IntoIterator<Item = &'a u8>,
1388     {
1389         self.extend(iter.into_iter().copied())
1390     }
1391 }
1392 
1393 impl Extend<Bytes> for BytesMut {
extend<T>(&mut self, iter: T) where T: IntoIterator<Item = Bytes>,1394     fn extend<T>(&mut self, iter: T)
1395     where
1396         T: IntoIterator<Item = Bytes>,
1397     {
1398         for bytes in iter {
1399             self.extend_from_slice(&bytes)
1400         }
1401     }
1402 }
1403 
1404 impl FromIterator<u8> for BytesMut {
from_iter<T: IntoIterator<Item = u8>>(into_iter: T) -> Self1405     fn from_iter<T: IntoIterator<Item = u8>>(into_iter: T) -> Self {
1406         BytesMut::from_vec(Vec::from_iter(into_iter))
1407     }
1408 }
1409 
1410 impl<'a> FromIterator<&'a u8> for BytesMut {
from_iter<T: IntoIterator<Item = &'a u8>>(into_iter: T) -> Self1411     fn from_iter<T: IntoIterator<Item = &'a u8>>(into_iter: T) -> Self {
1412         BytesMut::from_iter(into_iter.into_iter().copied())
1413     }
1414 }
1415 
1416 /*
1417  *
1418  * ===== Inner =====
1419  *
1420  */
1421 
increment_shared(ptr: *mut Shared)1422 unsafe fn increment_shared(ptr: *mut Shared) {
1423     let old_size = (*ptr).ref_count.fetch_add(1, Ordering::Relaxed);
1424 
1425     if old_size > isize::MAX as usize {
1426         crate::abort();
1427     }
1428 }
1429 
release_shared(ptr: *mut Shared)1430 unsafe fn release_shared(ptr: *mut Shared) {
1431     // `Shared` storage... follow the drop steps from Arc.
1432     if (*ptr).ref_count.fetch_sub(1, Ordering::Release) != 1 {
1433         return;
1434     }
1435 
1436     // This fence is needed to prevent reordering of use of the data and
1437     // deletion of the data.  Because it is marked `Release`, the decreasing
1438     // of the reference count synchronizes with this `Acquire` fence. This
1439     // means that use of the data happens before decreasing the reference
1440     // count, which happens before this fence, which happens before the
1441     // deletion of the data.
1442     //
1443     // As explained in the [Boost documentation][1],
1444     //
1445     // > It is important to enforce any possible access to the object in one
1446     // > thread (through an existing reference) to *happen before* deleting
1447     // > the object in a different thread. This is achieved by a "release"
1448     // > operation after dropping a reference (any access to the object
1449     // > through this reference must obviously happened before), and an
1450     // > "acquire" operation before deleting the object.
1451     //
1452     // [1]: (www.boost.org/doc/libs/1_55_0/doc/html/atomic/usage_examples.html)
1453     //
1454     // Thread sanitizer does not support atomic fences. Use an atomic load
1455     // instead.
1456     (*ptr).ref_count.load(Ordering::Acquire);
1457 
1458     // Drop the data
1459     drop(Box::from_raw(ptr));
1460 }
1461 
1462 impl Shared {
is_unique(&self) -> bool1463     fn is_unique(&self) -> bool {
1464         // The goal is to check if the current handle is the only handle
1465         // that currently has access to the buffer. This is done by
1466         // checking if the `ref_count` is currently 1.
1467         //
1468         // The `Acquire` ordering synchronizes with the `Release` as
1469         // part of the `fetch_sub` in `release_shared`. The `fetch_sub`
1470         // operation guarantees that any mutations done in other threads
1471         // are ordered before the `ref_count` is decremented. As such,
1472         // this `Acquire` will guarantee that those mutations are
1473         // visible to the current thread.
1474         self.ref_count.load(Ordering::Acquire) == 1
1475     }
1476 }
1477 
1478 #[inline]
original_capacity_to_repr(cap: usize) -> usize1479 fn original_capacity_to_repr(cap: usize) -> usize {
1480     let width = PTR_WIDTH - ((cap >> MIN_ORIGINAL_CAPACITY_WIDTH).leading_zeros() as usize);
1481     cmp::min(
1482         width,
1483         MAX_ORIGINAL_CAPACITY_WIDTH - MIN_ORIGINAL_CAPACITY_WIDTH,
1484     )
1485 }
1486 
original_capacity_from_repr(repr: usize) -> usize1487 fn original_capacity_from_repr(repr: usize) -> usize {
1488     if repr == 0 {
1489         return 0;
1490     }
1491 
1492     1 << (repr + (MIN_ORIGINAL_CAPACITY_WIDTH - 1))
1493 }
1494 
1495 #[cfg(test)]
1496 mod tests {
1497     use super::*;
1498 
1499     #[test]
test_original_capacity_to_repr()1500     fn test_original_capacity_to_repr() {
1501         assert_eq!(original_capacity_to_repr(0), 0);
1502 
1503         let max_width = 32;
1504 
1505         for width in 1..(max_width + 1) {
1506             let cap = 1 << width - 1;
1507 
1508             let expected = if width < MIN_ORIGINAL_CAPACITY_WIDTH {
1509                 0
1510             } else if width < MAX_ORIGINAL_CAPACITY_WIDTH {
1511                 width - MIN_ORIGINAL_CAPACITY_WIDTH
1512             } else {
1513                 MAX_ORIGINAL_CAPACITY_WIDTH - MIN_ORIGINAL_CAPACITY_WIDTH
1514             };
1515 
1516             assert_eq!(original_capacity_to_repr(cap), expected);
1517 
1518             if width > 1 {
1519                 assert_eq!(original_capacity_to_repr(cap + 1), expected);
1520             }
1521 
1522             //  MIN_ORIGINAL_CAPACITY_WIDTH must be bigger than 7 to pass tests below
1523             if width == MIN_ORIGINAL_CAPACITY_WIDTH + 1 {
1524                 assert_eq!(original_capacity_to_repr(cap - 24), expected - 1);
1525                 assert_eq!(original_capacity_to_repr(cap + 76), expected);
1526             } else if width == MIN_ORIGINAL_CAPACITY_WIDTH + 2 {
1527                 assert_eq!(original_capacity_to_repr(cap - 1), expected - 1);
1528                 assert_eq!(original_capacity_to_repr(cap - 48), expected - 1);
1529             }
1530         }
1531     }
1532 
1533     #[test]
test_original_capacity_from_repr()1534     fn test_original_capacity_from_repr() {
1535         assert_eq!(0, original_capacity_from_repr(0));
1536 
1537         let min_cap = 1 << MIN_ORIGINAL_CAPACITY_WIDTH;
1538 
1539         assert_eq!(min_cap, original_capacity_from_repr(1));
1540         assert_eq!(min_cap * 2, original_capacity_from_repr(2));
1541         assert_eq!(min_cap * 4, original_capacity_from_repr(3));
1542         assert_eq!(min_cap * 8, original_capacity_from_repr(4));
1543         assert_eq!(min_cap * 16, original_capacity_from_repr(5));
1544         assert_eq!(min_cap * 32, original_capacity_from_repr(6));
1545         assert_eq!(min_cap * 64, original_capacity_from_repr(7));
1546     }
1547 }
1548 
1549 unsafe impl Send for BytesMut {}
1550 unsafe impl Sync for BytesMut {}
1551 
1552 /*
1553  *
1554  * ===== PartialEq / PartialOrd =====
1555  *
1556  */
1557 
1558 impl PartialEq<[u8]> for BytesMut {
eq(&self, other: &[u8]) -> bool1559     fn eq(&self, other: &[u8]) -> bool {
1560         &**self == other
1561     }
1562 }
1563 
1564 impl PartialOrd<[u8]> for BytesMut {
partial_cmp(&self, other: &[u8]) -> Option<cmp::Ordering>1565     fn partial_cmp(&self, other: &[u8]) -> Option<cmp::Ordering> {
1566         (**self).partial_cmp(other)
1567     }
1568 }
1569 
1570 impl PartialEq<BytesMut> for [u8] {
eq(&self, other: &BytesMut) -> bool1571     fn eq(&self, other: &BytesMut) -> bool {
1572         *other == *self
1573     }
1574 }
1575 
1576 impl PartialOrd<BytesMut> for [u8] {
partial_cmp(&self, other: &BytesMut) -> Option<cmp::Ordering>1577     fn partial_cmp(&self, other: &BytesMut) -> Option<cmp::Ordering> {
1578         <[u8] as PartialOrd<[u8]>>::partial_cmp(self, other)
1579     }
1580 }
1581 
1582 impl PartialEq<str> for BytesMut {
eq(&self, other: &str) -> bool1583     fn eq(&self, other: &str) -> bool {
1584         &**self == other.as_bytes()
1585     }
1586 }
1587 
1588 impl PartialOrd<str> for BytesMut {
partial_cmp(&self, other: &str) -> Option<cmp::Ordering>1589     fn partial_cmp(&self, other: &str) -> Option<cmp::Ordering> {
1590         (**self).partial_cmp(other.as_bytes())
1591     }
1592 }
1593 
1594 impl PartialEq<BytesMut> for str {
eq(&self, other: &BytesMut) -> bool1595     fn eq(&self, other: &BytesMut) -> bool {
1596         *other == *self
1597     }
1598 }
1599 
1600 impl PartialOrd<BytesMut> for str {
partial_cmp(&self, other: &BytesMut) -> Option<cmp::Ordering>1601     fn partial_cmp(&self, other: &BytesMut) -> Option<cmp::Ordering> {
1602         <[u8] as PartialOrd<[u8]>>::partial_cmp(self.as_bytes(), other)
1603     }
1604 }
1605 
1606 impl PartialEq<Vec<u8>> for BytesMut {
eq(&self, other: &Vec<u8>) -> bool1607     fn eq(&self, other: &Vec<u8>) -> bool {
1608         *self == other[..]
1609     }
1610 }
1611 
1612 impl PartialOrd<Vec<u8>> for BytesMut {
partial_cmp(&self, other: &Vec<u8>) -> Option<cmp::Ordering>1613     fn partial_cmp(&self, other: &Vec<u8>) -> Option<cmp::Ordering> {
1614         (**self).partial_cmp(&other[..])
1615     }
1616 }
1617 
1618 impl PartialEq<BytesMut> for Vec<u8> {
eq(&self, other: &BytesMut) -> bool1619     fn eq(&self, other: &BytesMut) -> bool {
1620         *other == *self
1621     }
1622 }
1623 
1624 impl PartialOrd<BytesMut> for Vec<u8> {
partial_cmp(&self, other: &BytesMut) -> Option<cmp::Ordering>1625     fn partial_cmp(&self, other: &BytesMut) -> Option<cmp::Ordering> {
1626         other.partial_cmp(self)
1627     }
1628 }
1629 
1630 impl PartialEq<String> for BytesMut {
eq(&self, other: &String) -> bool1631     fn eq(&self, other: &String) -> bool {
1632         *self == other[..]
1633     }
1634 }
1635 
1636 impl PartialOrd<String> for BytesMut {
partial_cmp(&self, other: &String) -> Option<cmp::Ordering>1637     fn partial_cmp(&self, other: &String) -> Option<cmp::Ordering> {
1638         (**self).partial_cmp(other.as_bytes())
1639     }
1640 }
1641 
1642 impl PartialEq<BytesMut> for String {
eq(&self, other: &BytesMut) -> bool1643     fn eq(&self, other: &BytesMut) -> bool {
1644         *other == *self
1645     }
1646 }
1647 
1648 impl PartialOrd<BytesMut> for String {
partial_cmp(&self, other: &BytesMut) -> Option<cmp::Ordering>1649     fn partial_cmp(&self, other: &BytesMut) -> Option<cmp::Ordering> {
1650         <[u8] as PartialOrd<[u8]>>::partial_cmp(self.as_bytes(), other)
1651     }
1652 }
1653 
1654 impl<'a, T: ?Sized> PartialEq<&'a T> for BytesMut
1655 where
1656     BytesMut: PartialEq<T>,
1657 {
eq(&self, other: &&'a T) -> bool1658     fn eq(&self, other: &&'a T) -> bool {
1659         *self == **other
1660     }
1661 }
1662 
1663 impl<'a, T: ?Sized> PartialOrd<&'a T> for BytesMut
1664 where
1665     BytesMut: PartialOrd<T>,
1666 {
partial_cmp(&self, other: &&'a T) -> Option<cmp::Ordering>1667     fn partial_cmp(&self, other: &&'a T) -> Option<cmp::Ordering> {
1668         self.partial_cmp(*other)
1669     }
1670 }
1671 
1672 impl PartialEq<BytesMut> for &[u8] {
eq(&self, other: &BytesMut) -> bool1673     fn eq(&self, other: &BytesMut) -> bool {
1674         *other == *self
1675     }
1676 }
1677 
1678 impl PartialOrd<BytesMut> for &[u8] {
partial_cmp(&self, other: &BytesMut) -> Option<cmp::Ordering>1679     fn partial_cmp(&self, other: &BytesMut) -> Option<cmp::Ordering> {
1680         <[u8] as PartialOrd<[u8]>>::partial_cmp(self, other)
1681     }
1682 }
1683 
1684 impl PartialEq<BytesMut> for &str {
eq(&self, other: &BytesMut) -> bool1685     fn eq(&self, other: &BytesMut) -> bool {
1686         *other == *self
1687     }
1688 }
1689 
1690 impl PartialOrd<BytesMut> for &str {
partial_cmp(&self, other: &BytesMut) -> Option<cmp::Ordering>1691     fn partial_cmp(&self, other: &BytesMut) -> Option<cmp::Ordering> {
1692         other.partial_cmp(self)
1693     }
1694 }
1695 
1696 impl PartialEq<BytesMut> for Bytes {
eq(&self, other: &BytesMut) -> bool1697     fn eq(&self, other: &BytesMut) -> bool {
1698         other[..] == self[..]
1699     }
1700 }
1701 
1702 impl PartialEq<Bytes> for BytesMut {
eq(&self, other: &Bytes) -> bool1703     fn eq(&self, other: &Bytes) -> bool {
1704         other[..] == self[..]
1705     }
1706 }
1707 
1708 impl From<BytesMut> for Vec<u8> {
from(bytes: BytesMut) -> Self1709     fn from(bytes: BytesMut) -> Self {
1710         let kind = bytes.kind();
1711         let bytes = ManuallyDrop::new(bytes);
1712 
1713         let mut vec = if kind == KIND_VEC {
1714             unsafe {
1715                 let off = bytes.get_vec_pos();
1716                 rebuild_vec(bytes.ptr.as_ptr(), bytes.len, bytes.cap, off)
1717             }
1718         } else {
1719             let shared = bytes.data as *mut Shared;
1720 
1721             if unsafe { (*shared).is_unique() } {
1722                 let vec = mem::replace(unsafe { &mut (*shared).vec }, Vec::new());
1723 
1724                 unsafe { release_shared(shared) };
1725 
1726                 vec
1727             } else {
1728                 return ManuallyDrop::into_inner(bytes).deref().to_vec();
1729             }
1730         };
1731 
1732         let len = bytes.len;
1733 
1734         unsafe {
1735             ptr::copy(bytes.ptr.as_ptr(), vec.as_mut_ptr(), len);
1736             vec.set_len(len);
1737         }
1738 
1739         vec
1740     }
1741 }
1742 
1743 #[inline]
vptr(ptr: *mut u8) -> NonNull<u8>1744 fn vptr(ptr: *mut u8) -> NonNull<u8> {
1745     if cfg!(debug_assertions) {
1746         NonNull::new(ptr).expect("Vec pointer should be non-null")
1747     } else {
1748         unsafe { NonNull::new_unchecked(ptr) }
1749     }
1750 }
1751 
1752 /// Returns a dangling pointer with the given address. This is used to store
1753 /// integer data in pointer fields.
1754 ///
1755 /// It is equivalent to `addr as *mut T`, but this fails on miri when strict
1756 /// provenance checking is enabled.
1757 #[inline]
invalid_ptr<T>(addr: usize) -> *mut T1758 fn invalid_ptr<T>(addr: usize) -> *mut T {
1759     let ptr = core::ptr::null_mut::<u8>().wrapping_add(addr);
1760     debug_assert_eq!(ptr as usize, addr);
1761     ptr.cast::<T>()
1762 }
1763 
rebuild_vec(ptr: *mut u8, mut len: usize, mut cap: usize, off: usize) -> Vec<u8>1764 unsafe fn rebuild_vec(ptr: *mut u8, mut len: usize, mut cap: usize, off: usize) -> Vec<u8> {
1765     let ptr = ptr.sub(off);
1766     len += off;
1767     cap += off;
1768 
1769     Vec::from_raw_parts(ptr, len, cap)
1770 }
1771 
1772 // ===== impl SharedVtable =====
1773 
1774 static SHARED_VTABLE: Vtable = Vtable {
1775     clone: shared_v_clone,
1776     to_vec: shared_v_to_vec,
1777     to_mut: shared_v_to_mut,
1778     is_unique: shared_v_is_unique,
1779     drop: shared_v_drop,
1780 };
1781 
shared_v_clone(data: &AtomicPtr<()>, ptr: *const u8, len: usize) -> Bytes1782 unsafe fn shared_v_clone(data: &AtomicPtr<()>, ptr: *const u8, len: usize) -> Bytes {
1783     let shared = data.load(Ordering::Relaxed) as *mut Shared;
1784     increment_shared(shared);
1785 
1786     let data = AtomicPtr::new(shared as *mut ());
1787     Bytes::with_vtable(ptr, len, data, &SHARED_VTABLE)
1788 }
1789 
shared_v_to_vec(data: &AtomicPtr<()>, ptr: *const u8, len: usize) -> Vec<u8>1790 unsafe fn shared_v_to_vec(data: &AtomicPtr<()>, ptr: *const u8, len: usize) -> Vec<u8> {
1791     let shared: *mut Shared = data.load(Ordering::Relaxed).cast();
1792 
1793     if (*shared).is_unique() {
1794         let shared = &mut *shared;
1795 
1796         // Drop shared
1797         let mut vec = mem::replace(&mut shared.vec, Vec::new());
1798         release_shared(shared);
1799 
1800         // Copy back buffer
1801         ptr::copy(ptr, vec.as_mut_ptr(), len);
1802         vec.set_len(len);
1803 
1804         vec
1805     } else {
1806         let v = slice::from_raw_parts(ptr, len).to_vec();
1807         release_shared(shared);
1808         v
1809     }
1810 }
1811 
shared_v_to_mut(data: &AtomicPtr<()>, ptr: *const u8, len: usize) -> BytesMut1812 unsafe fn shared_v_to_mut(data: &AtomicPtr<()>, ptr: *const u8, len: usize) -> BytesMut {
1813     let shared: *mut Shared = data.load(Ordering::Relaxed).cast();
1814 
1815     if (*shared).is_unique() {
1816         let shared = &mut *shared;
1817 
1818         // The capacity is always the original capacity of the buffer
1819         // minus the offset from the start of the buffer
1820         let v = &mut shared.vec;
1821         let v_capacity = v.capacity();
1822         let v_ptr = v.as_mut_ptr();
1823         let offset = offset_from(ptr as *mut u8, v_ptr);
1824         let cap = v_capacity - offset;
1825 
1826         let ptr = vptr(ptr as *mut u8);
1827 
1828         BytesMut {
1829             ptr,
1830             len,
1831             cap,
1832             data: shared,
1833         }
1834     } else {
1835         let v = slice::from_raw_parts(ptr, len).to_vec();
1836         release_shared(shared);
1837         BytesMut::from_vec(v)
1838     }
1839 }
1840 
shared_v_is_unique(data: &AtomicPtr<()>) -> bool1841 unsafe fn shared_v_is_unique(data: &AtomicPtr<()>) -> bool {
1842     let shared = data.load(Ordering::Acquire);
1843     let ref_count = (*shared.cast::<Shared>()).ref_count.load(Ordering::Relaxed);
1844     ref_count == 1
1845 }
1846 
shared_v_drop(data: &mut AtomicPtr<()>, _ptr: *const u8, _len: usize)1847 unsafe fn shared_v_drop(data: &mut AtomicPtr<()>, _ptr: *const u8, _len: usize) {
1848     data.with_mut(|shared| {
1849         release_shared(*shared as *mut Shared);
1850     });
1851 }
1852 
1853 // compile-fails
1854 
1855 /// ```compile_fail
1856 /// use bytes::BytesMut;
1857 /// #[deny(unused_must_use)]
1858 /// {
1859 ///     let mut b1 = BytesMut::from("hello world");
1860 ///     b1.split_to(6);
1861 /// }
1862 /// ```
_split_to_must_use()1863 fn _split_to_must_use() {}
1864 
1865 /// ```compile_fail
1866 /// use bytes::BytesMut;
1867 /// #[deny(unused_must_use)]
1868 /// {
1869 ///     let mut b1 = BytesMut::from("hello world");
1870 ///     b1.split_off(6);
1871 /// }
1872 /// ```
_split_off_must_use()1873 fn _split_off_must_use() {}
1874 
1875 /// ```compile_fail
1876 /// use bytes::BytesMut;
1877 /// #[deny(unused_must_use)]
1878 /// {
1879 ///     let mut b1 = BytesMut::from("hello world");
1880 ///     b1.split();
1881 /// }
1882 /// ```
_split_must_use()1883 fn _split_must_use() {}
1884 
1885 // fuzz tests
1886 #[cfg(all(test, loom))]
1887 mod fuzz {
1888     use loom::sync::Arc;
1889     use loom::thread;
1890 
1891     use super::BytesMut;
1892     use crate::Bytes;
1893 
1894     #[test]
bytes_mut_cloning_frozen()1895     fn bytes_mut_cloning_frozen() {
1896         loom::model(|| {
1897             let a = BytesMut::from(&b"abcdefgh"[..]).split().freeze();
1898             let addr = a.as_ptr() as usize;
1899 
1900             // test the Bytes::clone is Sync by putting it in an Arc
1901             let a1 = Arc::new(a);
1902             let a2 = a1.clone();
1903 
1904             let t1 = thread::spawn(move || {
1905                 let b: Bytes = (*a1).clone();
1906                 assert_eq!(b.as_ptr() as usize, addr);
1907             });
1908 
1909             let t2 = thread::spawn(move || {
1910                 let b: Bytes = (*a2).clone();
1911                 assert_eq!(b.as_ptr() as usize, addr);
1912             });
1913 
1914             t1.join().unwrap();
1915             t2.join().unwrap();
1916         });
1917     }
1918 }
1919