tenant.openagents/omega

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 gpui::{
13    AtlasKey, AtlasTextureId, AtlasTextureKind, AtlasTextureList, AtlasTile, Bounds, DevicePixels,
14    PlatformAtlas, Point, Size,
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))
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);
94            Ok(Some(tile))
95        }
96    }
97
98    fn remove(&self, key: &AtlasKey) {
99        let mut lock = self.0.lock();
100
101        let Some(tile) = lock.tiles_by_key.remove(key) else {
102            return;
103        };
104        let id = tile.texture_id;
105
106        let textures = match id.kind {
107            AtlasTextureKind::Monochrome => &mut lock.monochrome_textures,
108            AtlasTextureKind::Polychrome => &mut lock.polychrome_textures,
109            AtlasTextureKind::Subpixel => &mut lock.subpixel_textures,
110        };
111
112        let Some(texture_slot) = textures.textures.get_mut(id.index as usize) else {
113            return;
114        };
115
116        if let Some(mut texture) = texture_slot.take() {
117            texture.allocator.deallocate(tile.tile_id.into());
118            texture.decrement_ref_count();
119            if texture.is_unreferenced() {
120                textures.free_list.push(texture.id.index as usize);
121            } else {
122                *texture_slot = Some(texture);
123            }
124        }
125    }
126}
127
128impl DirectXAtlasState {
129    fn allocate(
130        &mut self,
131        size: Size<DevicePixels>,
132        texture_kind: AtlasTextureKind,
133    ) -> Option<AtlasTile> {
134        {
135            let textures = match texture_kind {
136                AtlasTextureKind::Monochrome => &mut self.monochrome_textures,
137                AtlasTextureKind::Polychrome => &mut self.polychrome_textures,
138                AtlasTextureKind::Subpixel => &mut self.subpixel_textures,
139            };
140
141            if let Some(tile) = textures
142                .iter_mut()
143                .rev()
144                .find_map(|texture| texture.allocate(size))
145            {
146                return Some(tile);
147            }
148        }
149
150        let texture = self.push_texture(size, texture_kind)?;
151        texture.allocate(size)
152    }
153
154    fn push_texture(
155        &mut self,
156        min_size: Size<DevicePixels>,
157        kind: AtlasTextureKind,
158    ) -> Option<&mut DirectXAtlasTexture> {
159        const DEFAULT_ATLAS_SIZE: Size<DevicePixels> = Size {
160            width: DevicePixels(1024),
161            height: DevicePixels(1024),
162        };
163        // Max texture size for DirectX. See:
164        // https://learn.microsoft.com/en-us/windows/win32/direct3d11/overviews-direct3d-11-resources-limits
165        const MAX_ATLAS_SIZE: Size<DevicePixels> = Size {
166            width: DevicePixels(16384),
167            height: DevicePixels(16384),
168        };
169        let size = min_size.min(&MAX_ATLAS_SIZE).max(&DEFAULT_ATLAS_SIZE);
170        let pixel_format;
171        let bind_flag;
172        let bytes_per_pixel;
173        match kind {
174            AtlasTextureKind::Monochrome => {
175                pixel_format = DXGI_FORMAT_R8_UNORM;
176                bind_flag = D3D11_BIND_SHADER_RESOURCE;
177                bytes_per_pixel = 1;
178            }
179            AtlasTextureKind::Polychrome => {
180                pixel_format = DXGI_FORMAT_B8G8R8A8_UNORM;
181                bind_flag = D3D11_BIND_SHADER_RESOURCE;
182                bytes_per_pixel = 4;
183            }
184            AtlasTextureKind::Subpixel => {
185                pixel_format = DXGI_FORMAT_R8G8B8A8_UNORM;
186                bind_flag = D3D11_BIND_SHADER_RESOURCE;
187                bytes_per_pixel = 4;
188            }
189        }
190        let texture_desc = D3D11_TEXTURE2D_DESC {
191            Width: size.width.0 as u32,
192            Height: size.height.0 as u32,
193            MipLevels: 1,
194            ArraySize: 1,
195            Format: pixel_format,
196            SampleDesc: DXGI_SAMPLE_DESC {
197                Count: 1,
198                Quality: 0,
199            },
200            Usage: D3D11_USAGE_DEFAULT,
201            BindFlags: bind_flag.0 as u32,
202            CPUAccessFlags: 0,
203            MiscFlags: 0,
204        };
205        let mut texture: Option<ID3D11Texture2D> = None;
206        unsafe {
207            // This only returns None if the device is lost, which we will recreate later.
208            // So it's ok to return None here.
209            self.device
210                .CreateTexture2D(&texture_desc, None, Some(&mut texture))
211                .ok()?;
212        }
213        let texture = texture.unwrap();
214
215        let texture_list = match kind {
216            AtlasTextureKind::Monochrome => &mut self.monochrome_textures,
217            AtlasTextureKind::Polychrome => &mut self.polychrome_textures,
218            AtlasTextureKind::Subpixel => &mut self.subpixel_textures,
219        };
220        let index = texture_list.free_list.pop();
221        let view = unsafe {
222            let mut view = None;
223            self.device
224                .CreateShaderResourceView(&texture, None, Some(&mut view))
225                .ok()?;
226            [view]
227        };
228        let atlas_texture = DirectXAtlasTexture {
229            id: AtlasTextureId {
230                index: index.unwrap_or(texture_list.textures.len()) as u32,
231                kind,
232            },
233            bytes_per_pixel,
234            allocator: etagere::BucketedAtlasAllocator::new(device_size_to_etagere(size)),
235            texture,
236            view,
237            live_atlas_keys: 0,
238        };
239        if let Some(ix) = index {
240            texture_list.textures[ix] = Some(atlas_texture);
241            texture_list.textures.get_mut(ix).unwrap().as_mut()
242        } else {
243            texture_list.textures.push(Some(atlas_texture));
244            texture_list.textures.last_mut().unwrap().as_mut()
245        }
246    }
247
248    fn texture(&self, id: AtlasTextureId) -> &DirectXAtlasTexture {
249        match id.kind {
250            AtlasTextureKind::Monochrome => &self.monochrome_textures[id.index as usize]
251                .as_ref()
252                .unwrap(),
253            AtlasTextureKind::Polychrome => &self.polychrome_textures[id.index as usize]
254                .as_ref()
255                .unwrap(),
256            AtlasTextureKind::Subpixel => {
257                &self.subpixel_textures[id.index as usize].as_ref().unwrap()
258            }
259        }
260    }
261}
262
263impl DirectXAtlasTexture {
264    fn allocate(&mut self, size: Size<DevicePixels>) -> Option<AtlasTile> {
265        let allocation = self.allocator.allocate(device_size_to_etagere(size))?;
266        let tile = AtlasTile {
267            texture_id: self.id,
268            tile_id: allocation.id.into(),
269            bounds: Bounds {
270                origin: etagere_point_to_device(allocation.rectangle.min),
271                size,
272            },
273            padding: 0,
274        };
275        self.live_atlas_keys += 1;
276        Some(tile)
277    }
278
279    fn upload(
280        &self,
281        device_context: &ID3D11DeviceContext,
282        bounds: Bounds<DevicePixels>,
283        bytes: &[u8],
284    ) {
285        // `UpdateSubresource` reads `row_pitch * height` bytes from `bytes` based on the
286        // `D3D11_BOX` below. If the caller hands us a slice shorter than that, the driver would
287        // over-read past the end of the source buffer (potentially by multiple megabytes), so bail
288        // out instead. This is a first-insert path rather than a per-frame one, so the check is
289        // effectively free.
290        let row_bytes = bounds.size.width.to_bytes(self.bytes_per_pixel as u8) as usize;
291        let expected = row_bytes * bounds.size.height.0.max(0) as usize;
292        if bytes.len() < expected {
293            log::error!(
294                "DirectXAtlasTexture::upload: source slice is {} bytes but the {}x{} region \
295                 requires {} bytes; skipping upload to avoid a driver over-read",
296                bytes.len(),
297                bounds.size.width.0,
298                bounds.size.height.0,
299                expected,
300            );
301            return;
302        }
303        unsafe {
304            device_context.UpdateSubresource(
305                &self.texture,
306                0,
307                Some(&D3D11_BOX {
308                    left: bounds.left().0 as u32,
309                    top: bounds.top().0 as u32,
310                    front: 0,
311                    right: bounds.right().0 as u32,
312                    bottom: bounds.bottom().0 as u32,
313                    back: 1,
314                }),
315                bytes.as_ptr() as _,
316                bounds.size.width.to_bytes(self.bytes_per_pixel as u8),
317                0,
318            );
319        }
320    }
321
322    fn decrement_ref_count(&mut self) {
323        self.live_atlas_keys -= 1;
324    }
325
326    fn is_unreferenced(&mut self) -> bool {
327        self.live_atlas_keys == 0
328    }
329}
330
331fn device_size_to_etagere(size: Size<DevicePixels>) -> etagere::Size {
332    etagere::Size::new(size.width.into(), size.height.into())
333}
334
335fn etagere_point_to_device(value: etagere::Point) -> Point<DevicePixels> {
336    Point {
337        x: DevicePixels::from(value.x),
338        y: DevicePixels::from(value.y),
339    }
340}
341
342#[cfg(test)]
343mod tests {
344    use super::*;
345    use gpui::{ImageId, RenderImageParams};
346    use std::borrow::Cow;
347    use windows::Win32::{
348        Foundation::HMODULE,
349        Graphics::{
350            Direct3D::D3D_DRIVER_TYPE_WARP,
351            Direct3D11::{D3D11_CREATE_DEVICE_BGRA_SUPPORT, D3D11_SDK_VERSION, D3D11CreateDevice},
352        },
353    };
354
355    fn create_atlas() -> Option<DirectXAtlas> {
356        let mut device: Option<ID3D11Device> = None;
357        let mut device_context: Option<ID3D11DeviceContext> = None;
358        unsafe {
359            D3D11CreateDevice(
360                None,
361                D3D_DRIVER_TYPE_WARP,
362                HMODULE::default(),
363                D3D11_CREATE_DEVICE_BGRA_SUPPORT,
364                None,
365                D3D11_SDK_VERSION,
366                Some(&mut device),
367                None,
368                Some(&mut device_context),
369            )
370        }
371        .ok()?;
372        Some(DirectXAtlas::new(&device?, &device_context?))
373    }
374
375    fn make_image_key(image_id: usize) -> AtlasKey {
376        AtlasKey::Image(RenderImageParams {
377            image_id: ImageId(image_id),
378            frame_index: 0,
379        })
380    }
381
382    fn insert_tile(atlas: &DirectXAtlas, key: &AtlasKey, size: Size<DevicePixels>) -> AtlasTile {
383        atlas
384            .get_or_insert_with(key, &mut || {
385                let byte_count = (size.width.0 as usize) * (size.height.0 as usize) * 4;
386                Ok(Some((size, Cow::Owned(vec![0u8; byte_count]))))
387            })
388            .expect("allocation should succeed")
389            .expect("callback returns Some")
390    }
391
392    #[test]
393    fn test_remove_deallocates_tile_space_for_reuse() {
394        let Some(atlas) = create_atlas() else {
395            return;
396        };
397
398        let small = Size {
399            width: DevicePixels(64),
400            height: DevicePixels(64),
401        };
402        let big = Size {
403            width: DevicePixels(700),
404            height: DevicePixels(700),
405        };
406
407        let keeper_key = make_image_key(1);
408        let big_key_a = make_image_key(2);
409        let big_key_b = make_image_key(3);
410
411        let keeper_tile = insert_tile(&atlas, &keeper_key, small);
412        let tile_a = insert_tile(&atlas, &big_key_a, big);
413        assert_eq!(keeper_tile.texture_id, tile_a.texture_id);
414
415        atlas.remove(&big_key_a);
416
417        let tile_b = insert_tile(&atlas, &big_key_b, big);
418        assert_eq!(tile_b.texture_id, keeper_tile.texture_id);
419    }
420}
421