chunked-range-alloc 1.0.0

/*!
A simple range allocator for chunked external memory

[`ChunkedRangeAlloc`] was created for 2 use cases:

1. packing game assets into archive files
2. basis of specialized vulkan memory allocator

## Features:

- [`Allocation`] includes [`Allocation::chunk_index`] in addition to [`Allocation::offset`] and [`Allocation::len`]
- [`ChunkedRangeAlloc::from_allocations`] constructor for loading existing allocations (example: game assets index)
- optional `bincode` and `serde` support, see [`Allocation`]
- simple, safe code
- _good enough_ performance: allocator uses `BTree` internally, best-fit search strategy, immediately coalesces on free

## Non-goals:

- blazingly fast constant O(🚀) time complexity

## Example

```
use std::num::NonZeroU32;

use chunked_range_alloc::ChunkedRangeAlloc;

// create allocator with chunk size = 4
let mut alloc = ChunkedRangeAlloc::new(NonZeroU32::new(4).unwrap());

// allocate size 1 with align = 1
let one = alloc.alloc(NonZeroU32::MIN, NonZeroU32::MIN);
assert_eq!(one.chunk_index, 0);
assert_eq!(one.offset, 0);
assert_eq!(one.len.get(), 1);

// allocate size 1 with align 2
let two = alloc.alloc(NonZeroU32::MIN, NonZeroU32::new(2).unwrap());
assert_eq!(two.chunk_index, 0);
assert_eq!(two.offset, 2); // offset is not 1 because of alignment
assert_eq!(two.len.get(), 1);

let free_chunk = alloc.free(one);
assert!(!free_chunk);

// free returns true if chunk becomes completely free
// you can eagerly free external memory in that case
let free_chunk = alloc.free(two);
assert!(free_chunk);

```
*/

#![forbid(unsafe_code)]
#![warn(clippy::pedantic)]

use std::num::NonZeroU32;

use std::collections::{BTreeMap, BTreeSet};

/**
[`ChunkedRangeAlloc`] allocation

Serialization support:

- `bincode` feature: [`bincode::Encode`] and [`bincode::Decode`]
- `serde` feature: [`serde::Serialize`] and [`serde::Deserialize`]

Size is limited to [`u32::MAX`] (4GB) for a few reasons:

- packed 12-byte struct with 4-byte alignment
- _chunked_ allocation implies relatively small chunk sizes, not 4GB ones

Provides niche so that `Option<Allocation>` is the same size as `Allocation`

```
use std::mem::{align_of, size_of};
use chunked_range_alloc::Allocation;

assert_eq!(size_of::<Allocation>(), 12);
assert_eq!(align_of::<Allocation>(), 4);
assert_eq!(size_of::<Option<Allocation>>(), size_of::<Allocation>());
```
*/
#[must_use]
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[cfg_attr(feature = "bincode", derive(bincode::Encode, bincode::Decode))]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct Allocation {
    /// Index of chunk this allocation belongs to
    pub chunk_index: u32,
    /// Offset within chunk
    pub offset: u32,
    /// Size of allocation
    pub len: NonZeroU32,
}

/// A simple range allocator for chunked external memory, see [`Self::alloc`]
pub struct ChunkedRangeAlloc {
    chunk_size: NonZeroU32,
    /// Free ranges sorted by size -> chunk -> offset
    free_ranges: BTreeSet<FreeRange>,
    chunks: Vec<Chunk>,
    active_chunks: usize,
    allocated_memory: u64,
}

#[derive(Default)]
struct Chunk {
    free_offsets: BTreeMap<u32, NonZeroU32>,
}

impl Chunk {
    /// First free range before `offset`
    #[inline]
    fn previous_free_range(&self, offset: u32) -> Option<(u32, NonZeroU32)> {
        self.free_offsets
            .range(..=offset)
            .next_back()
            .map(|(&off, &len)| (off, len))
    }

    /// First free range after `offset`
    #[inline]
    fn next_free_range(&self, offset: u32) -> Option<(u32, NonZeroU32)> {
        self.free_offsets
            .range(offset..)
            .next()
            .map(|(&off, &len)| (off, len))
    }
}

/// Sorted by field order: len -> chunk -> offset
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
struct FreeRange {
    /// sort: first
    len: NonZeroU32,
    /// sort: second
    chunk: u32,
    /// sort: third
    offset: u32,
}

struct AlignedFreeRange {
    free_range: FreeRange,
    end: u32,
    aligned_offset: u32,
    aligned_end: u32,
}

impl ChunkedRangeAlloc {
    /// Constructs a new, empty [`ChunkedRangeAlloc`] with specified `chunk_size`
    #[must_use]
    #[inline]
    pub fn new(chunk_size: NonZeroU32) -> Self {
        ChunkedRangeAlloc {
            chunk_size,
            free_ranges: BTreeSet::new(),
            chunks: Vec::new(),
            active_chunks: 0,
            allocated_memory: 0,
        }
    }

