1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
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;

#[derive(Debug)]
pub struct VramAllocatorInUseError {}

static mut VRAM_ALLOCATOR: VramAllocatorSingleton = VramAllocatorSingleton {
    alloc: Some(VramAllocator::new()),
};

pub fn get_vram_allocator() -> Result<VramAllocator, VramAllocatorInUseError> {
    let opt_alloc = unsafe { VRAM_ALLOCATOR.get_vram_alloc() };
    opt_alloc.ok_or(VramAllocatorInUseError {})
}

pub struct VramAllocatorSingleton {
    alloc: Option<VramAllocator>,
}

impl VramAllocatorSingleton {
    pub fn get_vram_alloc(&mut self) -> Option<VramAllocator> {
        self.alloc.take()
    }
}

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<'_> {
    fn new(start: u32, len: u32) -> Self {
        Self {
            start,
            len,
            vram: PhantomData,
        }
    }

    pub fn as_mut_ptr_from_zero(&self) -> *mut u8 {
        unsafe { vram_start_addr_zero().add(self.start as usize) }
    }

    pub fn as_mut_ptr_direct_to_vram(&self) -> *mut u8 {
        unsafe { vram_start_addr_direct().add(self.start as usize) }
    }

    pub fn len(&self) -> u32 {
        self.len
    }

    pub fn is_empty(&self) -> bool {
        self.len == 0
    }
}

// A dead-simple VRAM bump allocator.
// TODO: pin?
#[derive(Debug)]
pub struct SimpleVramAllocator {
    offset: AtomicU32,
}

impl SimpleVramAllocator {
    const fn new() -> Self {
        Self {
            offset: AtomicU32::new(0),
        }
    }

    /// 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);
    }

    // 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(&self, size: u32) -> VramMemChunk<'_> {
        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!");
        }

        VramMemChunk::new(old_offset, size)
    }

    // TODO: ensure 16-bit alignment?
    pub fn alloc_sized<T: Sized>(&self, count: u32) -> VramMemChunk<'_> {
        let size = size_of::<T>() as u32;
        self.alloc(count * size)
    }

    pub fn alloc_texture_pixels(
        &self,
        width: u32,
        height: u32,
        psm: TexturePixelFormat,
    ) -> VramMemChunk<'_> {
        let size = get_memory_size(width, height, psm);
        self.alloc(size)
    }

    // TODO: write, or write_volatile?
    // TODO: result instead of unwrap?
    // TODO: Keep track of the allocated chunk
    // TODO: determine unsafety of this
    pub unsafe fn move_to_vram<T: Sized>(&mut self, obj: T) -> &mut T {
        let chunk = self.alloc_sized::<T>(1);
        let ptr = chunk.as_mut_ptr_direct_to_vram() as *mut T;
        ptr.write(obj);
        ptr.as_mut().unwrap()
    }

    fn total_mem(&self) -> u32 {
        total_vram_size()
    }
}

fn total_vram_size() -> u32 {
    unsafe { sceGeEdramGetSize() }
}

// NOTE: VRAM actually starts at 0x4000000, as returned by sceGeEdramGetAddr.
//       The Gu functions take that into account, and start their pointer
//       indices at 0. See GE_EDRAM_ADDRESS in gu.rs for that offset being used.
fn vram_start_addr_zero() -> *mut u8 {
    null_mut()
}

fn vram_start_addr_direct() -> *mut u8 {
    unsafe { sceGeEdramGetAddr() }
}

fn get_memory_size(width: u32, height: u32, psm: TexturePixelFormat) -> u32 {
    match psm {
        TexturePixelFormat::PsmT4 => (width * height) >> 1,
        TexturePixelFormat::PsmT8 => width * height,

        TexturePixelFormat::Psm5650
        | TexturePixelFormat::Psm5551
        | TexturePixelFormat::Psm4444
        | TexturePixelFormat::PsmT16 => 2 * width * height,

        TexturePixelFormat::Psm8888 | TexturePixelFormat::PsmT32 => 4 * width * height,

        _ => unimplemented!(),
    }
}