clone 与提升

const KIND_ARC: usize = 0b0;
const KIND_VEC: usize = 0b1;
const KIND_MASK: usize = 0b1;

// Assert that the alignment of `Shared` is divisible by 2.
// This is a necessary invariant since we depend on allocating `Shared` a
// shared object to implicitly carry the `KIND_ARC` flag in its pointer.
// This flag is set when the LSB is 0.
const _: [(); 0 - mem::align_of::<Shared>() % 2] = []; // Assert that the alignment of `Shared` is divisible by 2.

struct Shared {
    // Holds arguments to dealloc upon Drop, but otherwise doesn't use them
    buf: *mut u8,
    cap: usize,
    ref_cnt: AtomicUsize,
}

unsafe fn promotable_even_clone(data: &AtomicPtr<()>, ptr: *const u8, len: usize) -> Bytes {
    let shared = data.load(Ordering::Acquire);
    let kind = shared as usize & KIND_MASK;

    if kind == KIND_ARC {
        shallow_clone_arc(shared.cast(), ptr, len)
    } else {
        debug_assert_eq!(kind, KIND_VEC);
        let buf = ptr_map(shared.cast(), |addr| addr & !KIND_MASK);
        shallow_clone_vec(data, shared, buf, ptr, len)
    }
}

#[cold]
unsafe fn shallow_clone_vec(
    atom: &AtomicPtr<()>,
    ptr: *const (),
    buf: *mut u8,
    offset: *const u8,
    len: usize,
) -> Bytes {
    // ...
    let shared = Box::new(MaybeUninit::<Shared>::uninit());
    let shared = Shared::init_to_raw(
        shared,
        Shared {
            buf,
            cap: offset.offset_from(buf) as usize + len,
            // Initialize refcount to 2. One for this reference, and one
            // for the new clone that will be returned from
            // `shallow_clone`.
            ref_cnt: AtomicUsize::new(2),
        },
    );

    // ...
    match atom.compare_exchange(ptr as _, shared as _, Ordering::AcqRel, Ordering::Acquire) {
        Ok(actual) => {
            debug_assert!(core::ptr::eq(actual, ptr));
            // The upgrade was successful, the new handle can be
            // returned.
            Bytes {
                ptr: offset,
                len,
                data: AtomicPtr::new(shared as _),
                vtable: &SHARED_VTABLE,
            }
        }
        Err(actual) => {
            // The upgrade failed, a concurrent clone happened. Release
            // the allocation that was made in this thread, it will not
            // be needed.
            let shared = Box::from_raw(shared);
            mem::forget(*shared);

            // Buffer already promoted to shared storage, so increment ref
            // count.
            shallow_clone_arc(actual as _, offset, len)
        }
    }
}

bytes = "=1.12.1"

use std::alloc::{GlobalAlloc, Layout, System};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering::Relaxed};

// Counts every allocation while ON is set; the Shared header is [usize; 3].
pub struct Counting;

pub static ON: AtomicBool = AtomicBool::new(false);
pub static ALLOCS: AtomicUsize = AtomicUsize::new(0);
pub static DEALLOCS: AtomicUsize = AtomicUsize::new(0);
pub static SHARED_ALLOCS: AtomicUsize = AtomicUsize::new(0);
pub static SHARED_DEALLOCS: AtomicUsize = AtomicUsize::new(0);
pub static LIVE: AtomicUsize = AtomicUsize::new(0);

fn is_shared(layout: Layout) -> bool {
  layout == Layout::new::<[usize; 3]>()
}

unsafe impl GlobalAlloc for Counting {
  unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
    LIVE.fetch_add(layout.size(), Relaxed);
    if ON.load(Relaxed) {
      ALLOCS.fetch_add(1, Relaxed);
      if is_shared(layout) {
        SHARED_ALLOCS.fetch_add(1, Relaxed);
      }
    }
    unsafe { System.alloc(layout) }
  }

  unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
    // ... 同上，计入 DEALLOCS / SHARED_DEALLOCS，LIVE 减去 size
    unsafe { System.dealloc(ptr, layout) }
  }
}

#[global_allocator]
static GLOBAL: Counting = Counting;

/// Runs `f` with counting on and returns (allocs, deallocs).
pub fn count<R>(f: impl FnOnce() -> R) -> (R, usize, usize) {
  let a0 = ALLOCS.load(Relaxed);
  let d0 = DEALLOCS.load(Relaxed);
  ON.store(true, Relaxed);
  let r = f();
  ON.store(false, Relaxed);
  (r, ALLOCS.load(Relaxed) - a0, DEALLOCS.load(Relaxed) - d0)
}

