o3 0.33.0

use std::{marker, mem, num, ptr, slice};

use crate::{
    buffer::{resident, storage, storage::regions},
    mem::rc,
};

/// One allocation whose complete data region is initialized.
#[derive(Debug)]
#[repr(transparent)]
pub(in crate::buffer) struct Init<P: crate::buffer::raw::Policy = ()> {
    allocation: crate::buffer::raw::Heap<P>,
}

/// One complete allocation, with its plain or resident policy encoded in bit zero.
#[derive(Debug)]
#[repr(transparent)]
pub(in crate::buffer) struct Complete<'d> {
    pointer: ptr::NonNull<regions::Prefix<num::NonZero<u32>>>,
    ownership: marker::PhantomData<Init<resident::Budget<'d>>>,
}

impl<P: crate::buffer::raw::Policy> Init<P> {
    /// # Safety
    /// Every byte through `allocation.capacity()` must be initialized.
    pub(in crate::buffer) unsafe fn new(heap: crate::buffer::raw::Heap<P>) -> Self {
        Self { allocation: heap }
    }

    pub(in crate::buffer) fn capacity(&self) -> num::NonZero<u32> {
        self.allocation.capacity()
    }

    pub(in crate::buffer) fn as_slice(&self) -> &[u8] {
        // SAFETY: construction proves the complete allocation is initialized,
        // and this owner keeps the header and data live for the borrow.
        unsafe {
            slice::from_raw_parts(
                self.allocation.as_ptr().cast(),
                self.capacity().get() as usize,
            )
        }
    }

    pub(in crate::buffer) fn into_owned(self) -> storage::Owned<P> {
        storage::Owned {
            len: self.capacity().get(),
            allocation: Some(self.allocation),
        }
    }

    pub(in crate::buffer) fn into_buf(self) -> storage::Shared<crate::buffer::raw::Rc<P>> {
        let len = self.capacity();
        // SAFETY: Init transfers its fully initialized extent with its original owner.
        unsafe { self.allocation.into_shared().into_range(0, len.get()) }
    }
}

impl Init<()> {
    pub(in crate::buffer) fn bind<'d>(self) -> regions::Complete<'d> {
        Complete {
            pointer: self.allocation.into_shared().into_raw().cast(),
            ownership: marker::PhantomData,
        }
    }
}

impl<'d> Init<resident::Budget<'d>> {
    pub(in crate::buffer) fn bind(self) -> regions::Complete<'d> {
        Complete {
            pointer: self
                .allocation
                .into_shared()
                .into_raw()
                .cast()
                .map_addr(|address| address | 1),
            ownership: marker::PhantomData,
        }
    }
}

const _: () = assert!(size_of::<Init>() == size_of::<usize>());

impl<'d> Complete<'d> {
    pub(in crate::buffer) fn as_slice<'a>(&'a self) -> &'a [u8] {
        let pointer = self.pointer();
        let bytes = if self.pointer.addr().get() & 1 == 0 {
            // SAFETY: the plain Init constructor transfers its entire initialized allocation.
            unsafe { crate::buffer::raw::Header::<()>::slice(pointer.cast()) }
        } else {
            // SAFETY: bit zero encodes the transferred resident Init and its exact lifetime.
            unsafe { crate::buffer::raw::Header::<resident::Budget<'d>>::slice(pointer.cast()) }
        };
        // SAFETY: this owner borrow keeps the complete initialized region live.
        unsafe { bytes.as_ref() }
    }

    pub(in crate::buffer) fn into_buf(self) -> resident::Shared<'d> {
        let this = mem::ManuallyDrop::new(self);
        let pointer = this.pointer();
        let len = this.prefix().word().get();
        if this.pointer.addr().get() & 1 == 0 {
            // SAFETY: this Complete transfers one plain counted edge and its initialized extent.
            let bytes: storage::Shared = unsafe {
                crate::buffer::raw::Rc::<()>::from_raw(pointer.cast()).into_range(0, len)
            }
            .into();
            bytes.into()
        } else {
            // SAFETY: this Complete transfers one resident edge, its extent and policy lifetime.
            unsafe {
                crate::buffer::raw::Rc::<resident::Budget<'d>>::from_raw(pointer.cast())
                    .into_range(0, len)
            }
            .into()
        }
    }

    pub(in crate::buffer) fn bytes(&self) -> usize {
        self.prefix().word().get() as usize
    }

    fn pointer(&self) -> ptr::NonNull<regions::Prefix<num::NonZero<u32>>> {
        // SAFETY: both constructors own aligned nonzero headers; only bit zero is encoded.
        self.pointer
            .map_addr(|address| unsafe { num::NonZero::new_unchecked(address.get() & !1) })
    }

    fn prefix(&self) -> &regions::Prefix<num::NonZero<u32>> {
        // SAFETY: both complete allocation policies begin with a live Prefix.
        unsafe { self.pointer().as_ref() }
    }
}

impl Complete<'_> {
    pub(in crate::buffer) fn try_clone(&self) -> Result<Self, rc::Error> {
        use crate::mem::raw::Edge as _;

        // SAFETY: this complete owner retains the same initialized allocation during retention.
        unsafe { self.prefix().try_retain() }?;
        Ok(Self {
            pointer: self.pointer,
            ownership: marker::PhantomData,
        })
    }
}

impl<'d> Drop for Complete<'d> {
    fn drop(&mut self) {
        use crate::mem::raw::Edge as _;

        // SAFETY: Drop consumes this complete allocation's exact edge once.
        if unsafe { self.prefix().release() } != rc::Release::Last {
            return;
        }
        let pointer = self.pointer();
        if self.pointer.addr().get() & 1 == 0 {
            // SAFETY: the last plain allocation edge grants exclusive destruction.
            unsafe { crate::buffer::raw::Header::<()>::destroy(pointer.cast()) };
        } else {
            // SAFETY: the last resident allocation edge retains its exact policy lifetime.
            unsafe { crate::buffer::raw::Header::<resident::Budget<'d>>::destroy(pointer.cast()) };
        }
    }
}

const _: () = {
    assert!(align_of::<crate::buffer::raw::Header<()>>() > 1);
    assert!(align_of::<crate::buffer::raw::Header<resident::Budget<'static>>>() > 1);
    assert!(size_of::<Complete<'static>>() == size_of::<ptr::NonNull<()>>());
    assert!(size_of::<Option<Complete<'static>>>() == size_of::<ptr::NonNull<()>>());
};