directx_atlas.rs

  1use collections::FxHashMap;
  2use etagere::BucketedAtlasAllocator;
  3use parking_lot::Mutex;
  4use windows::Win32::Graphics::{
  5    Direct3D11::{
  6        D3D11_BIND_SHADER_RESOURCE, D3D11_BOX, D3D11_TEXTURE2D_DESC, D3D11_USAGE_DEFAULT,
  7        ID3D11Device, ID3D11DeviceContext, ID3D11ShaderResourceView, ID3D11Texture2D,
  8    },
  9    Dxgi::Common::*,
 10};
 11
 12use crate::{
 13    AtlasKey, AtlasTextureId, AtlasTextureKind, AtlasTile, Bounds, DevicePixels, PlatformAtlas,
 14    Point, Size, platform::AtlasTextureList,
 15};
 16
 17pub(crate) struct DirectXAtlas(Mutex<DirectXAtlasState>);
 18
 19struct DirectXAtlasState {
 20    device: ID3D11Device,
 21    device_context: ID3D11DeviceContext,
 22    monochrome_textures: AtlasTextureList<DirectXAtlasTexture>,
 23    polychrome_textures: AtlasTextureList<DirectXAtlasTexture>,
 24    subpixel_textures: AtlasTextureList<DirectXAtlasTexture>,
 25    tiles_by_key: FxHashMap<AtlasKey, AtlasTile>,
 26}
 27
 28struct DirectXAtlasTexture {
 29    id: AtlasTextureId,
 30    bytes_per_pixel: u32,
 31    allocator: BucketedAtlasAllocator,
 32    texture: ID3D11Texture2D,
 33    view: [Option<ID3D11ShaderResourceView>; 1],
 34    live_atlas_keys: u32,
 35}
 36
 37impl DirectXAtlas {
 38    pub(crate) fn new(device: &ID3D11Device, device_context: &ID3D11DeviceContext) -> Self {
 39        DirectXAtlas(Mutex::new(DirectXAtlasState {
 40            device: device.clone(),
 41            device_context: device_context.clone(),
 42            monochrome_textures: Default::default(),
 43            polychrome_textures: Default::default(),
 44            subpixel_textures: Default::default(),
 45            tiles_by_key: Default::default(),
 46        }))
 47    }
 48
 49    pub(crate) fn get_texture_view(
 50        &self,
 51        id: AtlasTextureId,
 52    ) -> [Option<ID3D11ShaderResourceView>; 1] {
 53        let lock = self.0.lock();
 54        let tex = lock.texture(id);
 55        tex.view.clone()
 56    }
 57
 58    pub(crate) fn handle_device_lost(
 59        &self,
 60        device: &ID3D11Device,
 61        device_context: &ID3D11DeviceContext,
 62    ) {
 63        let mut lock = self.0.lock();
 64        lock.device = device.clone();
 65        lock.device_context = device_context.clone();
 66        lock.monochrome_textures = AtlasTextureList::default();
 67        lock.polychrome_textures = AtlasTextureList::default();
 68        lock.subpixel_textures = AtlasTextureList::default();
 69        lock.tiles_by_key.clear();
 70    }
 71}
 72
 73impl PlatformAtlas for DirectXAtlas {
 74    fn get_or_insert_with<'a>(
 75        &self,
 76        key: &AtlasKey,
 77        build: &mut dyn FnMut() -> anyhow::Result<
 78            Option<(Size<DevicePixels>, std::borrow::Cow<'a, [u8]>)>,
 79        >,
 80    ) -> anyhow::Result<Option<AtlasTile>> {
 81        let mut lock = self.0.lock();
 82        if let Some(tile) = lock.tiles_by_key.get(key) {
 83            Ok(Some(tile.clone()))
 84        } else {
 85            let Some((size, bytes)) = build()? else {
 86                return Ok(None);
 87            };
 88            let tile = lock
 89                .allocate(size, key.texture_kind())
 90                .ok_or_else(|| anyhow::anyhow!("failed to allocate"))?;
 91            let texture = lock.texture(tile.texture_id);
 92            texture.upload(&lock.device_context, tile.bounds, &bytes);
 93            lock.tiles_by_key.insert(key.clone(), tile.clone());
 94            Ok(Some(tile))
 95        }
 96    }
 97
 98    fn remove(&self, key: &AtlasKey) {
 99        let mut lock = self.0.lock();
100
101        let Some(id) = lock.tiles_by_key.remove(key).map(|tile| tile.texture_id) else {
102            return;
103        };
104
105        let textures = match id.kind {
106            AtlasTextureKind::Monochrome => &mut lock.monochrome_textures,
107            AtlasTextureKind::Polychrome => &mut lock.polychrome_textures,
108            AtlasTextureKind::Subpixel => &mut lock.subpixel_textures,
109        };
110
111        let Some(texture_slot) = textures.textures.get_mut(id.index as usize) else {
112            return;
113        };
114
115        if let Some(mut texture) = texture_slot.take() {
116            texture.decrement_ref_count();
117            if texture.is_unreferenced() {
118                textures.free_list.push(texture.id.index as usize);
119                lock.tiles_by_key.remove(key);
120            } else {
121                *texture_slot = Some(texture);
122            }
123        }
124    }
125}
126
127impl DirectXAtlasState {
128    fn allocate(
129        &mut self,
130        size: Size<DevicePixels>,
131        texture_kind: AtlasTextureKind,
132    ) -> Option<AtlasTile> {
133        {
134            let textures = match texture_kind {
135                AtlasTextureKind::Monochrome => &mut self.monochrome_textures,
136                AtlasTextureKind::Polychrome => &mut self.polychrome_textures,
137                AtlasTextureKind::Subpixel => &mut self.subpixel_textures,
138            };
139
140            if let Some(tile) = textures
141                .iter_mut()
142                .rev()
143                .find_map(|texture| texture.allocate(size))
144            {
145                return Some(tile);
146            }
147        }
148
149        let texture = self.push_texture(size, texture_kind)?;
150        texture.allocate(size)
151    }
152
153    fn push_texture(
154        &mut self,
155        min_size: Size<DevicePixels>,
156        kind: AtlasTextureKind,
157    ) -> Option<&mut DirectXAtlasTexture> {
158        const DEFAULT_ATLAS_SIZE: Size<DevicePixels> = Size {
159            width: DevicePixels(1024),
160            height: DevicePixels(1024),
161        };
162        // Max texture size for DirectX. See:
163        // https://learn.microsoft.com/en-us/windows/win32/direct3d11/overviews-direct3d-11-resources-limits
164        const MAX_ATLAS_SIZE: Size<DevicePixels> = Size {
165            width: DevicePixels(16384),
166            height: DevicePixels(16384),
167        };
168        let size = min_size.min(&MAX_ATLAS_SIZE).max(&DEFAULT_ATLAS_SIZE);
169        let pixel_format;
170        let bind_flag;
171        let bytes_per_pixel;
172        match kind {
173            AtlasTextureKind::Monochrome => {
174                pixel_format = DXGI_FORMAT_R8_UNORM;
175                bind_flag = D3D11_BIND_SHADER_RESOURCE;
176                bytes_per_pixel = 1;
177            }
178            AtlasTextureKind::Polychrome => {
179                pixel_format = DXGI_FORMAT_B8G8R8A8_UNORM;
180                bind_flag = D3D11_BIND_SHADER_RESOURCE;
181                bytes_per_pixel = 4;
182            }
183            AtlasTextureKind::Subpixel => {
184                pixel_format = DXGI_FORMAT_R8G8B8A8_UNORM;
185                bind_flag = D3D11_BIND_SHADER_RESOURCE;
186                bytes_per_pixel = 4;
187            }
188        }
189        let texture_desc = D3D11_TEXTURE2D_DESC {
190            Width: size.width.0 as u32,
191            Height: size.height.0 as u32,
192            MipLevels: 1,
193            ArraySize: 1,
194            Format: pixel_format,
195            SampleDesc: DXGI_SAMPLE_DESC {
196                Count: 1,
197                Quality: 0,
198            },
199            Usage: D3D11_USAGE_DEFAULT,
200            BindFlags: bind_flag.0 as u32,
201            CPUAccessFlags: 0,
202            MiscFlags: 0,
203        };
204        let mut texture: Option<ID3D11Texture2D> = None;
205        unsafe {
206            // This only returns None if the device is lost, which we will recreate later.
207            // So it's ok to return None here.
208            self.device
209                .CreateTexture2D(&texture_desc, None, Some(&mut texture))
210                .ok()?;
211        }
212        let texture = texture.unwrap();
213
214        let texture_list = match kind {
215            AtlasTextureKind::Monochrome => &mut self.monochrome_textures,
216            AtlasTextureKind::Polychrome => &mut self.polychrome_textures,
217            AtlasTextureKind::Subpixel => &mut self.subpixel_textures,
218        };
219        let index = texture_list.free_list.pop();
220        let view = unsafe {
221            let mut view = None;
222            self.device
223                .CreateShaderResourceView(&texture, None, Some(&mut view))
224                .ok()?;
225            [view]
226        };
227        let atlas_texture = DirectXAtlasTexture {
228            id: AtlasTextureId {
229                index: index.unwrap_or(texture_list.textures.len()) as u32,
230                kind,
231            },
232            bytes_per_pixel,
233            allocator: etagere::BucketedAtlasAllocator::new(size.into()),
234            texture,
235            view,
236            live_atlas_keys: 0,
237        };
238        if let Some(ix) = index {
239            texture_list.textures[ix] = Some(atlas_texture);
240            texture_list.textures.get_mut(ix).unwrap().as_mut()
241        } else {
242            texture_list.textures.push(Some(atlas_texture));
243            texture_list.textures.last_mut().unwrap().as_mut()
244        }
245    }
246
247    fn texture(&self, id: AtlasTextureId) -> &DirectXAtlasTexture {
248        match id.kind {
249            crate::AtlasTextureKind::Monochrome => &self.monochrome_textures[id.index as usize]
250                .as_ref()
251                .unwrap(),
252            crate::AtlasTextureKind::Polychrome => &self.polychrome_textures[id.index as usize]
253                .as_ref()
254                .unwrap(),
255            crate::AtlasTextureKind::Subpixel => {
256                &self.subpixel_textures[id.index as usize].as_ref().unwrap()
257            }
258        }
259    }
260}
261
262impl DirectXAtlasTexture {
263    fn allocate(&mut self, size: Size<DevicePixels>) -> Option<AtlasTile> {
264        let allocation = self.allocator.allocate(size.into())?;
265        let tile = AtlasTile {
266            texture_id: self.id,
267            tile_id: allocation.id.into(),
268            bounds: Bounds {
269                origin: allocation.rectangle.min.into(),
270                size,
271            },
272            padding: 0,
273        };
274        self.live_atlas_keys += 1;
275        Some(tile)
276    }
277
278    fn upload(
279        &self,
280        device_context: &ID3D11DeviceContext,
281        bounds: Bounds<DevicePixels>,
282        bytes: &[u8],
283    ) {
284        unsafe {
285            device_context.UpdateSubresource(
286                &self.texture,
287                0,
288                Some(&D3D11_BOX {
289                    left: bounds.left().0 as u32,
290                    top: bounds.top().0 as u32,
291                    front: 0,
292                    right: bounds.right().0 as u32,
293                    bottom: bounds.bottom().0 as u32,
294                    back: 1,
295                }),
296                bytes.as_ptr() as _,
297                bounds.size.width.to_bytes(self.bytes_per_pixel as u8),
298                0,
299            );
300        }
301    }
302
303    fn decrement_ref_count(&mut self) {
304        self.live_atlas_keys -= 1;
305    }
306
307    fn is_unreferenced(&mut self) -> bool {
308        self.live_atlas_keys == 0
309    }
310}
311
312impl From<Size<DevicePixels>> for etagere::Size {
313    fn from(size: Size<DevicePixels>) -> Self {
314        etagere::Size::new(size.width.into(), size.height.into())
315    }
316}
317
318impl From<etagere::Point> for Point<DevicePixels> {
319    fn from(value: etagere::Point) -> Self {
320        Point {
321            x: DevicePixels::from(value.x),
322            y: DevicePixels::from(value.y),
323        }
324    }
325}