o3 0.15.0

use std::{array, cell};

use crate::collections;

/// Shared surplus and every lane state in one accounting domain.
pub struct Pool<const N: usize = 1> {
    shared: cell::Cell<[usize; N]>,
    states: Box<[State<N>]>,
    _thread: crate::ThreadBound,
}

/// One lane's protected reserve and current holdings.
struct State<const N: usize> {
    reserve: [usize; N],
    held: cell::Cell<[usize; N]>,
}

/// A borrowed lane view that updates all resource dimensions atomically.
#[derive(Clone, Copy)]
pub struct Lane<'a, const N: usize = 1> {
    shared: &'a cell::Cell<[usize; N]>,
    state: &'a State<N>,
}

/// Owned fair accounting with one uniform protected reserve per lane.
pub struct Credits<const N: usize = 1> {
    used: cell::Cell<[usize; N]>,
    pool: Pool<N>,
}

impl<const N: usize> Pool<N> {
    /// Builds one accounting domain whose protected total is split exactly
    /// across `lane_count` pool-owned states.
    pub fn split(shared: [usize; N], protected: [usize; N], lane_count: usize) -> Self {
        match Self::try_split(shared, protected, lane_count) {
            Ok(pool) => pool,
            Err(error) => error.abort(),
        }
    }

    pub fn try_split(
        shared: [usize; N],
        protected: [usize; N],
        lane_count: usize,
    ) -> Result<Self, collections::AllocationError> {
        Self::try_with_reserve(shared, lane_count, |lane| {
            array::from_fn(|dimension| {
                protected[dimension] / lane_count
                    + usize::from(lane < protected[dimension] % lane_count)
            })
        })
    }

    fn try_uniform(
        shared: [usize; N],
        lane_count: usize,
        reserve: [usize; N],
    ) -> Result<Self, collections::AllocationError> {
        Self::try_with_reserve(shared, lane_count, |_| reserve)
    }

    fn try_with_reserve(
        shared: [usize; N],
        lane_count: usize,
        mut reserve: impl FnMut(usize) -> [usize; N],
    ) -> Result<Self, collections::AllocationError> {
        use crate::ThreadBound;
        assert!(N > 0, "credit dimension count must be positive");
        Ok(Self {
            shared: cell::Cell::new(shared),
            states: collections::BoxSliceExt::try_box_with(lane_count, |lane| State {
                reserve: reserve(lane),
                held: cell::Cell::new([0; N]),
            })?,
            _thread: ThreadBound::NEW,
        })
    }

    /// Borrows one lane from this pool's accounting domain.
    pub fn lane(&self, lane: usize) -> Option<Lane<'_, N>> {
        Some(Lane {
            shared: &self.shared,
            state: self.states.get(lane)?,
        })
    }

    fn shared(&self) -> [usize; N] {
        self.shared.get()
    }
}

impl<const N: usize> Lane<'_, N> {
    fn can_acquire_all(self, amount: [usize; N]) -> bool {
        let held = self.state.held.get();
        let shared = self.shared.get();
        for dimension in 0..N {
            if held[dimension].checked_add(amount[dimension]).is_none() {
                return false;
            }
            let own = self.state.reserve[dimension].saturating_sub(held[dimension]);
            if amount[dimension].saturating_sub(own) > shared[dimension] {
                return false;
            }
        }
        true
    }

    pub fn try_acquire_all(self, amount: [usize; N]) -> bool {
        let held = self.state.held.get();
        let shared = self.shared.get();
        let mut next = [0; N];
        let mut borrowed = [0; N];
        for dimension in 0..N {
            let Some(next_held) = held[dimension].checked_add(amount[dimension]) else {
                return false;
            };
            next[dimension] = next_held;
            let own = self.state.reserve[dimension].saturating_sub(held[dimension]);
            borrowed[dimension] = amount[dimension].saturating_sub(own);
            if borrowed[dimension] > shared[dimension] {
                return false;
            }
        }
        self.state.held.set(next);
        self.shared.set(array::from_fn(|dimension| {
            shared[dimension] - borrowed[dimension]
        }));
        true
    }

    pub fn release_all(self, amount: [usize; N]) {
        let held = self.state.held.get();
        for dimension in 0..N {
            assert!(
                held[dimension] >= amount[dimension],
                "cannot release credits that are not held"
            );
        }
        let returned: [usize; N] = array::from_fn(|dimension| {
            amount[dimension].min(held[dimension].saturating_sub(self.state.reserve[dimension]))
        });
        self.state.held.set(array::from_fn(|dimension| {
            held[dimension] - amount[dimension]
        }));
        let shared = self.shared.get();
        self.shared.set(array::from_fn(|dimension| {
            shared[dimension] + returned[dimension]
        }));
    }
}

