Oregami
Repositories/oxedyne/fe2o3

oxedyne/fe2o3/fe2o3_text/tests/annealer_corpus/bytes_mut.rs

59.8 KiB, 1 run

created by r1870400018:11706, which is this file's identity for as long as the history lasts, whatever it is later renamed to

download · who wrote it · its history

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