use crate::{bitvec::BitVec, PAGE_SIZE_BITS};
use core::ptr::NonNull;

/// A single size class for a single region of the allocator. See the comment on the
/// `crate::allocators::buddy` module for more information.
pub struct Stripe<'buddy> {
    /// The base address of the tree this stripe is a part of.
    pub base: *const (),

    /// The order of the stripe. Order `n` means the subregions are `2ⁿ` pages in size.
    pub order: usize,

    /// The sentinel node of the free-list for this node. As an invariant of the type, there are no
    /// live references to any node in this list.
    pub free_list: NonNull<FreeListNode>,

    /// The bitset used to track whether a given subregion is allocated or not. A `true` bit
    /// corresponds to an allocated subregion.
    pub bitset: &'buddy mut BitVec,

    /// The offset from the start of the bitset to the region used by this stripe.
    pub bitset_offset: usize,
}

impl<'buddy> Stripe<'buddy> {
    /// Returns the buddy of the given pointer.
    ///
    /// ## Safety
    ///
    /// - The pointer must actually be part of this region.
    unsafe fn buddy_of(&self, ptr: NonNull<FreeListNode>) -> NonNull<FreeListNode> {
        let index = self.buddy_of_index(self.index_of(ptr));
        let addr = self.base as usize + (index << (PAGE_SIZE_BITS + self.order));
        NonNull::new_unchecked(addr as *mut FreeListNode)
    }

    /// Returns the buddy of the given index.
    fn buddy_of_index(&self, index: usize) -> usize {
        index ^ (1 << (PAGE_SIZE_BITS + self.order))
    }

    /// Returns the index the given pointer should have in the BitVec.
    fn index_of(&self, ptr: NonNull<FreeListNode>) -> usize {
        (ptr.as_ptr() as usize - self.base as usize) >> (PAGE_SIZE_BITS + self.order)
    }

    /// Pops a subregion from the free-list.
    pub fn pop(&mut self) -> Option<NonNull<FreeListNode>> {
        // SAFETY: The `free_list` is guaranteed to be valid by the invariants of the buddy
        // allocator. Retrieving the next pointer doesn't, on its own, break aliasing rules.
        let next = unsafe { self.free_list.read().next };

        // If the sentinel and the next pointer refer to the same spot, the free-list was empty, so
        // we can't pop from it.
        if self.free_list == next {
            return None;
        }

        // Otherwise, remove the node from the free-list.
        unsafe {
            FreeListNode::remove(next);
        }

        // Finally, mark the node as allocated in the bitvec.
        let index = self.index_of(next);
        assert!(self.bitset.get(self.bitset_offset + index));
        self.bitset.set(self.bitset_offset + index, true);

        Some(next)
    }

    /// Pushes a subregion back into the free-list.
    ///
    /// # Safety
    ///
    /// - There must be no live references to `subregion`.
    /// - `subregion` must not be a member of any list.
    pub unsafe fn push(&mut self, subregion: NonNull<FreeListNode>) {
        // Insert the subregion as the first element of the free-list.
        //
        // SAFETY: The free-list is guaranteed to be valid by the invariants of the buddy
        // allocator.
        unsafe {
            FreeListNode::insert(self.free_list, subregion);
        }

        // Mark the node as unallocated in the bitvec.
        let index = self.index_of(subregion);
        assert!(self.bitset.get(self.bitset_offset + index));
        self.bitset.set(self.bitset_offset + index, false);
    }

    /// Pushes a subregion into the free-list for the first time.
    ///
    /// # Safety
    ///
    /// - There must be no live references to `subregion`.
    /// - `subregion` must not be a member of any list.
    pub unsafe fn push_initial(&mut self, subregion: NonNull<FreeListNode>) {
        // Insert the subregion as the first element of the free-list.
        //
        // SAFETY: The free-list is guaranteed to be valid by the invariants of the buddy
        // allocator.
        unsafe {
            FreeListNode::insert(self.free_list, subregion);
        }

        // Mark the node as unallocated in the bitvec.
        let index = self.index_of(subregion);
        assert!(self.bitset.get(self.bitset_offset + index));
        self.bitset.set(self.bitset_offset + index, false);
    }
}

pub struct FreeListNode {
    next: NonNull<FreeListNode>,
    prev: NonNull<FreeListNode>,
}

impl FreeListNode {
    /// Inserts `new` after `prev`, initializing it. `prev` may be the sentinel.
    ///
    /// # Safety
    ///
    /// - There must be no live references to any node in the list that includes `prev`, including
    ///   the sentinel.
    /// - There must be no live references to `new`.
    /// - `new` must not be a member of any list.
    pub unsafe fn insert(prev: NonNull<FreeListNode>, new: NonNull<FreeListNode>) {
        let next = prev.read().next;
        (*prev.as_ptr()).next = new;
        (*next.as_ptr()).prev = new;
        new.write(FreeListNode { next, prev });
    }

    /// Removes this node from the free list it is a part of.
    ///
    /// # Safety
    ///
    /// - The pointer must point to a node that is part of a free list.
    /// - There must be no live references to any node in the list, including the sentinel.
    /// - The pointer must not point to the sentinel node.
    pub unsafe fn remove(ptr: NonNull<FreeListNode>) {
        let FreeListNode { next, prev } = ptr.read();
        (*next.as_ptr()).prev = prev;
        (*prev.as_ptr()).next = next;
    }
}