// offset-allocator/src/lib.rs

#![doc = include_str!("../README.md")]
#![deny(unsafe_code)]
#![warn(missing_docs)]

use std::fmt::{Debug, Display, Formatter, Result as FmtResult};

use log::debug;
use nonmax::{NonMaxU16, NonMaxU32};

pub mod ext;

mod small_float;

#[cfg(test)]
mod tests;

const NUM_TOP_BINS: usize = 32;
const BINS_PER_LEAF: usize = 8;
const TOP_BINS_INDEX_SHIFT: u32 = 3;
const LEAF_BINS_INDEX_MASK: u32 = 7;
const NUM_LEAF_BINS: usize = NUM_TOP_BINS * BINS_PER_LEAF;

/// Determines the number of allocations that the allocator supports.
///
/// By default, [`Allocator`] and related functions use `u32`, which allows for
/// `u32::MAX - 1` allocations. You can, however, use `u16` instead, which
/// causes the allocator to use less memory but limits the number of allocations
/// within a single allocator to at most 65,534.
pub trait NodeIndex: Clone + Copy + Default {
    /// The `NonMax` version of this type.
    ///
    /// This is used extensively to optimize `enum` representations.
    type NonMax: NodeIndexNonMax + TryFrom<Self> + Into<Self>;

    /// The maximum value representable in this type.
    const MAX: u32;

    /// Converts from a unsigned 32-bit integer to an instance of this type.
    fn from_u32(val: u32) -> Self;

    /// Converts this type to an unsigned machine word.
    fn to_usize(self) -> usize;
}

/// The `NonMax` version of the [`NodeIndex`].
///
/// For example, for `u32`, the `NonMax` version is [`NonMaxU32`].
pub trait NodeIndexNonMax: Clone + Copy + PartialEq + Default + Debug + Display {
    /// Converts this type to an unsigned machine word.
    fn to_usize(self) -> usize;
}

/// An allocator that manages a single contiguous chunk of space and hands out
/// portions of it as requested.
pub struct Allocator<NI = u32>
where
    NI: NodeIndex,
{
    size: u32,
    max_allocs: u32,
    free_storage: u32,

    used_bins_top: u32,
    used_bins: [u8; NUM_TOP_BINS],
    bin_indices: [Option<NI::NonMax>; NUM_LEAF_BINS],

    nodes: Vec<Node<NI>>,
    free_nodes: Vec<NI::NonMax>,
    free_offset: u32,
}

/// A single allocation.
#[derive(Clone, Copy)]
pub struct Allocation<NI = u32>
where
    NI: NodeIndex,
{
    /// The location of this allocation within the buffer.
    pub offset: NI,
    /// The node index associated with this allocation.
    metadata: NI::NonMax,
}

/// Provides a summary of the state of the allocator, including space remaining.
#[derive(Debug)]
pub struct StorageReport {
    /// The amount of free space left.
    pub total_free_space: u32,
    /// The maximum potential size of a single contiguous allocation.
    pub largest_free_region: u32,
}

/// Provides a detailed accounting of each bin within the allocator.
#[derive(Debug)]
pub struct StorageReportFull {
    /// Each bin within the allocator.
    pub free_regions: [StorageReportFullRegion; NUM_LEAF_BINS],
}

/// A detailed accounting of each allocator bin.
#[derive(Clone, Copy, Debug, Default)]
pub struct StorageReportFullRegion {
    /// The size of the bin, in units.
    pub size: u32,
    /// The number of allocations in the bin.
    pub count: u32,
}

#[derive(Clone, Copy, Default)]
struct Node<NI = u32>
where
    NI: NodeIndex,
{
    data_offset: u32,
    data_size: u32,
    bin_list_prev: Option<NI::NonMax>,
    bin_list_next: Option<NI::NonMax>,
    neighbor_prev: Option<NI::NonMax>,
    neighbor_next: Option<NI::NonMax>,
    used: bool, // TODO: Merge as bit flag
}

