klyff 0.1.3

//! Serializable snapshot of a [`TextureAtlas`], gated behind the `serde` feature.
use std::collections::HashMap;

use serde::{Deserialize, Serialize};

use super::TextureAtlas;
use super::types::{self, GlyphCacheKey, GlyphCacheValue};
use crate::Rect;

/// Bytes per pixel of the atlas texture (`Rgba8Unorm`).
const BYTES_PER_PIXEL: u32 = 4;

/// A fully baked atlas: the serializable metadata plus the rendered texture pixels.
///
/// The texture is kept as a separate `Vec<u8>` so callers may store it on its own (e.g. as a
/// PNG) if they prefer.
pub struct BakedAtlas {
    /// Tightly packed `Rgba8Unorm` pixels, `width * height * layers * 4` bytes, laid out
    /// layer-major then row-major.
    pub texture_data: Vec<u8>,
    /// Packing state, glyph cache and configuration.
    pub state: AtlasState,
}

/// Serializable metadata of a [`TextureAtlas`]: configuration, etagere packing state and the
/// glyph cache. Does not include the texture pixels (see [`BakedAtlas`]).
#[derive(Serialize, Deserialize)]
pub struct AtlasState {
    width: u32,
    height: u32,
    layers: u32,
    max_width: u32,
    max_height: u32,
    max_layers: u32,
    ppem: f32,
    padding_x: u32,
    padding_y: u32,
    allow_grow: bool,
    retain_glyphs_frame_duration: u64,
    allocators: Vec<etagere::BucketedAtlasAllocator>,
    glyphs: Vec<SerGlyph>,
    empty_glyphs: Vec<(StableFontKey, u32)>,
}

/// Stable, cross-run identity of a font face.
#[derive(Serialize, Deserialize, Clone, PartialEq, Eq, Hash)]
struct StableFontKey {
    post_script_name: String,
    index: u32,
}

/// Serializable mirror of a single glyph cache entry (key + value flattened together).
#[derive(Serialize, Deserialize)]
struct SerGlyph {
    font: StableFontKey,
    /// `Some` for rasterized bitmaps (size-specific), `None` for distance fields.
    font_size_bytes: Option<[u8; 4]>,
    glyph_id: u32,
    layer: u32,
    min_x: u32,
    min_y: u32,
    region_width: u32,
    region_height: u32,
    region_padding_x: u32,
    region_padding_y: u32,
    em_size: f32,
    /// `[min.x, min.y, max.x, max.y]` of the glyph bounds in em units.
    bounds_em: [f32; 4],
    alloc_id: u32,
}

/// Error returned by [`TextureAtlas::bake`] / [`TextureAtlas::from_baked`].
#[derive(Debug)]
pub enum BakeError {
    /// Mapping the GPU readback buffer failed while reading the texture back.
    BufferMapFailed,
    /// The texture pixel data does not match the size declared in [`AtlasState`].
    InvalidTextureData { expected: usize, got: usize },
}

impl std::fmt::Display for BakeError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            BakeError::BufferMapFailed => write!(f, "failed to map atlas readback buffer"),
            BakeError::InvalidTextureData { expected, got } => write!(
                f,
                "atlas texture data has wrong length: expected {expected} bytes, got {got}"
            ),
        }
    }
}
impl std::error::Error for BakeError {}

/// Resolve a font's stable identity from the database, or `None` if the face is unknown.
fn stable_key(db: &fontdb::Database, id: fontdb::ID) -> Option<StableFontKey> {
    let face = db.face(id)?;
    Some(StableFontKey {
        post_script_name: face.post_script_name.clone(),
        index: face.index,
    })
}

