dope-net 0.46.0

use std::{mem, ops::Deref};

use dope_core::driver::{
    receive::credit,
    route::{
        self, Identity, KeyTag,
        indices::SlotIndex,
        table::{self, Slab, borrows, conn, entries::occupied},
    },
};
use o3::{
    alloc::BoxSliceExt,
    depot::arena::slab::external::errors,
    mem::{raw, slot},
};

use crate::{
    link::{
        pool::{
            Key,
            recv::{
                self, input,
                storage::{self, Seal},
            },
        },
        slot::types,
    },
    wire,
    wire::receive::{Chunk, mode::Disabled},
};

/// Proof that one resident connection's selected tail state has been cleared.
#[repr(transparent)]
pub struct Cleared<'a, T, Tag: route::Tag> {
    entry: occupied::Entry<'a, T, Tag>,
}

impl<T, Tag: route::Tag> Cleared<'_, T, Tag> {
    pub(in crate::link::pool) fn remove(self) -> T {
        self.entry.remove()
    }

    pub(in crate::link::pool) fn try_remove<R, E>(
        self,
        take: impl FnOnce(T) -> Result<R, (T, E)>,
    ) -> Result<R, (Self, E)> {
        self.entry
            .try_remove(take)
            .map_err(|(entry, error)| (Self { entry }, error))
    }
}

impl<'a, T, Tag: route::Tag> Deref for Cleared<'a, T, Tag> {
    type Target = occupied::Entry<'a, T, Tag>;

    fn deref(&self) -> &Self::Target {
        &self.entry
    }
}

impl<'a, T, Tag: route::Tag> Cleared<'a, T, Tag> {
    pub(in crate::link::pool) fn get_mut(&mut self) -> &mut T {
        self.entry.get_mut()
    }
}

/// Connection slab and its statically selected tail storage.
pub struct Table<T, Tag: route::Tag, S, R = ()> {
    slab: Slab<T, Tag>,
    tail: S,
    storage: R,
}

impl<T, Tag: route::Tag, S, R> Table<T, Tag, S, R> {
    pub(in crate::link::pool) fn access(&mut self) -> borrows::Access<'_, T, Tag> {
        self.slab.access()
    }
}

impl<T, Tag: route::Tag, S, R> Deref for Table<T, Tag, S, R> {
    type Target = Slab<T, Tag>;

    fn deref(&self) -> &Self::Target {
        &self.slab
    }
}

impl<T, ID: Identity, S> Table<T, ID, S> {
    pub(in crate::link::pool) fn try_storage<'d, P: storage::raw::Mode, R>(
        self,
        storage: recv::State<'d, ID, P, R>,
    ) -> crate::Result<Stored<'d, T, ID, S, P, R>> {
        let entries = self.slab.capacity().count().get();
        if let storage::Extent::Lanes(lanes) = P::extent::<R>(&storage.storage)
            && lanes != entries
        {
            return Err(crate::Resource::Config(
                crate::ConfigError::ReceiveCapacity { entries, lanes },
            ));
        }
        let lanes = storage.rearm.capacity();
        if lanes != entries {
            return Err(crate::Resource::Config(crate::ConfigError::RearmCapacity {
                entries,
                lanes,
            }));
        }
        Ok(Table {
            slab: self.slab,
            tail: self.tail,
            storage,
        })
    }
}

impl<'w, 'd: 'w, ID: Identity + 'w, W: wire::Spec + 'w + 'd, U, S, P: storage::raw::Mode, R>
    Stored<'d, types::Connection<'w, 'd, ID, W, U>, ID, S, P, R>
{
    pub(in crate::link::pool) fn get_mut<'a>(
        &'a mut self,
        parts: table::Parts<ID>,
    ) -> Borrowed<'a, 'd, ID, P, &'a mut types::Connection<'w, 'd, ID, W, U>, R> {
        let entry = self.slab.access().get(parts)?;
        // SAFETY: construction pairs all fixed extents; this resident supplies the exact key.
        use o3::depot::raw::Proof;
        let bound = unsafe { Proof::new((&mut self.storage, entry.key)) }.into();
        Some((entry, bound))
    }

    pub(in crate::link::pool) fn try_entry<'a>(
        &'a mut self,
        parts: table::Parts<ID>,
    ) -> Occupied<'a, 'w, 'd, ID, W, U, P, R> {
        let entry = self.slab.access().try_occupied(parts)?;
        // SAFETY: the occupied entry belongs to this table's validated fixed domain.
        use o3::depot::raw::Proof;
        let bound = unsafe { Proof::new((&mut self.storage, entry.get().key)) }.into();
        Ok((entry, bound))
    }

    pub(in crate::link::pool) fn pending(&self) -> usize {
        self.storage.rearm.len()
    }

    pub(in crate::link::pool) fn pop(&mut self) -> Option<Key<'d, ID>> {
        self.storage.rearm.pop_front()
    }
}