impl<const N: usize> Credits<N> {
    /// Builds independent dimensions using the default balanced reserve.
    pub fn from_capacities(capacity: [usize; N], lane_count: usize) -> Self {
        match Self::try_from_capacities(capacity, lane_count) {
            Ok(credits) => credits,
            Err(error) => error.abort(),
        }
    }

    pub fn try_from_capacities(
        capacity: [usize; N],
        lane_count: usize,
    ) -> Result<Self, collections::AllocationError> {
        assert!(N > 0, "credit dimension count must be positive");
        assert!(lane_count > 0, "credit lane count must be positive");
        let reserve = capacity.map(|amount| {
            if lane_count == 1 {
                amount
            } else {
                amount / lane_count / 2
            }
        });
        Self::try_with_reserve_per_lane(capacity, lane_count, reserve)
    }

    fn try_with_reserve_per_lane(
        capacity: [usize; N],
        lane_count: usize,
        reserve: [usize; N],
    ) -> Result<Self, collections::AllocationError> {
        assert!(N > 0, "credit dimension count must be positive");
        assert!(lane_count > 0, "credit lane count must be positive");
        let reserved: [usize; N] = array::from_fn(|dimension| {
            assert!(
                reserve[dimension] <= capacity[dimension] / lane_count,
                "credit reserve exceeds capacity"
            );
            reserve[dimension] * lane_count
        });
        Ok(Self {
            used: cell::Cell::new([0; N]),
            pool: Pool::try_uniform(
                array::from_fn(|dimension| capacity[dimension] - reserved[dimension]),
                lane_count,
                reserve,
            )?,
        })
    }

    fn credit(&self, lane: usize) -> Option<Lane<'_, N>> {
        self.pool.lane(lane)
    }

    fn can_acquire_all(&self, lane: usize, amount: [usize; N]) -> bool {
        self.credit(lane)
            .is_some_and(|credit| credit.can_acquire_all(amount))
    }

    /// Atomically checks and acquires every resource dimension.
    pub fn try_acquire_all(&self, lane: usize, amount: [usize; N]) -> bool {
        let Some(credit) = self.credit(lane) else {
            return false;
        };
        if !credit.try_acquire_all(amount) {
            return false;
        }
        let used = self.used.get();
        self.used.set(array::from_fn(|dimension| {
            used[dimension] + amount[dimension]
        }));
        true
    }

    /// Releases every resource dimension as one state transition.
    pub fn release_all(&self, lane: usize, amount: [usize; N]) {
        let credit = self.pool.lane(lane).expect("credit lane out of bounds");
        credit.release_all(amount);
        let used = self.used.get();
        self.used.set(array::from_fn(|dimension| {
            used[dimension] - amount[dimension]
        }));
    }
}

impl Credits {
    pub fn with_reserve(capacity: usize, lane_count: usize, reserve_per_lane: usize) -> Self {
        match Self::try_with_reserve(capacity, lane_count, reserve_per_lane) {
            Ok(credits) => credits,
            Err(error) => error.abort(),
        }
    }

    pub fn try_with_reserve(
        capacity: usize,
        lane_count: usize,
        reserve_per_lane: usize,
    ) -> Result<Self, collections::AllocationError> {
        Self::try_with_reserve_per_lane([capacity], lane_count, [reserve_per_lane])
    }

    pub fn used(&self) -> usize {
        self.used.get()[0]
    }

    pub fn held_by(&self, lane: usize) -> Option<usize> {
        self.pool.states.get(lane).map(|state| state.held.get()[0])
    }

    pub fn reserved_for(&self, lane: usize) -> Option<usize> {
        self.pool.states.get(lane).map(|state| state.reserve[0])
    }

    pub fn shared_available(&self) -> usize {
        self.pool.shared()[0]
    }

    pub fn can_acquire(&self, lane: usize, amount: usize) -> bool {
        self.can_acquire_all(lane, [amount])
    }

    pub fn try_acquire(&self, lane: usize, amount: usize) -> bool {
        self.try_acquire_all(lane, [amount])
    }

    pub fn release(&self, lane: usize, amount: usize) {
        self.release_all(lane, [amount]);
    }
}