1use crate::geometry::{
  2    rect::RectI,
  3    vector::{vec2i, Vector2I},
  4};
  5use etagere::BucketedAtlasAllocator;
  6use foreign_types::ForeignType;
  7use metal::{Device, TextureDescriptor};
  8use objc::{msg_send, sel, sel_impl};
  9
 10pub struct AtlasAllocator {
 11    device: Device,
 12    texture_descriptor: TextureDescriptor,
 13    atlases: Vec<Atlas>,
 14    free_atlases: Vec<Atlas>,
 15}
 16
 17#[derive(Copy, Clone)]
 18pub struct AllocId {
 19    pub atlas_id: usize,
 20    alloc_id: etagere::AllocId,
 21}
 22
 23impl AtlasAllocator {
 24    pub fn new(device: Device, texture_descriptor: TextureDescriptor) -> Self {
 25        let mut me = Self {
 26            device,
 27            texture_descriptor,
 28            atlases: Vec::new(),
 29            free_atlases: Vec::new(),
 30        };
 31        let atlas = me.new_atlas(Vector2I::zero());
 32        me.atlases.push(atlas);
 33        me
 34    }
 35
 36    pub fn default_atlas_size(&self) -> Vector2I {
 37        vec2i(
 38            self.texture_descriptor.width() as i32,
 39            self.texture_descriptor.height() as i32,
 40        )
 41    }
 42
 43    pub fn allocate(&mut self, requested_size: Vector2I) -> (AllocId, Vector2I) {
 44        let (alloc_id, origin) = self
 45            .atlases
 46            .last_mut()
 47            .unwrap()
 48            .allocate(requested_size)
 49            .unwrap_or_else(|| {
 50                let mut atlas = self.new_atlas(requested_size);
 51                let (id, origin) = atlas.allocate(requested_size).unwrap();
 52                self.atlases.push(atlas);
 53                (id, origin)
 54            });
 55
 56        let id = AllocId {
 57            atlas_id: self.atlases.len() - 1,
 58            alloc_id,
 59        };
 60        (id, origin)
 61    }
 62
 63    pub fn upload(&mut self, size: Vector2I, bytes: &[u8]) -> (AllocId, RectI) {
 64        let (alloc_id, origin) = self.allocate(size);
 65        let bounds = RectI::new(origin, size);
 66        self.atlases[alloc_id.atlas_id].upload(bounds, bytes);
 67        (alloc_id, bounds)
 68    }
 69
 70    pub fn deallocate(&mut self, id: AllocId) {
 71        if let Some(atlas) = self.atlases.get_mut(id.atlas_id) {
 72            atlas.deallocate(id.alloc_id);
 73            if atlas.is_empty() {
 74                self.free_atlases.push(self.atlases.remove(id.atlas_id));
 75            }
 76        }
 77    }
 78
 79    pub fn clear(&mut self) {
 80        for atlas in &mut self.atlases {
 81            atlas.clear();
 82        }
 83        self.free_atlases.extend(self.atlases.drain(1..));
 84    }
 85
 86    pub fn texture(&self, atlas_id: usize) -> Option<&metal::TextureRef> {
 87        self.atlases.get(atlas_id).map(|a| a.texture.as_ref())
 88    }
 89
 90    fn new_atlas(&mut self, required_size: Vector2I) -> Atlas {
 91        if let Some(i) = self.free_atlases.iter().rposition(|atlas| {
 92            atlas.size().x() >= required_size.x() && atlas.size().y() >= required_size.y()
 93        }) {
 94            self.free_atlases.remove(i)
 95        } else {
 96            let size = self.default_atlas_size().max(required_size);
 97            let texture = if size.x() as u64 > self.texture_descriptor.width()
 98                || size.y() as u64 > self.texture_descriptor.height()
 99            {
100                let descriptor = unsafe {
101                    let descriptor_ptr: *mut metal::MTLTextureDescriptor =
102                        msg_send![self.texture_descriptor, copy];
103                    metal::TextureDescriptor::from_ptr(descriptor_ptr)
104                };
105                descriptor.set_width(size.x() as u64);
106                descriptor.set_height(size.y() as u64);
107                self.device.new_texture(&descriptor)
108            } else {
109                self.device.new_texture(&self.texture_descriptor)
110            };
111            Atlas::new(size, texture)
112        }
113    }
114}
115
116struct Atlas {
117    allocator: BucketedAtlasAllocator,
118    texture: metal::Texture,
119}
120
121impl Atlas {
122    fn new(size: Vector2I, texture: metal::Texture) -> Self {
123        Self {
124            allocator: BucketedAtlasAllocator::new(etagere::Size::new(size.x(), size.y())),
125            texture,
126        }
127    }
128
129    fn size(&self) -> Vector2I {
130        let size = self.allocator.size();
131        vec2i(size.width, size.height)
132    }
133
134    fn allocate(&mut self, size: Vector2I) -> Option<(etagere::AllocId, Vector2I)> {
135        let alloc = self
136            .allocator
137            .allocate(etagere::Size::new(size.x(), size.y()))?;
138        let origin = alloc.rectangle.min;
139        Some((alloc.id, vec2i(origin.x, origin.y)))
140    }
141
142    fn upload(&mut self, bounds: RectI, bytes: &[u8]) {
143        let region = metal::MTLRegion::new_2d(
144            bounds.origin().x() as u64,
145            bounds.origin().y() as u64,
146            bounds.size().x() as u64,
147            bounds.size().y() as u64,
148        );
149        self.texture.replace_region(
150            region,
151            0,
152            bytes.as_ptr() as *const _,
153            (bounds.size().x() * self.bytes_per_pixel() as i32) as u64,
154        );
155    }
156
157    fn bytes_per_pixel(&self) -> u8 {
158        use metal::MTLPixelFormat::*;
159        match self.texture.pixel_format() {
160            A8Unorm | R8Unorm => 1,
161            RGBA8Unorm | BGRA8Unorm => 4,
162            _ => unimplemented!(),
163        }
164    }
165
166    fn deallocate(&mut self, id: etagere::AllocId) {
167        self.allocator.deallocate(id);
168    }
169
170    fn is_empty(&self) -> bool {
171        self.allocator.is_empty()
172    }
173
174    fn clear(&mut self) {
175        self.allocator.clear();
176    }
177}