From 0536594241d9e5219b20d4e79a2f4d189b5e13a6 Mon Sep 17 00:00:00 2001
From: Ayrton <a.munoz3327@gmail.com>
Date: Sun, 5 Jun 2022 21:34:10 -0400
Subject: [PATCH] Add SimpleVramAllocator::free_all and attach a lifetime to
 VRAM chunks

---
 examples/cube/src/main.rs | 15 ++++++-----
 psp/src/vram_alloc.rs     | 53 ++++++++++++++++++++++++++++-----------
 2 files changed, 48 insertions(+), 20 deletions(-)

diff --git a/examples/cube/src/main.rs b/examples/cube/src/main.rs
index ab7f3970..81b46df2 100644
--- a/examples/cube/src/main.rs
+++ b/examples/cube/src/main.rs
@@ -89,18 +89,21 @@ unsafe fn psp_main_inner() {
     psp::enable_home_button();
 
     let mut allocator = get_vram_allocator().unwrap();
-    let fbp0 = allocator.alloc_texture_pixels(BUF_WIDTH, SCREEN_HEIGHT, TexturePixelFormat::Psm8888).as_mut_ptr_from_zero();
-    let fbp1 = allocator.alloc_texture_pixels(BUF_WIDTH, SCREEN_HEIGHT, TexturePixelFormat::Psm8888).as_mut_ptr_from_zero();
-    let zbp = allocator.alloc_texture_pixels(BUF_WIDTH, SCREEN_HEIGHT, TexturePixelFormat::Psm4444).as_mut_ptr_from_zero();
+    let fbp0 = allocator.alloc_texture_pixels(BUF_WIDTH, SCREEN_HEIGHT, TexturePixelFormat::Psm8888);
+    let fbp1 = allocator.alloc_texture_pixels(BUF_WIDTH, SCREEN_HEIGHT, TexturePixelFormat::Psm8888);
+    let zbp = allocator.alloc_texture_pixels(BUF_WIDTH, SCREEN_HEIGHT, TexturePixelFormat::Psm4444);
+    // Attempting to free the three VRAM chunks at this point would give a
+    // compile-time error since fbp0, fbp1 and zbp are used later on
+    //allocator.free_all();
 
     sys::sceGumLoadIdentity();
 
     sys::sceGuInit();
 
     sys::sceGuStart(GuContextType::Direct, &mut LIST.0 as *mut [u32; 0x40000] as *mut _);
-    sys::sceGuDrawBuffer(DisplayPixelFormat::Psm8888, fbp0 as _, BUF_WIDTH as i32);
-    sys::sceGuDispBuffer(SCREEN_WIDTH as i32, SCREEN_HEIGHT as i32, fbp1 as _, BUF_WIDTH as i32);
-    sys::sceGuDepthBuffer(zbp as _, BUF_WIDTH as i32);
+    sys::sceGuDrawBuffer(DisplayPixelFormat::Psm8888, fbp0.as_mut_ptr_from_zero() as _, BUF_WIDTH as i32);
+    sys::sceGuDispBuffer(SCREEN_WIDTH as i32, SCREEN_HEIGHT as i32, fbp1.as_mut_ptr_from_zero() as _, BUF_WIDTH as i32);
+    sys::sceGuDepthBuffer(zbp.as_mut_ptr_from_zero() as _, BUF_WIDTH as i32);
     sys::sceGuOffset(2048 - (SCREEN_WIDTH / 2), 2048 - (SCREEN_HEIGHT / 2));
     sys::sceGuViewport(2048, 2048, SCREEN_WIDTH as i32, SCREEN_HEIGHT as i32);
     sys::sceGuDepthRange(65535, 0);
diff --git a/psp/src/vram_alloc.rs b/psp/src/vram_alloc.rs
index b872cb36..f343bdcb 100644
--- a/psp/src/vram_alloc.rs
+++ b/psp/src/vram_alloc.rs
@@ -1,7 +1,9 @@
 use crate::sys::TexturePixelFormat;
 use crate::sys::{sceGeEdramGetAddr, sceGeEdramGetSize};
+use core::marker::PhantomData;
 use core::mem::size_of;
 use core::ptr::null_mut;
+use core::sync::atomic::{AtomicU32, Ordering};
 
 type VramAllocator = SimpleVramAllocator;
 
@@ -27,14 +29,20 @@ impl VramAllocatorSingleton {
     }
 }
 
-pub struct VramMemChunk {
+pub struct VramMemChunk<'a> {
     start: u32,
     len: u32,
+    // Needed since VramMemChunk has a lifetime, but doesn't contain references
+    vram: PhantomData<&'a mut ()>,
 }
 
-impl VramMemChunk {
+impl VramMemChunk<'_> {
     fn new(start: u32, len: u32) -> Self {
-        Self { start, len }
+        Self {
+            start,
+            len,
+            vram: PhantomData,
+        }
     }
 
     pub fn as_mut_ptr_from_zero(&self) -> *mut u8 {
@@ -54,20 +62,37 @@ impl VramMemChunk {
 // TODO: pin?
 #[derive(Debug)]
 pub struct SimpleVramAllocator {
-    offset: u32,
+    offset: AtomicU32,
 }
 
 impl SimpleVramAllocator {
     const fn new() -> Self {
-        Self { offset: 0 }
+        Self {
+            offset: AtomicU32::new(0),
+        }
     }
 
-    // TODO: return a Result instead of panicking
-    pub fn alloc(&mut self, size: u32) -> VramMemChunk {
-        let old_offset = self.offset;
-        self.offset += size;
+    /// Frees all previously allocated VRAM chunks.
+    ///
+    /// This resets the allocator's counter, but does not change the contents of
+    /// VRAM. Since this method requires `&mut Self`, it cannot overlap with any
+    /// previously allocated `VramMemChunk`s since they have the lifetime of the
+    /// `&Self` that allocated them.
+    pub fn free_all(&mut self) {
+        self.offset.store(0, Ordering::Relaxed);
+    }
 
-        if self.offset > self.total_mem() {
+    // TODO: return a Result instead of panicking
+    /// Allocates `size` bytes of VRAM
+    ///
+    /// The returned VRAM chunk has the same lifetime as the
+    /// `SimpleVramAllocator` borrow (i.e. `&self`) that allocated it.
+    pub fn alloc<'a>(&'a self, size: u32) -> VramMemChunk<'a> {
+        let old_offset = self.offset.load(Ordering::Relaxed);
+        let new_offset = old_offset + size;
+        self.offset.store(new_offset, Ordering::Relaxed);
+
+        if new_offset > self.total_mem() {
             panic!("Total VRAM size exceeded!");
         }
 
@@ -75,17 +100,17 @@ impl SimpleVramAllocator {
     }
 
     // TODO: ensure 16-bit alignment?
-    pub fn alloc_sized<T: Sized>(&mut self, count: u32) -> VramMemChunk {
+    pub fn alloc_sized<'a, T: Sized>(&'a self, count: u32) -> VramMemChunk<'a> {
         let size = size_of::<T>() as u32;
         self.alloc(count * size)
     }
 
-    pub fn alloc_texture_pixels(
-        &mut self,
+    pub fn alloc_texture_pixels<'a>(
+        &'a self,
         width: u32,
         height: u32,
         psm: TexturePixelFormat,
-    ) -> VramMemChunk {
+    ) -> VramMemChunk<'a> {
         let size = get_memory_size(width, height, psm);
         self.alloc(size)
     }