Skip to main content

pfx_text/
atlas.rs

1use std::collections::BTreeMap;
2
3use crate::{ATLAS_CELL_ALIGN, ATLAS_PADDING, Atlas, Error, GlyphImage, MipLevel, RasterKey};
4
5pub const DEFAULT_ATLAS_BUDGET: u64 = 32 << 20;
6pub const DEFAULT_ATLAS_PAGE: u32 = 1024;
7
8#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, PartialOrd, Ord)]
9pub enum PageKind {
10    Coverage,
11    Msdf,
12    Color,
13}
14
15impl PageKind {
16    pub fn channels(self) -> u32 {
17        match self {
18            Self::Coverage => 1,
19            Self::Msdf => 3,
20            Self::Color => 4,
21        }
22    }
23
24    pub fn levels(self) -> u32 {
25        match self {
26            Self::Msdf => 1,
27            Self::Coverage | Self::Color => 4,
28        }
29    }
30
31    pub fn texel_bytes(self) -> u64 {
32        match self {
33            Self::Coverage => 1,
34            Self::Msdf | Self::Color => 4,
35        }
36    }
37
38    pub fn page_bytes(self, size: u32) -> u64 {
39        (0..self.levels())
40            .map(|level| u64::from((size >> level).max(1)).pow(2) * self.texel_bytes())
41            .sum()
42    }
43}
44
45#[derive(Clone, Copy, Debug, PartialEq, Eq)]
46pub struct PageInfo {
47    pub id: u32,
48    pub kind: PageKind,
49    pub channels: u32,
50    pub size: u32,
51    pub levels: u32,
52}
53
54#[derive(Clone, Debug, PartialEq, Eq)]
55pub struct CellUpload {
56    pub page: u32,
57    pub level: u32,
58    pub rect: [u32; 4],
59    pub bytes: Vec<u8>,
60}
61
62#[derive(Clone, Debug, Default, PartialEq, Eq)]
63pub struct AtlasChanges {
64    pub created: Vec<PageInfo>,
65    pub dropped: Vec<u32>,
66    pub cells: Vec<CellUpload>,
67}
68
69impl AtlasChanges {
70    pub fn is_empty(&self) -> bool {
71        self.created.is_empty() && self.dropped.is_empty() && self.cells.is_empty()
72    }
73
74    pub fn bytes(&self) -> u64 {
75        self.cells.iter().map(|cell| cell.bytes.len() as u64).sum()
76    }
77}
78
79#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
80pub struct AtlasStats {
81    pub pages: usize,
82    pub cells: usize,
83    pub bytes: u64,
84    pub budget: u64,
85    pub shaped: u64,
86    pub fast: u64,
87    pub rasterized: u64,
88    pub evicted: u64,
89    pub uploaded: u64,
90}
91
92#[derive(Clone, Copy, Debug, PartialEq)]
93pub(crate) struct CellRef {
94    pub page: u32,
95    pub kind: PageKind,
96    pub rect: [u32; 4],
97    pub size: u32,
98    pub left: f32,
99    pub top: f32,
100}
101
102struct Shelf {
103    y: u32,
104    height: u32,
105    free: Vec<[u32; 2]>,
106}
107
108struct Shelves {
109    units: u32,
110    shelves: Vec<Shelf>,
111    rows: Vec<[u32; 2]>,
112}
113
114fn insert_span(spans: &mut Vec<[u32; 2]>, start: u32, length: u32) {
115    let at = spans.partition_point(|span| span[0] < start);
116    spans.insert(at, [start, length]);
117    if at + 1 < spans.len() && spans[at][0] + spans[at][1] == spans[at + 1][0] {
118        spans[at][1] += spans[at + 1][1];
119        spans.remove(at + 1);
120    }
121    if at > 0 && spans[at - 1][0] + spans[at - 1][1] == spans[at][0] {
122        spans[at - 1][1] += spans[at][1];
123        spans.remove(at);
124    }
125}
126
127fn shelf_height(height: u32) -> u32 {
128    if height <= 4 {
129        height
130    } else {
131        height.next_multiple_of(2)
132    }
133}
134
135impl Shelves {
136    fn new(units: u32) -> Self {
137        Self {
138            units,
139            shelves: Vec::new(),
140            rows: vec![[0, units]],
141        }
142    }
143
144    fn allocate(&mut self, width: u32, height: u32) -> Option<[u32; 2]> {
145        if width > self.units || height > self.units {
146            return None;
147        }
148        let wanted = shelf_height(height);
149        let mut best: Option<(u32, usize, usize)> = None;
150        for (index, shelf) in self.shelves.iter().enumerate() {
151            if shelf.height < height || shelf.height > wanted + wanted / 2 {
152                continue;
153            }
154            if best.is_some_and(|(h, _, _)| h <= shelf.height) {
155                continue;
156            }
157            if let Some(span) = shelf.free.iter().position(|span| span[1] >= width) {
158                best = Some((shelf.height, index, span));
159            }
160        }
161        if best.is_none() {
162            let row = self.rows.iter().position(|row| row[1] >= wanted)?;
163            let y = self.rows[row][0];
164            self.rows[row][0] += wanted;
165            self.rows[row][1] -= wanted;
166            if self.rows[row][1] == 0 {
167                self.rows.remove(row);
168            }
169            let at = self.shelves.partition_point(|shelf| shelf.y < y);
170            self.shelves.insert(
171                at,
172                Shelf {
173                    y,
174                    height: wanted,
175                    free: vec![[0, self.units]],
176                },
177            );
178            best = Some((wanted, at, 0));
179        }
180        let (_, index, span) = best?;
181        let shelf = &mut self.shelves[index];
182        let x = shelf.free[span][0];
183        shelf.free[span][0] += width;
184        shelf.free[span][1] -= width;
185        if shelf.free[span][1] == 0 {
186            shelf.free.remove(span);
187        }
188        Some([x, shelf.y])
189    }
190
191    fn release(&mut self, x: u32, y: u32, width: u32) {
192        let Some(index) = self.shelves.iter().position(|shelf| shelf.y == y) else {
193            return;
194        };
195        let shelf = &mut self.shelves[index];
196        insert_span(&mut shelf.free, x, width);
197        if shelf.free == [[0, self.units]] {
198            let shelf = self.shelves.remove(index);
199            insert_span(&mut self.rows, shelf.y, shelf.height);
200        }
201    }
202
203    fn is_empty(&self) -> bool {
204        self.shelves.is_empty()
205    }
206}
207
208struct Page {
209    kind: PageKind,
210    face: String,
211    atlas: Atlas,
212    shelves: Shelves,
213}
214
215struct Cell {
216    page: u32,
217    slot: [u32; 4],
218    rect: [u32; 4],
219    left: f32,
220    top: f32,
221}
222
223struct Entry {
224    cell: Option<Cell>,
225    stamp: u64,
226}
227
228pub(crate) struct SharedAtlas {
229    budget: u64,
230    size: u32,
231    pages: BTreeMap<u32, Page>,
232    next_page: u32,
233    entries: BTreeMap<RasterKey, Entry>,
234    tick: u64,
235    frame_start: u64,
236    changes: AtlasChanges,
237    pub(crate) stats: AtlasStats,
238}
239
240impl SharedAtlas {
241    pub fn new() -> Self {
242        Self {
243            budget: DEFAULT_ATLAS_BUDGET,
244            size: DEFAULT_ATLAS_PAGE,
245            pages: BTreeMap::new(),
246            next_page: 0,
247            entries: BTreeMap::new(),
248            tick: 0,
249            frame_start: 0,
250            changes: AtlasChanges::default(),
251            stats: AtlasStats {
252                budget: DEFAULT_ATLAS_BUDGET,
253                ..AtlasStats::default()
254            },
255        }
256    }
257
258    pub fn set_budget(&mut self, bytes: u64) {
259        self.budget = bytes;
260        self.stats.budget = bytes;
261    }
262
263    pub fn set_page_size(&mut self, size: u32) -> Result<(), Error> {
264        if !(64..=8192).contains(&size) || !size.is_multiple_of(ATLAS_CELL_ALIGN) {
265            return Err(Error::InvalidStyle);
266        }
267        if size != self.size {
268            self.clear();
269            self.size = size;
270        }
271        Ok(())
272    }
273
274    pub fn clear(&mut self) {
275        let ids = self.pages.keys().copied().collect::<Vec<_>>();
276        for id in ids {
277            self.drop_page(id);
278        }
279        self.entries.clear();
280    }
281
282    pub fn begin_frame(&mut self) {
283        self.frame_start = self.tick;
284        if self.entries.len() > 1 << 16 {
285            let start = self.frame_start;
286            self.entries
287                .retain(|_, entry| entry.cell.is_some() || entry.stamp >= start);
288        }
289    }
290
291    pub fn take_changes(&mut self) -> AtlasChanges {
292        std::mem::take(&mut self.changes)
293    }
294
295    pub fn page(&self, id: u32) -> Option<&Atlas> {
296        self.pages.get(&id).map(|page| &page.atlas)
297    }
298
299    pub fn stats(&self) -> AtlasStats {
300        AtlasStats {
301            pages: self.pages.len(),
302            cells: self
303                .entries
304                .values()
305                .filter(|entry| entry.cell.is_some())
306                .count(),
307            bytes: self.bytes(),
308            ..self.stats
309        }
310    }
311
312    fn bytes(&self) -> u64 {
313        self.pages
314            .values()
315            .map(|page| page.kind.page_bytes(self.size))
316            .sum()
317    }
318
319    pub fn touch(&mut self, key: &RasterKey) -> Option<Option<CellRef>> {
320        let entry = self.entries.get_mut(key)?;
321        self.tick += 1;
322        entry.stamp = self.tick;
323        Some(entry.cell.as_ref().map(|cell| {
324            let page = &self.pages[&cell.page];
325            CellRef {
326                page: cell.page,
327                kind: page.kind,
328                rect: cell.rect,
329                size: self.size,
330                left: cell.left,
331                top: cell.top,
332            }
333        }))
334    }
335
336    pub fn insert(
337        &mut self,
338        key: RasterKey,
339        kind: PageKind,
340        image: Option<GlyphImage>,
341    ) -> Result<Option<CellRef>, Error> {
342        self.stats.rasterized += 1;
343        self.tick += 1;
344        let stamp = self.tick;
345        let Some(image) = image else {
346            self.entries.insert(key, Entry { cell: None, stamp });
347            return Ok(None);
348        };
349        let width = (image.width + ATLAS_PADDING * 2).next_multiple_of(ATLAS_CELL_ALIGN);
350        let height = (image.height + ATLAS_PADDING * 2).next_multiple_of(ATLAS_CELL_ALIGN);
351        if width > self.size || height > self.size {
352            return Err(Error::AtlasTooLarge);
353        }
354        let units = [width / ATLAS_CELL_ALIGN, height / ATLAS_CELL_ALIGN];
355        let face = key.font_id.clone();
356        let (page, origin) = self.place(&face, kind, units)?;
357        let slot = [
358            origin[0] * ATLAS_CELL_ALIGN,
359            origin[1] * ATLAS_CELL_ALIGN,
360            width,
361            height,
362        ];
363        let rect = [
364            slot[0] + ATLAS_PADDING,
365            slot[1] + ATLAS_PADDING,
366            image.width,
367            image.height,
368        ];
369        self.write(page, slot, &image);
370        self.entries.insert(
371            key,
372            Entry {
373                cell: Some(Cell {
374                    page,
375                    slot,
376                    rect,
377                    left: image.left,
378                    top: image.top,
379                }),
380                stamp,
381            },
382        );
383        Ok(Some(CellRef {
384            page,
385            kind,
386            rect,
387            size: self.size,
388            left: image.left,
389            top: image.top,
390        }))
391    }
392
393    fn place(
394        &mut self,
395        face: &str,
396        kind: PageKind,
397        units: [u32; 2],
398    ) -> Result<(u32, [u32; 2]), Error> {
399        let mut victims: Option<Vec<(u64, RasterKey)>> = None;
400        loop {
401            for (&id, page) in &mut self.pages {
402                if page.kind == kind
403                    && page.face == face
404                    && let Some(origin) = page.shelves.allocate(units[0], units[1])
405                {
406                    return Ok((id, origin));
407                }
408            }
409            if self.bytes() + kind.page_bytes(self.size) <= self.budget {
410                let id = self.create_page(face, kind);
411                let origin = self
412                    .pages
413                    .get_mut(&id)
414                    .and_then(|page| page.shelves.allocate(units[0], units[1]))
415                    .ok_or(Error::AtlasTooLarge)?;
416                return Ok((id, origin));
417            }
418            let list = victims.get_or_insert_with(|| {
419                let start = self.frame_start;
420                let mut list = self
421                    .entries
422                    .iter()
423                    .filter(|(_, entry)| entry.cell.is_some() && entry.stamp <= start)
424                    .map(|(key, entry)| (entry.stamp, key.clone()))
425                    .collect::<Vec<_>>();
426                list.sort_by_key(|entry| std::cmp::Reverse(entry.0));
427                list
428            });
429            let (_, key) = list.pop().ok_or(Error::AtlasTooLarge)?;
430            self.evict(&key);
431        }
432    }
433
434    fn evict(&mut self, key: &RasterKey) {
435        let Some(Entry {
436            cell: Some(cell), ..
437        }) = self.entries.remove(key)
438        else {
439            return;
440        };
441        self.stats.evicted += 1;
442        let Some(page) = self.pages.get_mut(&cell.page) else {
443            return;
444        };
445        page.shelves.release(
446            cell.slot[0] / ATLAS_CELL_ALIGN,
447            cell.slot[1] / ATLAS_CELL_ALIGN,
448            cell.slot[2] / ATLAS_CELL_ALIGN,
449        );
450        if page.shelves.is_empty() {
451            self.drop_page(cell.page);
452        }
453    }
454
455    fn create_page(&mut self, face: &str, kind: PageKind) -> u32 {
456        let id = self.next_page;
457        self.next_page += 1;
458        let channels = kind.channels();
459        let levels = (0..kind.levels())
460            .map(|level| {
461                let side = (self.size >> level).max(1);
462                MipLevel {
463                    width: side,
464                    height: side,
465                    bytes: vec![0; (side * side * channels) as usize],
466                }
467            })
468            .collect();
469        self.pages.insert(
470            id,
471            Page {
472                kind,
473                face: face.to_owned(),
474                atlas: Atlas { channels, levels },
475                shelves: Shelves::new(self.size / ATLAS_CELL_ALIGN),
476            },
477        );
478        self.changes.created.push(PageInfo {
479            id,
480            kind,
481            channels,
482            size: self.size,
483            levels: kind.levels(),
484        });
485        id
486    }
487
488    fn drop_page(&mut self, id: u32) {
489        if self.pages.remove(&id).is_none() {
490            return;
491        }
492        self.entries
493            .retain(|_, entry| entry.cell.as_ref().is_none_or(|cell| cell.page != id));
494        self.changes.cells.retain(|cell| cell.page != id);
495        if let Some(at) = self.changes.created.iter().position(|page| page.id == id) {
496            self.changes.created.remove(at);
497        } else {
498            self.changes.dropped.push(id);
499        }
500    }
501
502    fn write(&mut self, id: u32, slot: [u32; 4], image: &GlyphImage) {
503        let page = self.pages.get_mut(&id).unwrap();
504        let channels = page.atlas.channels as usize;
505        let levels = &mut page.atlas.levels;
506        let width = levels[0].width;
507        for row in slot[1]..slot[1] + slot[3] {
508            let start = ((row * width + slot[0]) as usize) * channels;
509            levels[0].bytes[start..start + slot[2] as usize * channels].fill(0);
510        }
511        let padding = ATLAS_PADDING as i32;
512        let origin = [slot[0] as i32 + padding, slot[1] as i32 + padding];
513        for dy in -padding..image.height as i32 + padding {
514            for dx in -padding..image.width as i32 + padding {
515                let sx = dx.clamp(0, image.width as i32 - 1) as usize;
516                let sy = dy.clamp(0, image.height as i32 - 1) as usize;
517                let src = (sy * image.width as usize + sx) * channels;
518                let dst = (((origin[1] + dy) as u32 * width + (origin[0] + dx) as u32) as usize)
519                    * channels;
520                levels[0].bytes[dst..dst + channels]
521                    .copy_from_slice(&image.bytes[src..src + channels]);
522            }
523        }
524        for level in 1..levels.len() {
525            let (done, rest) = levels.split_at_mut(level);
526            let prev = &done[level - 1];
527            let next = &mut rest[0];
528            let shift = level as u32;
529            let [x0, y0, w, h] = [
530                slot[0] >> shift,
531                slot[1] >> shift,
532                slot[2] >> shift,
533                slot[3] >> shift,
534            ];
535            for row in y0..y0 + h {
536                for col in x0..x0 + w {
537                    for channel in 0..channels {
538                        let mut sum = 0u32;
539                        for dy in 0..2 {
540                            for dx in 0..2 {
541                                let sx = (2 * col + dx).min(prev.width - 1);
542                                let sy = (2 * row + dy).min(prev.height - 1);
543                                sum += u32::from(
544                                    prev.bytes
545                                        [(sy * prev.width + sx) as usize * channels + channel],
546                                );
547                            }
548                        }
549                        next.bytes[(row * next.width + col) as usize * channels + channel] =
550                            ((sum + 2) / 4) as u8;
551                    }
552                }
553            }
554        }
555        for (index, level) in levels.iter().enumerate() {
556            let shift = index as u32;
557            let rect = [
558                slot[0] >> shift,
559                slot[1] >> shift,
560                (slot[2] >> shift).max(1),
561                (slot[3] >> shift).max(1),
562            ];
563            let mut bytes = Vec::with_capacity((rect[2] * rect[3]) as usize * channels);
564            for row in rect[1]..rect[1] + rect[3] {
565                let start = ((row * level.width + rect[0]) as usize) * channels;
566                bytes.extend_from_slice(&level.bytes[start..start + rect[2] as usize * channels]);
567            }
568            self.stats.uploaded += (rect[2] * rect[3]) as u64 * page.kind.texel_bytes();
569            self.changes.cells.push(CellUpload {
570                page: id,
571                level: index as u32,
572                rect,
573                bytes,
574            });
575        }
576    }
577}
578
579#[cfg(test)]
580mod tests {
581    use super::*;
582
583    #[test]
584    fn shelves_reuse_freed_space_and_return_empty_rows() {
585        let mut shelves = Shelves::new(16);
586        let a = shelves.allocate(8, 4).unwrap();
587        let b = shelves.allocate(8, 4).unwrap();
588        assert_eq!(a, [0, 0]);
589        assert_eq!(b, [8, 0]);
590        assert!(shelves.allocate(1, 4).is_some_and(|at| at[1] == 4));
591        shelves.release(a[0], a[1], 8);
592        assert_eq!(shelves.allocate(8, 3), Some([0, 0]));
593        shelves.release(0, 0, 8);
594        shelves.release(8, 0, 8);
595        assert_eq!(shelves.shelves.len(), 1);
596        shelves.release(0, 4, 1);
597        assert!(shelves.is_empty());
598        assert_eq!(shelves.rows, vec![[0, 16]]);
599        assert_eq!(shelves.allocate(16, 16), Some([0, 0]));
600        assert_eq!(shelves.allocate(1, 1), None);
601    }
602
603    #[test]
604    fn page_bytes_count_every_level_at_its_gpu_format() {
605        assert_eq!(PageKind::Msdf.page_bytes(1024), 4 << 20);
606        assert_eq!(
607            PageKind::Coverage.page_bytes(1024),
608            (1 << 20) + (1 << 18) + (1 << 16) + (1 << 14)
609        );
610    }
611}