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}