impl TextureAtlas {
    /// Capture the full state of the atlas into a serializable [`BakedAtlas`].
    ///
    /// This reads the atlas texture back from the GPU (this submits work and blocks until it
    /// completes). Glyphs whose font is not present in `db` are skipped with a warning log.
    ///
    /// The returned [`BakedAtlas`] can later be passed to [`TextureAtlas::from_baked`] to
    /// reconstruct an equivalent atlas without regenerating any glyphs.
    pub fn bake(
        &self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        db: &fontdb::Database,
    ) -> Result<BakedAtlas, BakeError> {
        profiling::scope!("TextureAtlas::bake");

        let mut glyphs = Vec::new();
        for (key, value) in self.cache.iter() {
            let Some(font) = stable_key(db, key.font) else {
                log::warn!(
                    "bake: skipping cached glyph for font id {:?} not found in database",
                    key.font
                );
                continue;
            };
            let region = &value.region;
            glyphs.push(SerGlyph {
                font,
                font_size_bytes: key.font_size_bytes,
                glyph_id: key.id.to_u32(),
                layer: region.layer,
                min_x: region.min_x,
                min_y: region.min_y,
                region_width: region.width,
                region_height: region.height,
                region_padding_x: region.padding_x,
                region_padding_y: region.padding_y,
                em_size: value.em_size,
                bounds_em: [
                    value.bounds_em.min.x,
                    value.bounds_em.min.y,
                    value.bounds_em.max.x,
                    value.bounds_em.max.y,
                ],
                alloc_id: value.alloc_id.serialize(),
            });
        }

        let mut empty_glyphs = Vec::new();
        for (font_id, glyph_id) in self.empty_glyphs.iter() {
            let Some(font) = stable_key(db, *font_id) else {
                log::warn!(
                    "bake: skipping empty-glyph entry for font id {font_id:?} not found in database"
                );
                continue;
            };
            empty_glyphs.push((font, glyph_id.to_u32()));
        }

        let state = AtlasState {
            width: self.current_size.width,
            height: self.current_size.height,
            layers: self.current_size.depth_or_array_layers,
            max_width: self.max_size.width,
            max_height: self.max_size.height,
            max_layers: self.max_size.depth_or_array_layers,
            ppem: self.ppem,
            padding_x: self.padding_x as u32,
            padding_y: self.padding_y as u32,
            allow_grow: self.allow_grow,
            retain_glyphs_frame_duration: self.retain_glyphs_frame_duration,
            allocators: self.allocators.clone(),
            glyphs,
            empty_glyphs,
        };

        let texture_data = self.read_texture(device, queue)?;
        Ok(BakedAtlas {
            texture_data,
            state,
        })
    }