impl<'d, T, ID: Identity, S, P: storage::raw::Mode, R> Stored<'d, T, ID, S, P, R> {
    pub(in crate::link::pool) fn read(&self, parts: table::Parts<ID>) -> Option<(&T, bool)> {
        use o3::depot::raw::Proof;

        let entry = self.slab.get(parts)?;
        // SAFETY: try_storage paired the fixed extents; this lookup admits the exact index.
        let slot =
            unsafe { Proof::new((&self.storage.storage, parts.slot().index().into_usize())) };
        let lane = P::has_queued::<R>(slot);
        Some((entry, lane))
    }
}

impl<'d, T, ID: Identity, S, P, R> Stored<'d, T, ID, S, P, R>
where
    P: storage::raw::Mode<Storage<'d, R> = storage::Retained<'d, R>>,
{
    pub(in crate::link::pool) fn credit(&self) -> &credit::Storage<'d> {
        &self.storage.storage.credit
    }
}

/// Statically selected direct or deferred tail storage.
pub trait Store: Seal + Sized {
    fn try_new<T, Tag: route::Tag>(cap: conn::Capacity) -> crate::Result<Table<T, Tag, Self>>;

    /// Keeps the admitted resident while observing its deferred tail.
    fn pending<T, Tag: route::Tag>(
        slot: Slot<'_, T, Tag, Self>,
    ) -> (occupied::Entry<'_, T, Tag>, bool);

    fn clear<'a, 'w, 'd: 'w, ID: Identity + 'w, W: wire::Spec + 'w + 'd, T>(
        slot: Slot<'a, types::Connection<'w, 'd, ID, W, T>, ID, Self>,
    ) -> Cleared<'a, types::Connection<'w, 'd, ID, W, T>, ID>;
}

impl Store for () {
    fn try_new<T, Tag: route::Tag>(cap: conn::Capacity) -> crate::Result<Table<T, Tag, Self>> {
        Ok(Table {
            slab: Slab::try_new(cap.into()).map_err(crate::Resource::Allocation)?,
            tail: (),
            storage: (),
        })
    }

    fn pending<T, Tag: route::Tag>(
        slot: Slot<'_, T, Tag, Self>,
    ) -> (occupied::Entry<'_, T, Tag>, bool) {
        (slot.entry, false)
    }

    fn clear<'a, 'w, 'd: 'w, ID: Identity + 'w, W: wire::Spec + 'w + 'd, T>(
        slot: Slot<'a, types::Connection<'w, 'd, ID, W, T>, ID, Self>,
    ) -> Cleared<'a, types::Connection<'w, 'd, ID, W, T>, ID> {
        Cleared { entry: slot.entry }
    }
}

#[repr(transparent)]
pub struct Taken<'a, 'w, 'd: 'w, W: wire::Spec + 'w> {
    occupied: raw::Occupied<'a, Chunk<'w, 'd, W, Disabled>>,
}

/// One deferred tail per connection.
pub struct Queued<'w, 'd: 'w, W: wire::Spec + 'w> {
    chunks: Box<[slot::Storage<Chunk<'w, 'd, W, Disabled>>]>,
}

impl<'w, 'd: 'w, W: wire::Spec + 'w> Seal for Queued<'w, 'd, W> {}

impl<'w, 'd: 'w, W: wire::Spec + 'w> Store for Queued<'w, 'd, W> {
    fn try_new<T, Tag: route::Tag>(cap: conn::Capacity) -> crate::Result<Table<T, Tag, Self>> {
        Ok(Table {
            slab: Slab::try_new(cap.into()).map_err(crate::Resource::Allocation)?,
            tail: Self {
                chunks: Box::try_init(cap.entries().get(), || slot::Storage::new(()))
                    .map_err(crate::Resource::Allocation)?,
            },
            storage: (),
        })
    }

    fn pending<T, Tag: route::Tag>(
        slot: Slot<'_, T, Tag, Self>,
    ) -> (occupied::Entry<'_, T, Tag>, bool) {
        let (entry, chunk) = slot.split();
        (entry, chunk.get().is_some())
    }

    fn clear<'a, 'v, 'owner: 'v, ID: Identity + 'v, V: wire::Spec + 'v + 'owner, T>(
        slot: Slot<'a, types::Connection<'v, 'owner, ID, V, T>, ID, Self>,
    ) -> Cleared<'a, types::Connection<'v, 'owner, ID, V, T>, ID> {
        let (entry, chunk) = slot.split();
        if let Some(occupied) = chunk.occupied() {
            occupied.clear();
        }
        Cleared { entry }
    }
}

impl<'a, 'w, 'd: 'w, W: wire::Spec + 'w> Taken<'a, 'w, 'd, W> {
    pub(in crate::link::pool::recv) fn new(
        occupied: raw::Occupied<'a, Chunk<'w, 'd, W, Disabled>>,
    ) -> Self {
        Self { occupied }
    }

    pub fn commit(self) -> Chunk<'w, 'd, W, Disabled> {
        self.occupied.take().0
    }
}

