atlas.rs

  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    free_atlases: Vec<Atlas>,
 16}
 17
 18#[derive(Copy, Clone)]
 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 me = Self {
 27            device,
 28            texture_descriptor,
 29            atlases: Vec::new(),
 30            free_atlases: Vec::new(),
 31        };
 32        let atlas = me.new_atlas(Vector2I::zero());
 33        me.atlases.push(atlas);
 34        me
 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 allocation = self
 46            .atlases
 47            .last_mut()
 48            .unwrap()
 49            .allocate(requested_size)
 50            .or_else(|| {
 51                let mut atlas = self.new_atlas(requested_size);
 52                let (id, origin) = atlas.allocate(requested_size)?;
 53                self.atlases.push(atlas);
 54                Some((id, origin))
 55            });
 56
 57        if allocation.is_none() {
 58            warn!(
 59                "allocation of size {:?} could not be created",
 60                requested_size,
 61            );
 62        }
 63
 64        let (alloc_id, origin) = allocation?;
 65
 66        let id = AllocId {
 67            atlas_id: self.atlases.len() - 1,
 68            alloc_id,
 69        };
 70        Some((id, origin))
 71    }
 72
 73    pub fn upload(&mut self, size: Vector2I, bytes: &[u8]) -> Option<(AllocId, RectI)> {
 74        let (alloc_id, origin) = self.allocate(size)?;
 75        let bounds = RectI::new(origin, size);
 76        self.atlases[alloc_id.atlas_id].upload(bounds, bytes);
 77        Some((alloc_id, bounds))
 78    }
 79
 80    pub fn deallocate(&mut self, id: AllocId) {
 81        if let Some(atlas) = self.atlases.get_mut(id.atlas_id) {
 82            atlas.deallocate(id.alloc_id);
 83            if atlas.is_empty() {
 84                self.free_atlases.push(self.atlases.remove(id.atlas_id));
 85            }
 86        }
 87    }
 88
 89    pub fn clear(&mut self) {
 90        for atlas in &mut self.atlases {
 91            atlas.clear();
 92        }
 93        self.free_atlases.extend(self.atlases.drain(1..));
 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        if let Some(i) = self.free_atlases.iter().rposition(|atlas| {
102            atlas.size().x() >= required_size.x() && atlas.size().y() >= required_size.y()
103        }) {
104            self.free_atlases.remove(i)
105        } else {
106            let size = self.default_atlas_size().max(required_size);
107            let texture = if size.x() as u64 > self.texture_descriptor.width()
108                || size.y() as u64 > self.texture_descriptor.height()
109            {
110                let descriptor = unsafe {
111                    let descriptor_ptr: *mut metal::MTLTextureDescriptor =
112                        msg_send![self.texture_descriptor, copy];
113                    metal::TextureDescriptor::from_ptr(descriptor_ptr)
114                };
115                descriptor.set_width(size.x() as u64);
116                descriptor.set_height(size.y() as u64);
117                self.device.new_texture(&descriptor)
118            } else {
119                self.device.new_texture(&self.texture_descriptor)
120            };
121            Atlas::new(size, texture)
122        }
123    }
124}
125
126struct Atlas {
127    allocator: BucketedAtlasAllocator,
128    texture: metal::Texture,
129}
130
131impl Atlas {
132    fn new(size: Vector2I, texture: metal::Texture) -> Self {
133        Self {
134            allocator: BucketedAtlasAllocator::new(etagere::Size::new(size.x(), size.y())),
135            texture,
136        }
137    }
138
139    fn size(&self) -> Vector2I {
140        let size = self.allocator.size();
141        vec2i(size.width, size.height)
142    }
143
144    fn allocate(&mut self, size: Vector2I) -> Option<(etagere::AllocId, Vector2I)> {
145        let alloc = self
146            .allocator
147            .allocate(etagere::Size::new(size.x(), size.y()))?;
148        let origin = alloc.rectangle.min;
149        Some((alloc.id, vec2i(origin.x, origin.y)))
150    }
151
152    fn upload(&mut self, bounds: RectI, bytes: &[u8]) {
153        let region = metal::MTLRegion::new_2d(
154            bounds.origin().x() as u64,
155            bounds.origin().y() as u64,
156            bounds.size().x() as u64,
157            bounds.size().y() as u64,
158        );
159        self.texture.replace_region(
160            region,
161            0,
162            bytes.as_ptr() as *const _,
163            (bounds.size().x() * self.bytes_per_pixel() as i32) as u64,
164        );
165    }
166
167    fn bytes_per_pixel(&self) -> u8 {
168        use metal::MTLPixelFormat::*;
169        match self.texture.pixel_format() {
170            A8Unorm | R8Unorm => 1,
171            RGBA8Unorm | BGRA8Unorm => 4,
172            _ => unimplemented!(),
173        }
174    }
175
176    fn deallocate(&mut self, id: etagere::AllocId) {
177        self.allocator.deallocate(id);
178    }
179
180    fn is_empty(&self) -> bool {
181        self.allocator.is_empty()
182    }
183
184    fn clear(&mut self) {
185        self.allocator.clear();
186    }
187}