/// Reads the four words of a `Bytes`. Depends on private layout: observation only.
pub fn words(b: &bytes::Bytes) -> [usize; 4] {
  assert_eq!(size_of::<bytes::Bytes>(), size_of::<[usize; 4]>());
  unsafe { std::mem::transmute_copy::<bytes::Bytes, [usize; 4]>(b) }
}

/// Index of the `data` word, found on a fresh promotable Bytes, where data == ptr | 1.
pub fn data_index() -> usize {
  let b = bytes::Bytes::from(vec![7u8; 1000]);
  let w = words(&b);
  let p = b.as_ptr() as usize;
  assert_eq!(p & 1, 0);
  w.iter().position(|&x| x == p | 1).expect("data word")
}

use bytes::Bytes;
use exp_bytes_promote::{count, data_index, words};

fn main() {
  let di = data_index();
  let b = Bytes::from(vec![7u8; 1024]); // len == cap: promotable storage
  let p = b.as_ptr() as usize;
  println!("buf = {p:#x}, data = {:#x}", words(&b)[di]);

  let (c1, a, d) = count(|| b.clone());
  println!("clone #1: allocs={a} deallocs={d}, b.data = {:#x}, c1.data = {:#x}", words(&b)[di], words(&c1)[di]);
  let (c2, a, d) = count(|| b.clone());
  println!("clone #2: allocs={a} deallocs={d}, c2.data = {:#x}", words(&c2)[di]);
  let (c3, a, d) = count(|| c1.clone());
  println!("clone #3 (of c1): allocs={a} deallocs={d}, c3.data = {:#x}", words(&c3)[di]);
  println!("same ptr: {}", c1.as_ptr() == b.as_ptr() && c3.as_ptr() == b.as_ptr());
  let wb = words(&b);
  let vi = (0..4).find(|&i| i != di && wb[i] != p && wb[i] != b.len()).unwrap();
  println!("vtable: b == c1? {}, c1 == c2? {}", wb[vi] == words(&c1)[vi], words(&c1)[vi] == words(&c2)[vi]);

  let (_, a, d) = count(|| drop(c1));
  println!("drop c1: allocs={a} deallocs={d}");
  let (_, a, d) = count(|| drop(c2));
  println!("drop c2: allocs={a} deallocs={d}");
  let (_, a, d) = count(|| drop(c3));
  println!("drop c3: allocs={a} deallocs={d}, b.is_unique() = {}", b.is_unique());
  let (_, a, d) = count(|| drop(b));
  println!("drop b : allocs={a} deallocs={d}");
}

buf = 0x10521d8d0, data = 0x10521d8d1
clone #1: allocs=1 deallocs=0, b.data = 0x10521d0f0, c1.data = 0x10521d0f0
clone #2: allocs=0 deallocs=0, c2.data = 0x10521d0f0
clone #3 (of c1): allocs=0 deallocs=0, c3.data = 0x10521d0f0
same ptr: true
vtable: b == c1? false, c1 == c2? true
drop c1: allocs=0 deallocs=0
drop c2: allocs=0 deallocs=0
drop c3: allocs=0 deallocs=0, b.is_unique() = true
drop b : allocs=0 deallocs=2

unsafe fn shallow_clone_arc(shared: *mut Shared, ptr: *const u8, len: usize) -> Bytes {
    let old_size = (*shared).ref_cnt.fetch_add(1, Ordering::Relaxed);

    if old_size > usize::MAX >> 1 {
        crate::abort();
    }

    Bytes {
        ptr,
        len,
        data: AtomicPtr::new(shared as _),
        vtable: &SHARED_VTABLE,
    }
}

unsafe fn release_shared(ptr: *mut Shared) {
    // `Shared` storage... follow the drop steps from Arc.
    if (*ptr).ref_cnt.fetch_sub(1, Ordering::Release) != 1 {
        return;
    }

    // ...
    // Thread sanitizer does not support atomic fences. Use an atomic load
    // instead.
    (*ptr).ref_cnt.load(Ordering::Acquire);

    // Drop the data
    drop(Box::from_raw(ptr));
}