const _: () = {
    assert!(
        mem::size_of::<Cleared<'static, (), KeyTag<0>>>()
            == mem::size_of::<occupied::Entry<'static, (), KeyTag<0>>>()
    );
    assert!(
        mem::align_of::<Cleared<'static, (), KeyTag<0>>>()
            == mem::align_of::<occupied::Entry<'static, (), KeyTag<0>>>()
    );
    assert!(mem::size_of::<Table<(), KeyTag<0>, ()>>() == mem::size_of::<Slab<(), KeyTag<0>>>());
    assert!(mem::align_of::<Table<(), KeyTag<0>, ()>>() == mem::align_of::<Slab<(), KeyTag<0>>>());
};

pub struct Slot<'a, T, Tag: route::Tag, S> {
    entry: occupied::Entry<'a, T, Tag>,
    tail: &'a mut S,
}

impl<'a, T, Tag: route::Tag, S> Deref for Slot<'a, T, Tag, S> {
    type Target = occupied::Entry<'a, T, Tag>;

    fn deref(&self) -> &Self::Target {
        &self.entry
    }
}

impl<T, Tag: route::Tag, S> Slot<'_, T, Tag, S> {
    pub(in crate::link::pool) fn get_mut(&mut self) -> &mut T {
        self.entry.get_mut()
    }
}

impl<'a, T, Tag: route::Tag, S> Slot<'a, T, Tag, S> {
    pub(in crate::link::pool) fn new<R>(
        table: &'a mut Table<T, Tag, S, R>,
        parts: table::Parts<Tag>,
    ) -> Option<Self> {
        let entry = table.slab.access().occupied(parts)?;
        Some(Self {
            entry,
            tail: &mut table.tail,
        })
    }

    pub(in crate::link::pool::recv) fn queue<'w, 'd: 'w, W: wire::Spec + 'w, M>(
        self,
    ) -> Slot<'a, T, Tag, Queued<'w, 'd, W>>
    where
        M: input::Replay<W> + input::Kind<W, Tail<'w, 'd> = S>,
    {
        Slot {
            entry: self.entry,
            tail: M::queue(self.tail),
        }
    }
}

impl<'a, 'w, 'd: 'w, T, Tag: route::Tag, W: wire::Spec + 'w> Slot<'a, T, Tag, Queued<'w, 'd, W>> {
    pub(in crate::link::pool::recv) fn split(
        self,
    ) -> (
        occupied::Entry<'a, T, Tag>,
        &'a mut slot::Storage<Chunk<'w, 'd, W, Disabled>>,
    ) {
        use o3::depot::raw::Proof;

        let Self { entry, tail } = self;
        // SAFETY: Table retains both fixed allocations and never lends either
        // owner mutably. The occupied slot belongs to this original table.
        let chunk = unsafe {
            Proof::new((
                tail.chunks.as_mut(),
                entry.key().slot().index().get() as usize,
            ))
        }
        .into();
        (entry, chunk)
    }
}

impl<'a, 'w: 'a, 'd: 'w, ID: Identity + 'w, W: wire::Spec + 'w + 'd, U, S>
    Slot<'a, types::Connection<'w, 'd, ID, W, U>, ID, S>
{
    pub(in crate::link::pool) fn receive<P: storage::raw::Mode, R>(
        table: &'a mut Stored<'d, types::Connection<'w, 'd, ID, W, U>, ID, S, P, R>,
        parts: table::Parts<ID>,
    ) -> Borrowed<'a, 'd, ID, P, Self, R> {
        let entry = table.slab.access().occupied(parts)?;
        // SAFETY: Table owns matching fixed extents and this exact occupied resident.
        use o3::depot::raw::Proof;
        let bound = unsafe { Proof::new((&mut table.storage, entry.get().key)) }.into();
        Some((
            Self {
                entry,
                tail: &mut table.tail,
            },
            bound,
        ))
    }

    pub(in crate::link::pool) fn at<P: storage::raw::Mode, R>(
        table: &'a mut Stored<'d, types::Connection<'w, 'd, ID, W, U>, ID, S, P, R>,
        index: SlotIndex,
    ) -> Borrowed<'a, 'd, ID, P, Self, R> {
        let entry = table.slab.access().at(index)?;
        // SAFETY: the occupied resident comes from this same paired table.
        use o3::depot::raw::Proof;
        let bound = unsafe { Proof::new((&mut table.storage, entry.get().key)) }.into();
        Some((
            Self {
                entry,
                tail: &mut table.tail,
            },
            bound,
        ))
    }
}

type Stored<'d, T, ID, S, P, R> = Table<T, ID, S, recv::State<'d, ID, P, R>>;

type Occupied<'a, 'w, 'd, ID, W, U, P, R> = Result<
    (
        occupied::Entry<'a, types::Connection<'w, 'd, ID, W, U>, ID>,
        recv::Borrow<'a, 'd, ID, P, R>,
    ),
    errors::Lookup<route::Epoch, ID>,
>;

type Borrowed<'a, 'd, ID, P, T, R> = Option<(T, recv::Borrow<'a, 'd, ID, P, R>)>;