// Utility functions
fn find_lowest_bit_set_after(bit_mask: u32, start_bit_index: u32) -> Option<NonMaxU32> {
    let mask_before_start_index = (1 << start_bit_index) - 1;
    let mask_after_start_index = !mask_before_start_index;
    let bits_after = bit_mask & mask_after_start_index;
    if bits_after == 0 {
        None
    } else {
        NonMaxU32::try_from(bits_after.trailing_zeros()).ok()
    }
}

impl<NI> Allocator<NI>
where
    NI: NodeIndex,
{
    /// Creates a new allocator, managing a contiguous block of memory of `size`
    /// units, with a default reasonable number of maximum allocations.
    pub fn new(size: u32) -> Self {
        Allocator::with_max_allocs(size, u32::min(128 * 1024, NI::MAX - 1))
    }

    /// Creates a new allocator, managing a contiguous block of memory of `size`
    /// units, with the given number of maximum allocations.
    ///
    /// Note that the maximum number of allocations must be less than
    /// [`NodeIndex::MAX`] minus one. If this restriction is violated, this
    /// constructor will panic.
    pub fn with_max_allocs(size: u32, max_allocs: u32) -> Self {
        assert!(max_allocs < NI::MAX - 1);

        let mut this = Self {
            size,
            max_allocs,
            free_storage: 0,
            used_bins_top: 0,
            free_offset: 0,
            used_bins: [0; NUM_TOP_BINS],
            bin_indices: [None; NUM_LEAF_BINS],
            nodes: vec![],
            free_nodes: vec![],
        };
        this.reset();
        this
    }

    /// Clears out all allocations.
    pub fn reset(&mut self) {
        self.free_storage = 0;
        self.used_bins_top = 0;
        self.free_offset = self.max_allocs - 1;

        self.used_bins.iter_mut().for_each(|bin| *bin = 0);

        self.bin_indices.iter_mut().for_each(|index| *index = None);

        self.nodes = vec![Node::default(); self.max_allocs as usize];

        // Freelist is a stack. Nodes in inverse order so that [0] pops first.
        self.free_nodes = (0..self.max_allocs)
            .map(|i| {
                NI::NonMax::try_from(NI::from_u32(self.max_allocs - i - 1)).unwrap_or_default()
            })
            .collect();

        // Start state: Whole storage as one big node
        // Algorithm will split remainders and push them back as smaller nodes
        self.insert_node_into_bin(self.size, 0);
    }

    /// Allocates a block of `size` elements and returns its allocation.
    ///
    /// If there's not enough contiguous space for this allocation, returns
    /// None.
    pub fn allocate(&mut self, size: u32) -> Option<Allocation<NI>> {
        // Out of allocations?
        if self.free_offset == 0 {
            return None;
        }

        // Round up to bin index to ensure that alloc >= bin
        // Gives us min bin index that fits the size
        let min_bin_index = small_float::uint_to_float_round_up(size);

        let min_top_bin_index = min_bin_index >> TOP_BINS_INDEX_SHIFT;
        let min_leaf_bin_index = min_bin_index & LEAF_BINS_INDEX_MASK;

        let mut top_bin_index = min_top_bin_index;
        let mut leaf_bin_index = None;

        // If top bin exists, scan its leaf bin. This can fail (NO_SPACE).
        if (self.used_bins_top & (1 << top_bin_index)) != 0 {
            leaf_bin_index = find_lowest_bit_set_after(
                self.used_bins[top_bin_index as usize] as _,
                min_leaf_bin_index,
            );
        }

        // If we didn't find space in top bin, we search top bin from +1
        let leaf_bin_index = match leaf_bin_index {
            Some(leaf_bin_index) => leaf_bin_index,
            None => {
                top_bin_index =
                    find_lowest_bit_set_after(self.used_bins_top, min_top_bin_index + 1)?.into();

                // All leaf bins here fit the alloc, since the top bin was
                // rounded up. Start leaf search from bit 0.
                //
                // NOTE: This search can't fail since at least one leaf bit was
                // set because the top bit was set.
                NonMaxU32::try_from(self.used_bins[top_bin_index as usize].trailing_zeros())
                    .unwrap()
            }
        };

        let bin_index = (top_bin_index << TOP_BINS_INDEX_SHIFT) | u32::from(leaf_bin_index);

        // Pop the top node of the bin. Bin top = node.next.
        let node_index = self.bin_indices[bin_index as usize].unwrap();
        let node = &mut self.nodes[node_index.to_usize()];
        let node_total_size = node.data_size;
        node.data_size = size;
        node.used = true;
        self.bin_indices[bin_index as usize] = node.bin_list_next;
        if let Some(bin_list_next) = node.bin_list_next {
            self.nodes[bin_list_next.to_usize()].bin_list_prev = None;
        }
        self.free_storage -= node_total_size;
        debug!(
            "Free storage: {} (-{}) (allocate)",
            self.free_storage, node_total_size
        );

        // Bin empty?
        if self.bin_indices[bin_index as usize].is_none() {
            // Remove a leaf bin mask bit
            self.used_bins[top_bin_index as usize] &= !(1 << u32::from(leaf_bin_index));

            // All leaf bins empty?
            if self.used_bins[top_bin_index as usize] == 0 {
                // Remove a top bin mask bit
                self.used_bins_top &= !(1 << top_bin_index);
            }
        }

        // Push back remainder N elements to a lower bin
        let remainder_size = node_total_size - size;
        if remainder_size > 0 {
            let Node {
                data_offset,
                neighbor_next,
                ..
            } = self.nodes[node_index.to_usize()];

            let new_node_index = self.insert_node_into_bin(remainder_size, data_offset + size);

            // Link nodes next to each other so that we can merge them later if both are free
            // And update the old next neighbor to point to the new node (in middle)
            let node = &mut self.nodes[node_index.to_usize()];
            if let Some(neighbor_next) = node.neighbor_next {
                self.nodes[neighbor_next.to_usize()].neighbor_prev = Some(new_node_index);
            }
            self.nodes[new_node_index.to_usize()].neighbor_prev = Some(node_index);
            self.nodes[new_node_index.to_usize()].neighbor_next = neighbor_next;
            self.nodes[node_index.to_usize()].neighbor_next = Some(new_node_index);
        }

        let node = &mut self.nodes[node_index.to_usize()];
        Some(Allocation {
            offset: NI::from_u32(node.data_offset),
            metadata: node_index,
        })
    }

    /// Frees an allocation, returning the data to the heap.
    ///
    /// If the allocation has already been freed, the behavior is unspecified.
    /// It may or may not panic. Note that, because this crate contains no
    /// unsafe code, the memory safe of the allocator *itself* will be
    /// uncompromised, even on double free.
    pub fn free(&mut self, allocation: Allocation<NI>) {
        let node_index = allocation.metadata;

        // Merge with neighbors…
        let Node {
            data_offset: mut offset,
            data_size: mut size,
            used,
            ..
        } = self.nodes[node_index.to_usize()];

        // Double delete check
        assert!(used);

        if let Some(neighbor_prev) = self.nodes[node_index.to_usize()].neighbor_prev {
            if !self.nodes[neighbor_prev.to_usize()].used {
                // Previous (contiguous) free node: Change offset to previous
                // node offset. Sum sizes
                let prev_node = &self.nodes[neighbor_prev.to_usize()];
                offset = prev_node.data_offset;
                size += prev_node.data_size;

                // Remove node from the bin linked list and put it in the
                // freelist
                self.remove_node_from_bin(neighbor_prev);

                let prev_node = &self.nodes[neighbor_prev.to_usize()];
                debug_assert_eq!(prev_node.neighbor_next, Some(node_index));
                self.nodes[node_index.to_usize()].neighbor_prev = prev_node.neighbor_prev;
            }
        }

        if let Some(neighbor_next) = self.nodes[node_index.to_usize()].neighbor_next {
            if !self.nodes[neighbor_next.to_usize()].used {
                // Next (contiguous) free node: Offset remains the same. Sum
                // sizes.
                let next_node = &self.nodes[neighbor_next.to_usize()];
                size += next_node.data_size;

                // Remove node from the bin linked list and put it in the
                // freelist
                self.remove_node_from_bin(neighbor_next);

                let next_node = &self.nodes[neighbor_next.to_usize()];
                debug_assert_eq!(next_node.neighbor_prev, Some(node_index));
                self.nodes[node_index.to_usize()].neighbor_next = next_node.neighbor_next;
            }
        }

        let Node {
            neighbor_next,
            neighbor_prev,
            ..
        } = self.nodes[node_index.to_usize()];

        // Insert the removed node to freelist
        debug!(
            "Putting node {} into freelist[{}] (free)",
            node_index,
            self.free_offset + 1
        );
        self.free_offset += 1;
        self.free_nodes[self.free_offset as usize] = node_index;

        // Insert the (combined) free node to bin
        let combined_node_index = self.insert_node_into_bin(size, offset);

        // Connect neighbors with the new combined node
        if let Some(neighbor_next) = neighbor_next {
            self.nodes[combined_node_index.to_usize()].neighbor_next = Some(neighbor_next);
            self.nodes[neighbor_next.to_usize()].neighbor_prev = Some(combined_node_index);
        }
        if let Some(neighbor_prev) = neighbor_prev {
            self.nodes[combined_node_index.to_usize()].neighbor_prev = Some(neighbor_prev);
            self.nodes[neighbor_prev.to_usize()].neighbor_next = Some(combined_node_index);
        }
    }

    fn insert_node_into_bin(&mut self, size: u32, data_offset: u32) -> NI::NonMax {
        // Round down to bin index to ensure that bin >= alloc
        let bin_index = small_float::uint_to_float_round_down(size);

        let top_bin_index = bin_index >> TOP_BINS_INDEX_SHIFT;
        let leaf_bin_index = bin_index & LEAF_BINS_INDEX_MASK;

        // Bin was empty before?
        if self.bin_indices[bin_index as usize].is_none() {
            // Set bin mask bits
            self.used_bins[top_bin_index as usize] |= 1 << leaf_bin_index;
            self.used_bins_top |= 1 << top_bin_index;
        }

        // Take a freelist node and insert on top of the bin linked list (next = old top)
        let top_node_index = self.bin_indices[bin_index as usize];
        let free_offset = self.free_offset;
        let node_index = self.free_nodes[free_offset as usize];
        self.free_offset -= 1;
        debug!(
            "Getting node {} from freelist[{}]",
            node_index,
            self.free_offset + 1
        );
        self.nodes[node_index.to_usize()] = Node {
            data_offset,
            data_size: size,
            bin_list_next: top_node_index,
            ..Node::default()
        };
        if let Some(top_node_index) = top_node_index {
            self.nodes[top_node_index.to_usize()].bin_list_prev = Some(node_index);
        }
        self.bin_indices[bin_index as usize] = Some(node_index);

        self.free_storage += size;
        debug!(
            "Free storage: {} (+{}) (insert_node_into_bin)",
            self.free_storage, size
        );
        node_index
    }

    fn remove_node_from_bin(&mut self, node_index: NI::NonMax) {
        // Copy the node to work around borrow check.
        let node = self.nodes[node_index.to_usize()];

        match node.bin_list_prev {
            Some(bin_list_prev) => {
                // Easy case: We have previous node. Just remove this node from the middle of the list.
                self.nodes[bin_list_prev.to_usize()].bin_list_next = node.bin_list_next;
                if let Some(bin_list_next) = node.bin_list_next {
                    self.nodes[bin_list_next.to_usize()].bin_list_prev = node.bin_list_prev;
                }
            }
            None => {
                // Hard case: We are the first node in a bin. Find the bin.

                // Round down to bin index to ensure that bin >= alloc
                let bin_index = small_float::uint_to_float_round_down(node.data_size);

                let top_bin_index = (bin_index >> TOP_BINS_INDEX_SHIFT) as usize;
                let leaf_bin_index = (bin_index & LEAF_BINS_INDEX_MASK) as usize;

                self.bin_indices[bin_index as usize] = node.bin_list_next;
                if let Some(bin_list_next) = node.bin_list_next {
                    self.nodes[bin_list_next.to_usize()].bin_list_prev = None;
                }

                // Bin empty?
                if self.bin_indices[bin_index as usize].is_none() {
                    // Remove a leaf bin mask bit
                    self.used_bins[top_bin_index as usize] &= !(1 << leaf_bin_index);

                    // All leaf bins empty?
                    if self.used_bins[top_bin_index as usize] == 0 {
                        // Remove a top bin mask bit
                        self.used_bins_top &= !(1 << top_bin_index);
                    }
                }
            }
        }

        // Insert the node to freelist
        debug!(
            "Putting node {} into freelist[{}] (remove_node_from_bin)",
            node_index,
            self.free_offset + 1
        );
        self.free_offset += 1;
        self.free_nodes[self.free_offset as usize] = node_index;

        self.free_storage -= node.data_size;
        debug!(
            "Free storage: {} (-{}) (remove_node_from_bin)",
            self.free_storage, node.data_size
        );
    }

    /// Returns the *used* size of an allocation.
    ///
    /// Note that this may be larger than the size requested at allocation time,
    /// due to rounding.
    pub fn allocation_size(&self, allocation: Allocation<NI>) -> u32 {
        self.nodes
            .get(allocation.metadata.to_usize())
            .map(|node| node.data_size)
            .unwrap_or_default()
    }

    /// Returns a structure containing the amount of free space remaining, as
    /// well as the largest amount that can be allocated at once.
    pub fn storage_report(&self) -> StorageReport {
        let mut largest_free_region = 0;
        let mut free_storage = 0;

        // Out of allocations? -> Zero free space
        if self.free_offset > 0 {
            free_storage = self.free_storage;
            if self.used_bins_top > 0 {
                let top_bin_index = 31 - self.used_bins_top.leading_zeros();
                let leaf_bin_index =
                    31 - (self.used_bins[top_bin_index as usize] as u32).leading_zeros();
                largest_free_region = small_float::float_to_uint(
                    (top_bin_index << TOP_BINS_INDEX_SHIFT) | leaf_bin_index,
                );
                debug_assert!(free_storage >= largest_free_region);
            }
        }

        StorageReport {
            total_free_space: free_storage,
            largest_free_region,
        }
    }

    /// Returns detailed information about the number of allocations in each
    /// bin.
    pub fn storage_report_full(&self) -> StorageReportFull {
        let mut report = StorageReportFull::default();
        for i in 0..NUM_LEAF_BINS {
            let mut count = 0;
            let mut maybe_node_index = self.bin_indices[i];
            while let Some(node_index) = maybe_node_index {
                maybe_node_index = self.nodes[node_index.to_usize()].bin_list_next;
                count += 1;
            }
            report.free_regions[i] = StorageReportFullRegion {
                size: small_float::float_to_uint(i as u32),
                count,
            }
        }
        report
    }
}

impl Default for StorageReportFull {
    fn default() -> Self {
        Self {
            free_regions: [Default::default(); NUM_LEAF_BINS],
        }
    }
}

impl<NI> Debug for Allocator<NI>
where
    NI: NodeIndex,
{
    fn fmt(&self, f: &mut Formatter<'_>) -> FmtResult {
        self.storage_report().fmt(f)
    }
}

impl NodeIndex for u32 {
    type NonMax = NonMaxU32;
    const MAX: u32 = u32::MAX;

    fn from_u32(val: u32) -> Self {
        val
    }

    fn to_usize(self) -> usize {
        self as usize
    }
}

impl NodeIndex for u16 {
    type NonMax = NonMaxU16;
    const MAX: u32 = u16::MAX as u32;

    fn from_u32(val: u32) -> Self {
        val as u16
    }

    fn to_usize(self) -> usize {
        self as usize
    }
}

impl NodeIndexNonMax for NonMaxU32 {
    fn to_usize(self) -> usize {
        u32::from(self) as usize
    }
}

impl NodeIndexNonMax for NonMaxU16 {
    fn to_usize(self) -> usize {
        u16::from(self) as usize
    }
}