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
113 self.device.new_texture(&descriptor)
114 } else {
115 self.device.new_texture(&self.texture_descriptor)
116 };
117 Atlas::new(size, texture)
118 }
119}
120
121struct Atlas {
122 allocator: BucketedAtlasAllocator,
123 texture: metal::Texture,
124}
125
126impl Atlas {
127 fn new(size: Vector2I, texture: metal::Texture) -> Self {
128 Self {
129 allocator: BucketedAtlasAllocator::new(etagere::Size::new(size.x(), size.y())),
130 texture,
131 }
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 clear(&mut self) {
171 self.allocator.clear();
172 }
173}