add will_truncate to Bytes vtable to prevent special casing promotable ones

Author: HyeonuPark

src/bytes.rs

@@ -110,6 +110,11 @@ pub(crate) struct Vtable {
     pub clone: unsafe fn(&AtomicPtr<()>, *const u8, usize) -> Bytes,
     /// fn(data, ptr, len)
     ///
+    /// Called before the `Bytes::truncate` is processed.
+    /// Useful if the implementation needs some preparation step for it.
+    pub will_truncate: unsafe fn(&mut AtomicPtr<()>, *const u8, usize),
+    /// fn(data, ptr, len)
+    ///
     /// Consumes `Bytes` to return `Vec<u8>`
     pub into_vec: unsafe fn(&mut AtomicPtr<()>, *const u8, usize) -> Vec<u8>,
     /// fn(data, ptr, len)
@@ -455,16 +460,10 @@ impl Bytes {
     #[inline]
     pub fn truncate(&mut self, len: usize) {
         if len < self.len {
-            // The Vec "promotable" vtables do not store the capacity,
-            // so we cannot truncate while using this repr. We *have* to
-            // promote using `split_off` so the capacity can be stored.
-            if self.vtable as *const Vtable == &PROMOTABLE_EVEN_VTABLE
-                || self.vtable as *const Vtable == &PROMOTABLE_ODD_VTABLE
-            {
-                drop(self.split_off(len));
-            } else {
-                self.len = len;
+            unsafe {
+                (self.vtable.will_truncate)(&mut self.data, self.ptr, self.len);
             }
+            self.len = len;
         }
     }
 
@@ -862,6 +861,7 @@ impl fmt::Debug for Vtable {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("Vtable")
             .field("clone", &(self.clone as *const ()))
+            .field("will_truncate", &(self.will_truncate as *const ()))
             .field("into_vec", &(self.into_vec as *const ()))
             .field("drop", &(self.drop as *const ()))
             .finish()
@@ -872,6 +872,7 @@ impl fmt::Debug for Vtable {
 
 const STATIC_VTABLE: Vtable = Vtable {
     clone: static_clone,
+    will_truncate: static_will_truncate,
     into_vec: static_into_vec,
     drop: static_drop,
 };
@@ -881,6 +882,10 @@ unsafe fn static_clone(_: &AtomicPtr<()>, ptr: *const u8, len: usize) -> Bytes {
     Bytes::from_static(slice)
 }
 
+unsafe fn static_will_truncate(_: &mut AtomicPtr<()>, _: *const u8, _: usize) {
+    // nothing to do before truncate for &'static [u8]
+}
+
 unsafe fn static_into_vec(_: &mut AtomicPtr<()>, ptr: *const u8, len: usize) -> Vec<u8> {
     let slice = slice::from_raw_parts(ptr, len);
     slice.to_vec()
@@ -894,12 +899,14 @@ unsafe fn static_drop(_: &mut AtomicPtr<()>, _: *const u8, _: usize) {
 
 static PROMOTABLE_EVEN_VTABLE: Vtable = Vtable {
     clone: promotable_even_clone,
+    will_truncate: promotable_even_will_truncate,
     into_vec: promotable_even_into_vec,
     drop: promotable_even_drop,
 };
 
 static PROMOTABLE_ODD_VTABLE: Vtable = Vtable {
     clone: promotable_odd_clone,
+    will_truncate: promotable_odd_will_truncate,
     into_vec: promotable_odd_into_vec,
     drop: promotable_odd_drop,
 };
@@ -917,6 +924,13 @@ unsafe fn promotable_even_clone(data: &AtomicPtr<()>, ptr: *const u8, len: usize
     }
 }
 
+unsafe fn promotable_even_will_truncate(data: &mut AtomicPtr<()>, ptr: *const u8, len: usize) {
+    // The Vec "promotable" vtables do not store the capacity,
+    // so we cannot truncate while using this repr. We *have* to
+    // promote using `clone` so the capacity can be stored.
+    drop(promotable_even_clone(&*data, ptr, len));
+}
+
 unsafe fn promotable_into_vec(
     data: &mut AtomicPtr<()>,
     ptr: *const u8,
@@ -980,6 +994,13 @@ unsafe fn promotable_odd_clone(data: &AtomicPtr<()>, ptr: *const u8, len: usize)
     }
 }
 
+unsafe fn promotable_odd_will_truncate(data: &mut AtomicPtr<()>, ptr: *const u8, len: usize) {
+    // The Vec "promotable" vtables do not store the capacity,
+    // so we cannot truncate while using this repr. We *have* to
+    // promote using `clone` so the capacity can be stored.
+    drop(promotable_odd_clone(&*data, ptr, len));
+}
+
 unsafe fn promotable_odd_into_vec(data: &mut AtomicPtr<()>, ptr: *const u8, len: usize) -> Vec<u8> {
     promotable_into_vec(data, ptr, len, |shared| shared.cast())
 }
@@ -1027,6 +1048,7 @@ const _: [(); 0 - mem::align_of::<Shared>() % 2] = []; // Assert that the alignm
 
 static SHARED_VTABLE: Vtable = Vtable {
     clone: shared_clone,
+    will_truncate: shared_will_truncate,
     into_vec: shared_into_vec,
     drop: shared_drop,
 };
@@ -1040,6 +1062,10 @@ unsafe fn shared_clone(data: &AtomicPtr<()>, ptr: *const u8, len: usize) -> Byte
     shallow_clone_arc(shared as _, ptr, len)
 }
 
+unsafe fn shared_will_truncate(_: &mut AtomicPtr<()>, _: *const u8, _: usize) {
+    // nothing to do before truncate for Shared
+}
+
 unsafe fn shared_into_vec_impl(shared: *mut Shared, ptr: *const u8, len: usize) -> Vec<u8> {
     // Check that the ref_cnt is 1 (unique).
     //

src/bytes_mut.rs

@@ -1675,6 +1675,7 @@ unsafe fn rebuild_vec(ptr: *mut u8, mut len: usize, mut cap: usize, off: usize)
 
 static SHARED_VTABLE: Vtable = Vtable {
     clone: shared_v_clone,
+    will_truncate: shared_v_will_truncate,
     into_vec: shared_v_into_vec,
     drop: shared_v_drop,
 };
@@ -1687,6 +1688,10 @@ unsafe fn shared_v_clone(data: &AtomicPtr<()>, ptr: *const u8, len: usize) -> By
     Bytes::with_vtable(ptr, len, data, &SHARED_VTABLE)
 }
 
+unsafe fn shared_v_will_truncate(_: &mut AtomicPtr<()>, _: *const u8, _: usize) {
+    // nothing to do before truncate for Shared
+}
+
 unsafe fn shared_v_into_vec(data: &mut AtomicPtr<()>, ptr: *const u8, len: usize) -> Vec<u8> {
     let shared: *mut Shared = (data.with_mut(|p| *p)).cast();