//! Bounded ASID ownership. A recycled identifier cannot be activated until
//! the caller confirms an ASID-scoped TLBI; this module does not execute TLBI.

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Lease {
    pub id: u8,
    pub generation: u64,
}

impl Lease {
    pub fn ttbr0_value(self, root: u64) -> Option<u64> {
        (self.id != 0 && root & 0xfff == 0 && root < (1u64 << 48))
            .then_some(root | ((self.id as u64) << 48))
    }

    pub fn tlbi_operand(self) -> u64 {
        (self.id as u64) << 48
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct Grant {
    pub lease: Lease,
    pub recycled: bool,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Error {
    InvalidCapacity,
    Full,
    GenerationExhausted,
    WrongLease,
    TlbiRequired,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum Slot {
    Virgin,
    Active(u64),
    Released(u64),
    PendingTlbi(u64),
}

pub struct Pool<const N: usize> {
    slots: [Slot; N],
    next_generation: u64,
}

impl<const N: usize> Default for Pool<N> {
    fn default() -> Self {
        Self::new()
    }
}

impl<const N: usize> Pool<N> {
    pub const fn new() -> Self {
        Self {
            slots: [Slot::Virgin; N],
            next_generation: 0,
        }
    }

    pub fn acquire(&mut self) -> Result<Grant, Error> {
        if N == 0 || N > u8::MAX as usize {
            return Err(Error::InvalidCapacity);
        }
        let index = self
            .slots
            .iter()
            .position(|slot| matches!(slot, Slot::Virgin | Slot::Released(_)))
            .ok_or(Error::Full)?;
        let generation = self
            .next_generation
            .checked_add(1)
            .ok_or(Error::GenerationExhausted)?;
        let recycled = matches!(self.slots[index], Slot::Released(_));
        self.slots[index] = if recycled {
            Slot::PendingTlbi(generation)
        } else {
            Slot::Active(generation)
        };
        self.next_generation = generation;
        Ok(Grant {
            lease: Lease {
                id: index as u8 + 1,
                generation,
            },
            recycled,
        })
    }

    pub fn ready(&self, lease: Lease) -> bool {
        self.slot(lease)
            .is_some_and(|slot| *slot == Slot::Active(lease.generation))
    }

    pub fn confirm_tlbi(&mut self, lease: Lease) -> Result<(), Error> {
        let slot = self.slot_mut(lease).ok_or(Error::WrongLease)?;
        if *slot != Slot::PendingTlbi(lease.generation) {
            return Err(Error::WrongLease);
        }
        *slot = Slot::Active(lease.generation);
        Ok(())
    }

    pub fn release(&mut self, lease: Lease) -> Result<(), Error> {
        let slot = self.slot_mut(lease).ok_or(Error::WrongLease)?;
        if *slot == Slot::PendingTlbi(lease.generation) {
            return Err(Error::TlbiRequired);
        }
        if *slot != Slot::Active(lease.generation) {
            return Err(Error::WrongLease);
        }
        *slot = Slot::Released(lease.generation);
        Ok(())
    }

    fn slot(&self, lease: Lease) -> Option<&Slot> {
        self.slots.get((lease.id as usize).checked_sub(1)?)
    }

    fn slot_mut(&mut self, lease: Lease) -> Option<&mut Slot> {
        self.slots.get_mut((lease.id as usize).checked_sub(1)?)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn recycled_asid_requires_targeted_tlbi_before_activation() {
        let mut pool = Pool::<1>::new();
        let first = pool.acquire().unwrap();
        assert_eq!(
            first.lease,
            Lease {
                id: 1,
                generation: 1
            }
        );
        assert!(!first.recycled);
        assert!(pool.ready(first.lease));
        assert_eq!(pool.acquire(), Err(Error::Full));
        pool.release(first.lease).unwrap();
        let second = pool.acquire().unwrap();
        assert_eq!(
            second.lease,
            Lease {
                id: 1,
                generation: 2
            }
        );
        assert!(second.recycled);
        assert!(!pool.ready(second.lease));
        assert_eq!(pool.release(second.lease), Err(Error::TlbiRequired));
        assert_eq!(pool.confirm_tlbi(first.lease), Err(Error::WrongLease));
        pool.confirm_tlbi(second.lease).unwrap();
        assert!(pool.ready(second.lease));
        assert_eq!(pool.release(first.lease), Err(Error::WrongLease));
        pool.release(second.lease).unwrap();
    }

    #[test]
    fn two_active_identifiers_are_distinct_and_stale_tokens_fail() {
        let mut pool = Pool::<2>::new();
        let a = pool.acquire().unwrap();
        let b = pool.acquire().unwrap();
        assert_eq!((a.lease.id, b.lease.id), (1, 2));
        assert_eq!(pool.acquire(), Err(Error::Full));
        assert_eq!(pool.confirm_tlbi(a.lease), Err(Error::WrongLease));
        assert_eq!(
            pool.release(Lease {
                id: 0,
                generation: 1
            }),
            Err(Error::WrongLease)
        );
        pool.release(a.lease).unwrap();
        assert!(pool.ready(b.lease));
        let reused = pool.acquire().unwrap();
        assert_eq!(reused.lease.id, 1);
        assert_eq!(reused.lease.generation, 3);
        assert!(reused.recycled);
        assert!(!pool.ready(reused.lease));
    }

    #[test]
    fn ttbr_and_tlbi_operands_contain_only_legal_asid_bits() {
        let lease = Lease {
            id: 7,
            generation: 1,
        };
        assert_eq!(lease.ttbr0_value(0x4000_5000), Some(0x0007_0000_4000_5000));
        assert_eq!(lease.tlbi_operand(), 0x0007_0000_0000_0000);
        assert_eq!(lease.ttbr0_value(0x4000_5001), None);
        assert_eq!(lease.ttbr0_value(1u64 << 48), None);
        assert_eq!(
            Lease {
                id: 0,
                generation: 1
            }
            .ttbr0_value(0x4000_5000),
            None
        );
    }

    #[test]
    fn invalid_capacity_and_generation_exhaustion_do_not_change_ownership() {
        assert_eq!(Pool::<0>::new().acquire(), Err(Error::InvalidCapacity));
        assert_eq!(Pool::<256>::new().acquire(), Err(Error::InvalidCapacity));
        let mut pool = Pool::<1>::new();
        pool.next_generation = u64::MAX;
        assert_eq!(pool.acquire(), Err(Error::GenerationExhausted));
        assert_eq!(pool.slots[0], Slot::Virgin);
    }
}