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