pbrt/src/utils/arena.rs

use crate::utils::backend::GpuAllocator;
use crate::utils::mipmap::MIPMap;
use parking_lot::Mutex;
use shared::utils::soa::SoAAllocator;
use shared::Ptr;
use std::alloc::Layout;
use std::collections::HashMap;
use std::panic::Location;
use std::sync::Arc;

struct Chunk {
    ptr: *mut u8,
    layout: Layout,
}

struct GpuBump<A: GpuAllocator> {
    allocator: A,
    current: *mut u8,
    end: *mut u8,
    chunks: Vec<Chunk>,
}

const CHUNK_SIZE: usize = 256 * 1024;

impl<A: GpuAllocator> GpuBump<A> {
    fn new(allocator: A) -> Self {
        Self {
            allocator,
            current: std::ptr::null_mut(),
            end: std::ptr::null_mut(),
            chunks: Vec::new(),
        }
    }

    fn alloc<T>(&mut self, value: T) -> *mut T {
        let layout = Layout::new::<T>();
        let ptr = self.alloc_layout(layout) as *mut T;
        unsafe { ptr.write(value) };
        ptr
    }

    fn alloc_slice<T: Copy>(&mut self, values: &[T]) -> (*mut T, usize) {
        if values.is_empty() {
            return (std::ptr::null_mut(), 0);
        }
        let layout = Layout::array::<T>(values.len()).unwrap();
        let ptr = self.alloc_layout(layout) as *mut T;
        unsafe {
            std::ptr::copy_nonoverlapping(values.as_ptr(), ptr, values.len());
        }
        (ptr, values.len())
    }

    fn alloc_layout(&mut self, layout: Layout) -> *mut u8 {
        let size = layout.size();
        let align = layout.align();

        if size == 0 {
            return align as *mut u8;
        }

        // Fast path: bump from current chunk
        let start = self.current as usize;
        let aligned = (start + align - 1) & !(align - 1);
        let end = aligned + size;

        if end <= self.end as usize {
            self.current = end as *mut u8;
            return aligned as *mut u8;
        }

        let chunk_size = if size > CHUNK_SIZE {
            size.next_multiple_of(align.max(16))
        } else {
            CHUNK_SIZE.max(size.next_multiple_of(align.max(16)))
        };

        let chunk_layout = Layout::from_size_align(chunk_size, align.max(16)).unwrap();
        let chunk = unsafe { self.allocator.alloc(chunk_layout) };

        let caller = Location::caller();
        if chunk.is_null() {
            panic!(
                "GpuBump OOM {} {}: chunk_size={} align={} backend={}",
                caller.file(),
                caller.line(),
                chunk_size,
                chunk_layout.align(),
                std::any::type_name::<A>()
            );
        }

        self.chunks.push(Chunk {
            ptr: chunk,
            layout: chunk_layout,
        });

        // If this object alone fills the chunk, mark it consumed and return.
        if size > CHUNK_SIZE {
            self.current = unsafe { chunk.add(chunk_size) };
            self.end = self.current;
            return chunk;
        }

        // Set up bump pointers inside the new chunk.
        self.current = chunk;
        self.end = unsafe { chunk.add(chunk_size) };

        let aligned = (self.current as usize + align - 1) & !(align - 1);
        self.current = (aligned + size) as *mut u8;
        aligned as *mut u8
    }
}

impl<A: GpuAllocator> Drop for GpuBump<A> {
    fn drop(&mut self) {
        for chunk in self.chunks.drain(..) {
            unsafe { self.allocator.dealloc(chunk.ptr, chunk.layout) };
        }
    }
}

unsafe impl<A: GpuAllocator> Send for GpuBump<A> {}
unsafe impl<A: GpuAllocator> Sync for GpuBump<A> {}

pub struct Arena<A: GpuAllocator> {
    bump: Mutex<GpuBump<A>>,
    texture_cache: Mutex<HashMap<usize, u64>>,
}

impl<A: GpuAllocator> Arena<A> {
    pub fn new(allocator: A) -> Self {
        Self {
            bump: Mutex::new(GpuBump::new(allocator)),
            texture_cache: Mutex::new(HashMap::new()),
        }
    }

    pub fn alloc<T>(&self, value: T) -> Ptr<T> {
        let mut bump = self.bump.lock();
        let ptr = bump.alloc(value);
        Ptr::from_raw(ptr)
    }

    pub fn alloc_opt<T>(&self, value: Option<T>) -> Ptr<T> {
        match value {
            Some(v) => self.alloc(v),
            None => Ptr::null(),
        }
    }

    pub fn alloc_arc<T: Clone>(&self, value: Arc<T>) -> Ptr<T> {
        match Arc::try_unwrap(value) {
            Ok(inner) => self.alloc(inner),
            Err(arc) => self.alloc((*arc).clone()),
        }
    }

    pub fn alloc_slice<T: Copy>(&self, values: &[T]) -> (Ptr<T>, usize) {
        let mut bump = self.bump.lock();
        let (ptr, len) = bump.alloc_slice(values);
        (Ptr::from_raw(ptr), len)
    }

    pub fn alloc_layout(&self, layout: Layout) -> *mut u8 {
        let mut bump = self.bump.lock();
        bump.alloc_layout(layout)
    }

    pub fn get_texture_object(&self, mipmap: &Arc<MIPMap>) -> u64 {
        let key = Arc::as_ptr(mipmap) as usize;
        let mut cache = self.texture_cache.lock();
        if let Some(&tex_obj) = cache.get(&key) {
            return tex_obj;
        }
        let tex_obj = 0u64; // TODO: backend-specific creation
        cache.insert(key, tex_obj);
        tex_obj
    }
}

impl<A: GpuAllocator + Default> Default for Arena<A> {
    fn default() -> Self {
        Self::new(A::default())
    }
}

impl<A: GpuAllocator> SoAAllocator for Arena<A> {
    fn alloc_raw(&self, layout: core::alloc::Layout) -> *mut u8 {
        self.alloc_layout(layout)
    }
}