unsafe fn promotable_is_unique(data: &AtomicPtr<()>) -> bool {
    let shared = data.load(Ordering::Acquire);
    let kind = shared as usize & KIND_MASK;

    if kind == KIND_ARC {
        let ref_cnt = (*shared.cast::<Shared>()).ref_cnt.load(Ordering::Relaxed);
        ref_cnt == 1
    } else {
        true
    }
}

    pub fn try_into_mut(self) -> Result<BytesMut, Bytes> {
        if self.is_unique() {
            Ok(self.into())
        } else {
            Err(self)
        }
    }

    if (*shared).ref_cnt.load(Ordering::Acquire) == 1 {
        // Deallocate the `Shared` instance without running its destructor.
        let shared = *Box::from_raw(shared);
        let shared = ManuallyDrop::new(shared);
        let buf = shared.buf;
        let cap = shared.cap;

        // Rebuild Vec
        let off = ptr.offset_from(buf) as usize;
        let v = Vec::from_raw_parts(buf, len + off, cap);

        let mut b = BytesMut::from_vec(v);
        b.advance_unchecked(off);
        b
    } else {
        // Copy the data from Shared in a new Vec, then release it
        let v = slice::from_raw_parts(ptr, len).to_vec();
        release_shared(shared);
        BytesMut::from_vec(v)
    }

use bytes::{Buf, Bytes};
use exp_bytes_promote::count;

fn report(name: &str, b: Bytes) {
  let ptr = b.as_ptr();
  let unique = b.is_unique();
  let (r, a, d) = count(|| b.try_into_mut());
  match r {
    Ok(m) => println!(
      "{name:<28} is_unique={unique:<5} Ok  same_ptr={:<5} len={} cap={} allocs={a} deallocs={d}",
      m.as_ptr() == ptr, m.len(), m.capacity()
    ),
    Err(b) => println!("{name:<28} is_unique={unique:<5} Err same_ptr={:<5} allocs={a} deallocs={d}", b.as_ptr() == ptr),
  }
}

fn main() {
  report("promotable, never cloned", Bytes::from(vec![1u8; 1024]));

  let mut b = Bytes::from(vec![1u8; 1024]);
  b.advance(1000);
  report("promotable, advance(1000)", b);

  let b = Bytes::from(vec![1u8; 1024]);
  drop(b.clone());
  report("promoted, clone dropped", b);

  let b = Bytes::from(vec![1u8; 1024]);
  let keep = b.clone();
  report("promoted, clone alive", b);
  drop(keep);

  let b = Bytes::from(vec![1u8; 1024]);
  let s = b.slice(10..20);
  drop(b);
  report("slice(10..20), only handle", s);

  let mut v = Vec::with_capacity(2048);
  v.resize(1024, 1u8);
  report("Vec len != cap (Shared)", Bytes::from(v));

  report("from_static", Bytes::from_static(&[1u8; 1024]));
  report("from_owner(Vec)", Bytes::from_owner(vec![1u8; 1024]));
}

promotable, never cloned     is_unique=true  Ok  same_ptr=true  len=1024 cap=1024 allocs=0 deallocs=0
promotable, advance(1000)    is_unique=true  Ok  same_ptr=true  len=24 cap=24 allocs=0 deallocs=0
promoted, clone dropped      is_unique=true  Ok  same_ptr=true  len=1024 cap=1024 allocs=0 deallocs=1
promoted, clone alive        is_unique=false Err same_ptr=true  allocs=0 deallocs=0
slice(10..20), only handle   is_unique=true  Ok  same_ptr=true  len=10 cap=1014 allocs=0 deallocs=1
Vec len != cap (Shared)      is_unique=true  Ok  same_ptr=true  len=1024 cap=2048 allocs=0 deallocs=1
from_static                  is_unique=false Err same_ptr=true  allocs=0 deallocs=0
from_owner(Vec)              is_unique=false Err same_ptr=true  allocs=0 deallocs=0

threads=8 iterations=100000
  1 Shared allocated (and 0 freed by losers): 1 iterations
  2 Shared allocated (and 1 freed by losers): 8 iterations
  3 Shared allocated (and 2 freed by losers): 17 iterations
  4 Shared allocated (and 3 freed by losers): 7 iterations
  5 Shared allocated (and 4 freed by losers): 70 iterations
  6 Shared allocated (and 5 freed by losers): 332 iterations
  7 Shared allocated (and 6 freed by losers): 2272 iterations
  8 Shared allocated (and 7 freed by losers): 97293 iterations
iterations where survivors != 1: 0
iterations where a clone saw a different Shared: 0
iterations where final drop did not free exactly 1 Shared: 0