realloc 0.1.0

use core::ptr::NonNull;

use crate::lock::Mutex;
use crate::{Align, Alloc, Allocator, Brand, Error, Layout};

// If changed, update Arena::set_chunk_size docs
const DEFAULT_CHUNK_SIZE: usize = 1024 * 1024; // 1 MiB

/// An allocator that uses another allocator to obtain large blocks of memory,
/// then hands those out in smaller chunks. Only once all chunks have been freed
/// can the memory as a whole be freed.
///
/// Note that the underlying memory is not freed on drop; only once you call
/// [`dealloc_inner`](Arena::dealloc_inner) will the underlying memory get freed.
pub struct Arena<'alloc> {
    /// The inner allocator this delegates to.
    pub allocator: &'alloc dyn Allocator,
    inner: Mutex<ArenaInner<'alloc>>,
    brand: Brand,
}

struct ArenaInner<'alloc> {
    chunk_size: usize,
    buf: Option<Alloc<'alloc, u8>>,
    next: usize,
    count: usize,
}

#[cfg(feature = "global")]
impl Arena<'static> {
    /// Create a new arena allocator in the global allocator, if any.
    ///
    /// Only available on feature `global`.
    pub fn new() -> Option<Self> {
        crate::global::get().map(Self::new_in)
    }
}

impl<'alloc> Arena<'alloc> {
    /// Create a new arena allocator in a given allocator.
    pub fn new_in(allocator: &'alloc dyn Allocator) -> Self {
        Self {
            allocator,
            inner: Mutex::new(ArenaInner {
                chunk_size: DEFAULT_CHUNK_SIZE,
                buf: None,
                next: 0,
                count: 0,
            }),
            brand: Brand::new(),
        }
    }

    /// Get the chunk size to use when allocating the underlying memory.
    ///
    /// The default size is subject to change; at the moment, it is 1MiB. Each
    /// time it's forced to grow, it doubles the size.
    pub fn chunk_size(&self) -> usize {
        self.inner.lock().chunk_size
    }

    /// Set the chunk size to use when allocating the underlying memory.
    ///
    /// The default size is subject to change; at the moment, it is 1MiB. Each
    /// time it's forced to grow, it doubles the size.
    pub fn set_chunk_size(&mut self, chunk_size: usize) {
        let this = self.inner.get_mut();
        if chunk_size > this.chunk_size {
            this.chunk_size = chunk_size;
        }
    }

    /// Deallocate all the data that's been allocated so far.
    ///
    /// Requires (through borrow-checking) that all outstanding allocations have
    /// been either freed or forgotten.
    pub fn dealloc_inner(&mut self) {
        let this = self.inner.get_mut();
        // This is sound because the exclusivity of `&mut self` requires no
        // outstanding references, which `Alloc`s contain. Therefore, all
        // existing allocations must be either freed or forgotten.
        if let Some(buf) = this.buf.take() {
            buf.dealloc();
            this.chunk_size = DEFAULT_CHUNK_SIZE;
        }
    }
}

unsafe impl Allocator for Arena<'_> {
    fn brand(&self) -> &Brand {
        &self.brand
    }

    fn alloc_bytes(&self, layout: Layout) -> Result<Alloc<'_, u8>, Error> {
        let mut this_lock = self.inner.lock();
        let this = &mut *this_lock;

        let buf = match &mut this.buf {
            Some(b) => b,
            None => {
                let layout = Layout {
                    size: this.chunk_size,
                    align: Align::new(1).unwrap(),
                };
                let alloc = self.allocator.alloc_bytes(layout)?;
                this.buf.insert(alloc)
            }
        };

        if layout.size > this.chunk_size - this.next {
            // TODO: is there a more efficient way?
            let mut chunk_size = this.chunk_size * 2;
            while layout.size > chunk_size - this.next {
                chunk_size *= 2;
            }

            let layout = Layout {
                size: chunk_size,
                align: Align::new(1).unwrap(),
            };
            *buf = self.allocator.alloc_bytes(layout)?;

            // If we update `this.chunk_size` before allocating, and allocating
            // fails, then the next attempted allocation could mistakenly
            // believe we have `chunk_size` memory, which we do not have. This
            // would cause very bad things to happen.
            this.chunk_size = chunk_size;
        }

        let align = layout.align.align();
        // SAFETY: no isizes are overflowed, nor bounds overreached
        let start = unsafe { buf.as_ptr().add(this.next) };
        let start = start.map_addr(|addr| addr.div_ceil(align) * align);
        let end = start as usize + layout.size;

        debug_assert!(start as usize % align == 0);
        debug_assert!(end <= buf.as_ptr() as usize + this.chunk_size);

        this.next = end - buf.as_ptr() as usize;
        this.count += 1;

        drop(this_lock);

        let ptr = NonNull::new(start).unwrap();
        // SAFETY:
        // - `ptr` was just allocated in this allocator, and thus not already
        //   deallocated.
        // - `layout` is the layout used to allocate `ptr`.
        // - This is not being used to double-drop an allocation.
        Ok(unsafe { Alloc::new(ptr, layout, self) })
    }

    fn dealloc_bytes<'this>(&'this self, alloc: Alloc<'this, u8>) {
        if alloc.allocator().brand() != &self.brand {
            return;
        }

        let mut this = self.inner.lock();

        this.count -= 1;
        if this.count == 0 {
            this.next = 0;
        }
    }

    fn dangling(&self, align: Align) -> Alloc<'_, u8> {
        let layout = Layout { size: 0, align };
        unsafe {
            Alloc::new(
                NonNull::new(align.align() as *mut u8).unwrap(),
                layout,
                self,
            )
        }
    }

    fn total_size(&self) -> Option<usize> {
        self.allocator.total_size()
    }

    fn remaining_size(&self) -> Option<usize> {
        self.allocator.remaining_size()
    }
}