    /// Constructs [`ChunkedRangeAlloc`] with specified `chunk_size` and inserts `allocations`
    ///
    /// Returns [`None`] if any allocation is invalid or overlaps
    #[must_use]
    #[inline]
    pub fn from_allocations<'a>(
        chunk_size: NonZeroU32,
        allocations: impl IntoIterator<Item = &'a Allocation>,
    ) -> Option<Self> {
        let mut a = Self::new(chunk_size);

        for allocation in allocations {
            let chunk_index = usize::try_from(allocation.chunk_index).ok()?;

            // create chunks up to `chunk`
            let chunk = loop {
                if let Some(chunk) = a.chunks.get(chunk_index) {
                    break chunk;
                }

                let chunk = u32::try_from(a.chunks.len()).ok()?;

                // create new uninitialized chunk
                a.chunks.push(Chunk::default());

                // immediately insert free range covering entire chunk
                a.insert_free_range(FreeRange {
                    len: chunk_size,
                    chunk,
                    offset: 0,
                });
            };

            let (free_offset, free_len) = chunk.previous_free_range(allocation.offset)?;

            // encure allocation is in free range
            let free_end = free_offset + free_len.get();
            let allocation_end = allocation.offset + allocation.len.get();

            if !(free_offset <= allocation.offset && free_end >= allocation_end) {
                return None;
            }

            a.alloc_from(&AlignedFreeRange {
                free_range: FreeRange {
                    len: free_len,
                    chunk: allocation.chunk_index,
                    offset: free_offset,
                },
                end: free_offset + free_len.get(),
                aligned_offset: allocation.offset,
                aligned_end: allocation.offset + allocation.len.get(),
            });
        }

        Some(a)
    }

    /// Chunk size
    #[must_use]
    #[inline]
    pub fn chunk_size(&self) -> NonZeroU32 {
        self.chunk_size
    }

    /**
    Allocates memory with specified `size` and `align`, alignment MUST be power of two

    Use [`NonZeroU32::MIN`] as `align` if alignment is not required

    Attempts to best-fit find suitable contiguous range, otherwise creates new chunk

    Caller is responsible for allocating external memory when encountering new chunk index, chunk indices are contiguous

    # Panics

    - panics if `size` exceeds [`Self::chunk_size`]
    - panics if `align` is not power of two
    - panics if chunk count exceeds [`u32::MAX`]
     */
    #[inline]
    pub fn alloc(&mut self, size: NonZeroU32, align: NonZeroU32) -> Allocation {
        if let Some(aligned_free_range) = self.find_free_range(size, align) {
            let allocation = Allocation {
                chunk_index: aligned_free_range.free_range.chunk,
                offset: aligned_free_range.aligned_offset,
                len: size,
            };

            self.alloc_from(&aligned_free_range);

            return allocation;
        }

        let len = size;
        let size = size.get();

        let chunk = u32::try_from(self.chunks.len()).expect("too many chunks");

        let tail_len = self
            .chunk_size
            .get()
            .checked_sub(size)
            .expect("size exceeds chunk_size");

        // create new uninitialized chunk
        self.chunks.push(Chunk::default());

        // insert tail [size, end) if required
        if let Some(len) = NonZeroU32::new(tail_len) {
            self.insert_free_range(FreeRange {
                len,
                chunk,
                offset: size,
            });
        }

        self.active_chunks += 1;

        self.allocated_memory += u64::from(size);

        Allocation {
            chunk_index: chunk,
            offset: 0,
            len,
        }
    }

    /**
    Frees specified `allocation`

    Returns `true` if chunk with index specified in `allocation` becomes completely free

    # Panics

    Panics if allocation is invalid or double freed
     */
    #[must_use]
    #[inline]
    pub fn free(&mut self, allocation: Allocation) -> bool {
        let chunk_index = usize::try_from(allocation.chunk_index).expect("invalid chunk index");
        let chunk = self.chunks.get(chunk_index).expect("invalid chunk index");

        let previous_free_range = chunk.previous_free_range(allocation.offset);
        let next_free_range = chunk.next_free_range(allocation.offset);

        let mut free_range = FreeRange {
            len: allocation.len,
            chunk: allocation.chunk_index,
            offset: allocation.offset,
        };

        // coalesce before
        if let Some((offset, len)) = previous_free_range {
            let previous_end = offset + len.get();

            assert!(previous_end <= free_range.offset, "double free");

            if previous_end == free_range.offset {
                free_range.offset = offset;
                free_range.len = free_range
                    .len
                    .checked_add(len.get())
                    .expect("invalid allocation");

                self.remove_free_range(FreeRange {
                    len,
                    chunk: free_range.chunk,
                    offset,
                });
            }
        }

        // coalesce after
        if let Some((offset, len)) = next_free_range {
            let end = free_range.offset + free_range.len.get();

            assert!(end <= offset, "double free");

            if end == offset {
                free_range.len = free_range
                    .len
                    .checked_add(len.get())
                    .expect("invalid allocation");

                self.remove_free_range(FreeRange {
                    len,
                    chunk: free_range.chunk,
                    offset,
                });
            }
        }

        self.insert_free_range(free_range);

        self.allocated_memory -= u64::from(allocation.len.get());

        let free_chunk = free_range.len >= self.chunk_size;
        self.active_chunks -= usize::from(free_chunk);

        free_chunk
    }

    /// Amount of tracked chunks, both active and free/unused
    #[must_use]
    #[inline]
    pub fn chunk_count(&self) -> usize {
        self.chunks.len()
    }

    /// Amount of chunks containing live allocations
    #[must_use]
    #[inline]
    pub fn active_chunks(&self) -> usize {
        self.active_chunks
    }

    /// Shorthand for `chunk_count() - active_chunks()`
    #[must_use]
    #[inline]
    pub fn unused_chunks(&self) -> usize {
        self.chunk_count() - self.active_chunks
    }

    /// Shorthand for `chunk_count() * chunk_size()`
    #[must_use]
    #[inline]
    pub fn capacity(&self) -> u64 {
        u64::try_from(self.chunk_count()).unwrap_or(u64::MAX) * u64::from(self.chunk_size.get())
    }

    /// Sum of currently allocated [`Allocation::len`]s
    #[must_use]
    #[inline]
    pub fn allocated_memory(&self) -> u64 {
        self.allocated_memory
    }

    /// Shorthand for `capacity() - allocated_memory()`
    #[must_use]
    #[inline]
    pub fn remaining(&self) -> u64 {
        self.capacity() - self.allocated_memory
    }

    /// Shorthand for `active_chunks() * chunk_size()`
    #[must_use]
    #[inline]
    pub fn used_memory(&self) -> u64 {
        u64::try_from(self.active_chunks).unwrap_or(u64::MAX) * u64::from(self.chunk_size.get())
    }

    /// Shorthand for `used_memory() - allocated_memory()`
    #[must_use]
    #[inline]
    pub fn unused_memory(&self) -> u64 {
        self.used_memory() - self.allocated_memory
    }

    #[inline]
    fn alloc_from(&mut self, aligned_free_range: &AlignedFreeRange) {
        let &AlignedFreeRange {
            free_range,
            end,
            aligned_offset,
            aligned_end,
        } = aligned_free_range;

        self.remove_free_range(free_range);

        // insert head [free_range.offset, aligned_offset) if required
        if let Some(len) = NonZeroU32::new(aligned_offset - free_range.offset) {
            self.insert_free_range(FreeRange {
                len,
                chunk: free_range.chunk,
                offset: free_range.offset,
            });
        }

        // insert tail [aligned_end, end) if required
        if let Some(len) = NonZeroU32::new(end - aligned_end) {
            self.insert_free_range(FreeRange {
                len,
                chunk: free_range.chunk,
                offset: aligned_end,
            });
        }

        let free_chunk = free_range.len >= self.chunk_size;
        self.active_chunks += usize::from(free_chunk);

        self.allocated_memory += u64::from(aligned_end - aligned_offset);
    }

    /**
    Finds first free range which could fit aligned size

    # Panics

    - panics if `size` exceeds [`Self::chunk_size`]
    - panics if `align` is not power of two
    */
    #[inline]
    fn find_free_range(&self, size: NonZeroU32, align: NonZeroU32) -> Option<AlignedFreeRange> {
        assert!(size <= self.chunk_size);
        assert!(align.is_power_of_two());

        let align = align.get();

        let start = FreeRange {
            len: size,
            chunk: 0,
            offset: 0,
        };

        let size = size.get();

        // finding free range suitable for aligned allocation is not as simple as picking first big enough
        self.free_ranges
            .range(start..)
            .copied()
            .find_map(|free_range| {
                let end = free_range.offset + free_range.len.get();

                // align  offset and end withing free range, skip on overflow
                let aligned_offset = free_range.offset.checked_next_multiple_of(align)?;
                let aligned_end = aligned_offset.checked_add(size)?;

                // check aligned fit
                if aligned_end > end {
                    return None;
                }

                Some(AlignedFreeRange {
                    free_range,
                    end,
                    aligned_offset,
                    aligned_end,
                })
            })
    }

    /// Inserts free range into `free_ranges` and `chunks`, updates `active_chunks` and `allocated_memory`
    ///
    /// Returns true if chunk is entirely free
    #[inline]
    fn insert_free_range(&mut self, free_range: FreeRange) {
        let chunk_index = usize::try_from(free_range.chunk).expect("invalid chunk index");
        let chunk = self
            .chunks
            .get_mut(chunk_index)
            .expect("invalid chunk index");

        let inserted = self.free_ranges.insert(free_range);
        assert!(inserted);

        let last = chunk.free_offsets.insert(free_range.offset, free_range.len);
        assert!(last.is_none());
    }

    /// Removes free range from `free_ranges` and `chunks`, updates `active_chunks` and `allocated_memory`
    #[inline]
    fn remove_free_range(&mut self, free_range: FreeRange) {
        let chunk_index = usize::try_from(free_range.chunk).expect("invalid chunk index");
        let chunk = self
            .chunks
            .get_mut(chunk_index)
            .expect("invalid chunk index");

        let removed = self.free_ranges.remove(&free_range);
        assert!(removed);

        let removed = chunk.free_offsets.remove(&free_range.offset);
        assert_eq!(removed, Some(free_range.len));
    }
}