    /// Read the atlas texture back into a tightly packed `Rgba8Unorm` buffer.
    fn read_texture(
        &self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
    ) -> Result<Vec<u8>, BakeError> {
        let size = self.current_size;
        let width = size.width;
        let height = size.height;
        let layers = size.depth_or_array_layers;

        // Pad each row to COPY_BYTES_PER_ROW_ALIGNMENT (256) for copy_texture_to_buffer.
        let unaligned = width * BYTES_PER_PIXEL;
        let align = wgpu::COPY_BYTES_PER_ROW_ALIGNMENT;
        let padded_bytes_per_row = unaligned.div_ceil(align) * align;
        let buffer_size = padded_bytes_per_row as u64 * height as u64 * layers as u64;

        let readback = device.create_buffer(&wgpu::BufferDescriptor {
            label: Some("klyff atlas bake readback"),
            size: buffer_size,
            usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
            mapped_at_creation: false,
        });

        let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
            label: Some("klyff atlas bake encoder"),
        });
        encoder.copy_texture_to_buffer(
            self.texture.as_image_copy(),
            wgpu::TexelCopyBufferInfo {
                buffer: &readback,
                layout: wgpu::TexelCopyBufferLayout {
                    offset: 0,
                    bytes_per_row: Some(padded_bytes_per_row),
                    rows_per_image: Some(height),
                },
            },
            size,
        );
        queue.submit([encoder.finish()]);

        let slice = readback.slice(..);
        slice.map_async(wgpu::MapMode::Read, |_| {});
        device
            .poll(wgpu::PollType::wait_indefinitely())
            .map_err(|_| BakeError::BufferMapFailed)?;

        let mapped = slice.get_mapped_range();
        let row_bytes = (width * BYTES_PER_PIXEL) as usize;
        let mut tight = vec![0u8; row_bytes * height as usize * layers as usize];
        let layer_stride = padded_bytes_per_row as usize * height as usize;
        for layer in 0..layers as usize {
            for y in 0..height as usize {
                let src = layer * layer_stride + y * padded_bytes_per_row as usize;
                let dst = (layer * height as usize + y) * row_bytes;
                tight[dst..dst + row_bytes].copy_from_slice(&mapped[src..src + row_bytes]);
            }
        }
        drop(mapped);
        readback.unmap();
        Ok(tight)
    }

    /// Reconstruct an atlas from a [`BakedAtlas`] previously produced by [`TextureAtlas::bake`].
    ///
    /// This recreates the GPU texture from saved pixels and rebuilds the glyph cache.
    /// Glyphs whose font is not present in `db` are skipped with a warning log (their reserved
    /// atlas space remains allocated).
    pub fn from_baked(
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        db: &fontdb::Database,
        baked: BakedAtlas,
    ) -> Result<TextureAtlas, BakeError> {
        profiling::scope!("TextureAtlas::from_baked");
        let BakedAtlas {
            texture_data,
            state,
        } = baked;

        let expected = state.width as usize
            * state.height as usize
            * state.layers as usize
            * BYTES_PER_PIXEL as usize;
        if texture_data.len() != expected {
            return Err(BakeError::InvalidTextureData {
                expected,
                got: texture_data.len(),
            });
        }

        let size = wgpu::Extent3d {
            width: state.width,
            height: state.height,
            depth_or_array_layers: state.layers,
        };
        let usage = types::texture_usage(state.allow_grow);
        let (texture, sampler) = types::create_texture(device, size, usage);
        let bind_group = TextureAtlas::create_bind_group(device, &texture, &sampler);

        queue.write_texture(
            texture.as_image_copy(),
            &texture_data,
            wgpu::TexelCopyBufferLayout {
                offset: 0,
                bytes_per_row: Some(state.width * BYTES_PER_PIXEL),
                rows_per_image: Some(state.height),
            },
            size,
        );

        // Reverse lookup from stable identity to this run's font ids.
        let mut font_ids: HashMap<StableFontKey, fontdb::ID> = HashMap::new();
        for face in db.faces() {
            font_ids
                .entry(StableFontKey {
                    post_script_name: face.post_script_name.clone(),
                    index: face.index,
                })
                .or_insert(face.id);
        }

        let mut cache = super::cache::Cache::new();
        for g in state.glyphs {
            let Some(&font) = font_ids.get(&g.font) else {
                log::warn!(
                    "from_baked: skipping glyph for font \"{}\" (index {}) not found in database",
                    g.font.post_script_name,
                    g.font.index
                );
                continue;
            };
            let key = GlyphCacheKey {
                font,
                font_size_bytes: g.font_size_bytes,
                id: skrifa::GlyphId::new(g.glyph_id),
            };
            let value = GlyphCacheValue {
                region: super::AtlasRegion {
                    layer: g.layer,
                    min_x: g.min_x,
                    min_y: g.min_y,
                    width: g.region_width,
                    height: g.region_height,
                    padding_x: g.region_padding_x,
                    padding_y: g.region_padding_y,
                },
                em_size: g.em_size,
                bounds_em: Rect {
                    min: glam::vec2(g.bounds_em[0], g.bounds_em[1]),
                    max: glam::vec2(g.bounds_em[2], g.bounds_em[3]),
                },
                alloc_id: etagere::AllocId::deserialize(g.alloc_id),
            };
            cache.insert(key, value);
        }

        let mut empty_glyphs = std::collections::HashSet::new();
        for (font, glyph_id) in state.empty_glyphs {
            let Some(&font) = font_ids.get(&font) else {
                log::warn!(
                    "from_baked: skipping empty-glyph entry for font \"{}\" (index {}) not found in database",
                    font.post_script_name,
                    font.index
                );
                continue;
            };
            empty_glyphs.insert((font, skrifa::GlyphId::new(glyph_id)));
        }

        Ok(TextureAtlas {
            bind_group,
            texture,
            sampler,
            allocators: state.allocators,
            cache,
            empty_glyphs,
            gpu_write: klyff_msdf::MtsdfGpuWriter::new(device),
            current_size: size,
            max_size: wgpu::Extent3d {
                width: state.max_width,
                height: state.max_height,
                depth_or_array_layers: state.max_layers,
            },
            ppem: state.ppem,
            padding_x: state.padding_x as usize,
            padding_y: state.padding_y as usize,
            allow_grow: state.allow_grow,
            retain_glyphs_frame_duration: state.retain_glyphs_frame_duration,
            frozen: false,
            scale_context: swash::scale::ScaleContext::new(),
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use skrifa::MetadataProvider;

    const OPEN_SANS: &[u8] = include_bytes!("../../../../fonts/OpenSans-Regular.ttf");

    async fn request_device() -> Option<(wgpu::Device, wgpu::Queue)> {
        let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor::default());
        let adapter = instance
            .request_adapter(&wgpu::RequestAdapterOptions {
                power_preference: wgpu::PowerPreference::None,
                compatible_surface: None,
                force_fallback_adapter: false,
            })
            .await
            .ok()?;
        adapter
            .request_device(&wgpu::DeviceDescriptor::default())
            .await
            .ok()
    }

    /// Flatten the glyph cache into an order-independent, comparable map.
    type Entry = (fontdb::ID, Option<[u8; 4]>, u32);
    type RegionTuple = (u32, u32, u32, u32, u32, u32, u32, u32);
    fn cache_entries(atlas: &TextureAtlas) -> HashMap<Entry, RegionTuple> {
        atlas
            .cache
            .iter()
            .map(|(k, v)| {
                let r = &v.region;
                (
                    (k.font, k.font_size_bytes, k.id.to_u32()),
                    (
                        r.layer,
                        r.min_x,
                        r.min_y,
                        r.width,
                        r.height,
                        r.padding_x,
                        r.padding_y,
                        v.alloc_id.serialize(),
                    ),
                )
            })
            .collect()
    }

    #[test]
    fn bake_round_trip_preserves_cache_and_pixels() {
        let Some((device, queue)) = pollster::block_on(request_device()) else {
            eprintln!("no wgpu device available; skipping bake round-trip test");
            return;
        };

        // Single-face database so the bake/load font translation is unambiguous.
        let mut db = fontdb::Database::new();
        db.load_font_data(OPEN_SANS.to_vec());
        let font_id = db.faces().next().unwrap().id;
        let face_index = db.face(font_id).unwrap().index;
        let mut font_system =
            cosmic_text::FontSystem::new_with_locale_and_db("en-US".to_string(), db);
        let font = font_system
            .get_font(font_id, fontdb::Weight::NORMAL)
            .unwrap();
        let font_data = super::super::FontData::new(font, face_index);

        // Resolve real glyph ids (with outlines) for a handful of letters.
        let font_ref = skrifa::FontRef::new(OPEN_SANS).unwrap();
        let charmap = font_ref.charmap();
        let glyph_ids: Vec<u16> = "klyff"
            .chars()
            .map(|c| charmap.map(c).unwrap().to_u32() as u16)
            .collect();

        // Small, fixed, non-growable atlas so the texture readback stays cheap and exercises
        // the serde-only COPY_SRC path.
        let mut atlas = TextureAtlas::new(
            &device,
            super::super::TextureAtlasDescriptor {
                initial_size: super::super::AtlasSize::SingleLayer {
                    width: 256,
                    height: 256,
                },
                max_size: Some(super::super::AtlasSize::SingleLayer {
                    width: 256,
                    height: 256,
                }),
                allow_grow: false,
                ..Default::default()
            },
        );

        let mut generator = klyff_msdf::MsdfGenerator::new();
        let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
        atlas.start_of_frame();
        for &gid in &glyph_ids {
            atlas
                .retrieve_or_generate_glyph(
                    super::super::GlyphKey::new(font_id, 32.0, gid),
                    &font_data,
                    &mut generator,
                    None,
                    (&device, &queue, &mut encoder),
                )
                .unwrap();
        }
        atlas.write_glyphs(&device, &queue, &mut encoder);
        queue.submit([encoder.finish()]);
        device.poll(wgpu::PollType::wait_indefinitely()).unwrap();

        assert!(
            !cache_entries(&atlas).is_empty(),
            "expected some glyphs to be cached as distance fields"
        );

        // Bake -> serialize -> deserialize -> reload.
        let baked = atlas.bake(&device, &queue, font_system.db()).unwrap();
        assert_eq!(baked.texture_data.len(), 256 * 256 * 4);
        let json = serde_json::to_vec(&baked.state).unwrap();
        let restored_baked: AtlasState = serde_json::from_slice(&json).unwrap();
        let restored = TextureAtlas::from_baked(
            &device,
            &queue,
            font_system.db(),
            BakedAtlas {
                texture_data: baked.texture_data.clone(),
                state: restored_baked,
            },
        )
        .unwrap();

        assert_eq!(restored.atlas_size(), atlas.atlas_size());
        assert_eq!(cache_entries(&restored), cache_entries(&atlas));

        // Re-baking the restored atlas must reproduce the exact same pixels.
        let rebaked = restored.bake(&device, &queue, font_system.db()).unwrap();
        assert_eq!(rebaked.texture_data, baked.texture_data);
    }
}