klyff_msdf 0.1.3

//! GPU-accelerated glyph generation utilities.

#![cfg(feature = "wgpu")]
use bytemuck::{Pod, Zeroable};

use crate::generator::GlyphOutline;
use crate::math_segment::Segment;
use crate::segment_soa::SegmentSoa;

/// Axis-aligned region of an atlas texture to render a glyph into.
///
/// Coordinates are in pixels, on a specific texture array layer. For a plain
/// 2D (non-array) texture, `layer` must be 0.
#[derive(Debug, Clone, Copy)]
pub struct AtlasRegion {
    pub layer: u32,
    pub min_x: u32,
    pub min_y: u32,
    pub width: u32,
    pub height: u32,
    pub padding_x: u32,
    pub padding_y: u32,
}

#[repr(C)]
#[derive(Clone, Copy, Pod, Zeroable)]
struct SegmentGpu {
    from: [f32; 2],
    to: [f32; 2],
    mid: [f32; 2],
    tangent: [f32; 2],
    tangent_unit: [f32; 2],
    inv_tangent_len_sq: f32,
    color_mask: u32,
}

#[repr(C)]
#[derive(Clone, Copy, Pod, Zeroable)]
struct Vertex {
    clip_pos: [f32; 2],
    bound_pos: [f32; 2],
    seg_from: u32,
    seg_to: u32,
    units_per_em: f32,
}

#[derive(Default)]
struct LayerBatch {
    vertices: Vec<Vertex>,
    indices: Vec<u32>,
}

/// Helper for generating MTSDF directly into a texture atlas on the GPU.
///
/// Create this once and stores across the app lifetime.
///
/// The shader always emits full MTSDF, however you can still write MSDF if the target
/// texture only has RGB channel. Typically RGBA texture is more common however, and as such
/// there is no reason not to store the full MTSDF.
pub struct MtsdfGpuWriter {
    pipeline: wgpu::RenderPipeline,
    bgl: wgpu::BindGroupLayout,

    segments: Vec<SegmentGpu>,
    layers: Vec<LayerBatch>,

    segment_buf: Option<wgpu::Buffer>,
    segment_cap: u64,
    vertex_buf: Option<wgpu::Buffer>,
    vertex_cap: u64,
    index_buf: Option<wgpu::Buffer>,
    index_cap: u64,
    bind_group: Option<wgpu::BindGroup>,
}

impl MtsdfGpuWriter {
    pub fn new(device: &wgpu::Device) -> Self {
        let shader = device.create_shader_module(wgpu::ShaderModuleDescriptor {
            label: Some("klyff_msdf gpu shader"),
            source: wgpu::ShaderSource::Wgsl(std::borrow::Cow::Borrowed(include_str!(
                "msdf_gen.wgsl"
            ))),
        });

        let bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
            label: Some("klyff_msdf gpu bgl"),
            entries: &[wgpu::BindGroupLayoutEntry {
                binding: 0,
                visibility: wgpu::ShaderStages::FRAGMENT,
                ty: wgpu::BindingType::Buffer {
                    ty: wgpu::BufferBindingType::Storage { read_only: true },
                    has_dynamic_offset: false,
                    min_binding_size: None,
                },
                count: None,
            }],
        });

        let pipeline_layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: Some("klyff_msdf gpu pipeline layout"),
            bind_group_layouts: &[&bgl],
            push_constant_ranges: &[],
        });

        let vertex_layout = wgpu::VertexBufferLayout {
            array_stride: std::mem::size_of::<Vertex>() as u64,
            step_mode: wgpu::VertexStepMode::Vertex,
            attributes: &wgpu::vertex_attr_array![
                0 => Float32x2, // clip_pos
                1 => Float32x2, // bound_pos
                2 => Uint32,    // seg_from
                3 => Uint32,    // seg_to
                4 => Float32,   // em_size
            ],
        };

        let pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
            label: Some("klyff_msdf gpu pipeline"),
            layout: Some(&pipeline_layout),
            vertex: wgpu::VertexState {
                module: &shader,
                entry_point: Some("vs_main"),
                buffers: &[vertex_layout],
                compilation_options: Default::default(),
            },
            fragment: Some(wgpu::FragmentState {
                module: &shader,
                entry_point: Some("fs_main"),
                targets: &[Some(wgpu::ColorTargetState {
                    format: wgpu::TextureFormat::Rgba8Unorm,
                    blend: None,
                    write_mask: wgpu::ColorWrites::ALL,
                })],
                compilation_options: Default::default(),
            }),
            primitive: wgpu::PrimitiveState {
                topology: wgpu::PrimitiveTopology::TriangleList,
                cull_mode: None,
                ..Default::default()
            },
            depth_stencil: None,
            multisample: wgpu::MultisampleState::default(),
            multiview: None,
            cache: None,
        });

        Self {
            pipeline,
            bgl,
            segments: Vec::new(),
            layers: Vec::new(),
            segment_buf: None,
            segment_cap: 0,
            vertex_buf: None,
            vertex_cap: 0,
            index_buf: None,
            index_cap: 0,
            bind_group: None,
        }
    }

    /// Queue a glyph to be rasterized into `region` on the next call to [`Self::write_glyphs`].
    pub fn add_glyph(&mut self, outline: GlyphOutline<'_>, region: AtlasRegion) {
        profiling::scope!("GpuMsdfWriter::add_glyph");
        let segs = &outline.generator.cached_segment_vec;
        let soa = &outline.generator.cached_soa;
        let units_per_em = outline.units_per_em;
        let bound = outline.bound;

        let seg_from = self.segments.len() as u32;
        push_segments(&mut self.segments, segs, soa);
        let seg_to = self.segments.len() as u32;

        // Per-vertex bound positions at the 4 region corners. These extend
        // beyond `bound` by the padding-in-em on each side; interpolating them
        // across the quad gives exactly the per-pixel `pos` the CPU loop
        // computes from `xt` / `yt`.
        let total_w = region.width as f32;
        let total_h = region.height as f32;
        let inner_w = (region.width - 2 * region.padding_x).max(1) as f32;
        let inner_h = (region.height - 2 * region.padding_y).max(1) as f32;
        let px = region.padding_x as f32;
        let py = region.padding_y as f32;

        let bound_size = bound.max - bound.min;
        let bx = |uv_x: f32| bound.min.x + bound_size.x * (uv_x - px) / inner_w;
        // Y flip: memory row 0 corresponds to `total_h - 1` in CPU `y` indexing,
        // so the top of the framebuffer (uv.y = 0) maps to bound.max.y.
        let by = |uv_y: f32| bound.min.y + bound_size.y * (total_h - uv_y - py) / inner_h;

        // Clip-space corners of the region within its atlas layer.
        // Will be finalised in `write_glyphs` once the atlas size is known,
        // so for now store atlas-pixel coordinates in `clip_pos` and convert
        // at upload time. Use NaN guard to fail loudly if forgotten.
        let pxmin = region.min_x as f32;
        let pymin = region.min_y as f32;
        let pxmax = (region.min_x + region.width) as f32;
        let pymax = (region.min_y + region.height) as f32;

        let mk = |px_atlas: f32, py_atlas: f32, uv_x: f32, uv_y: f32| Vertex {
            clip_pos: [px_atlas, py_atlas],
            bound_pos: [bx(uv_x), by(uv_y)],
            seg_from,
            seg_to,
            units_per_em,
        };

        let layer = region.layer as usize;
        if self.layers.len() <= layer {
            self.layers.resize_with(layer + 1, LayerBatch::default);
        }
        let batch = &mut self.layers[layer];
        let v0 = batch.vertices.len() as u32;
        batch.vertices.push(mk(pxmin, pymin, 0.0, 0.0));
        batch.vertices.push(mk(pxmax, pymin, total_w, 0.0));
        batch.vertices.push(mk(pxmin, pymax, 0.0, total_h));
        batch.vertices.push(mk(pxmax, pymax, total_w, total_h));
        batch.indices.push(v0);
        batch.indices.push(v0 + 2);
        batch.indices.push(v0 + 1);
        batch.indices.push(v0 + 1);
        batch.indices.push(v0 + 2);
        batch.indices.push(v0 + 3);
    }

    /// Upload all queued glyphs and run a render pass per touched atlas layer.
    ///
    /// `texture` may be either a `D2` or `D2Array` texture, as long as every
    /// queued [`AtlasRegion::layer`] is less than `texture.depth_or_array_layers()`.
    /// For a plain 2D texture this means all regions must have `layer == 0`.
    pub fn write_glyphs(
        &mut self,
        device: &wgpu::Device,
        queue: &wgpu::Queue,
        encoder: &mut wgpu::CommandEncoder,
        texture: &wgpu::Texture,
    ) {
        profiling::scope!("GpuMsdfWriter::write_glyphs");
        if self.layers.iter().all(|l| l.indices.is_empty()) {
            return;
        }

        debug_assert!(
            self.layers.len() as u32 <= texture.depth_or_array_layers(),
            "AtlasRegion.layer ({}) exceeds texture layer count ({})",
            self.layers.len() - 1,
            texture.depth_or_array_layers(),
        );

        // Convert atlas-pixel `clip_pos` to NDC now that the atlas size is known.
        let tw = texture.width() as f32;
        let th = texture.height() as f32;
        for layer in &mut self.layers {
            for v in &mut layer.vertices {
                let px = v.clip_pos[0];
                let py = v.clip_pos[1];
                v.clip_pos = [(px / tw) * 2.0 - 1.0, 1.0 - (py / th) * 2.0];
            }
        }

        // Concatenate vertex/index data, recording per-layer offsets.
        let mut vertices_all: Vec<Vertex> = Vec::new();
        let mut indices_all: Vec<u32> = Vec::new();
        // (layer_index, base_vertex, index_start, index_count)
        let mut layer_draws: Vec<(u32, i32, u32, u32)> = Vec::new();
        for (layer_idx, layer) in self.layers.iter().enumerate() {
            if layer.indices.is_empty() {
                continue;
            }
            let base_vertex = vertices_all.len() as i32;
            let index_start = indices_all.len() as u32;
            vertices_all.extend_from_slice(&layer.vertices);
            indices_all.extend_from_slice(&layer.indices);
            let index_count = layer.indices.len() as u32;
            layer_draws.push((layer_idx as u32, base_vertex, index_start, index_count));
        }

        // Upload segment buffer.
        let seg_bytes: &[u8] = bytemuck::cast_slice(&self.segments);
        let seg_size = seg_bytes.len() as u64;
        if self.segment_buf.is_none() || self.segment_cap < seg_size {
            let new_cap = seg_size.next_power_of_two().max(256);
            self.segment_buf = Some(device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("klyff_msdf segment storage"),
                size: new_cap,
                usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            }));
            self.segment_cap = new_cap;
            self.bind_group = None;
        }
        queue.write_buffer(self.segment_buf.as_ref().unwrap(), 0, seg_bytes);

        // Upload vertex/index buffers.
        let v_bytes: &[u8] = bytemuck::cast_slice(&vertices_all);
        let v_size = v_bytes.len() as u64;
        if self.vertex_buf.is_none() || self.vertex_cap < v_size {
            let new_cap = v_size.next_power_of_two().max(256);
            self.vertex_buf = Some(device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("klyff_msdf vertices"),
                size: new_cap,
                usage: wgpu::BufferUsages::VERTEX | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            }));
            self.vertex_cap = new_cap;
        }
        queue.write_buffer(self.vertex_buf.as_ref().unwrap(), 0, v_bytes);

        let i_bytes: &[u8] = bytemuck::cast_slice(&indices_all);
        let i_size = i_bytes.len() as u64;
        if self.index_buf.is_none() || self.index_cap < i_size {
            let new_cap = i_size.next_power_of_two().max(256);
            self.index_buf = Some(device.create_buffer(&wgpu::BufferDescriptor {
                label: Some("klyff_msdf indices"),
                size: new_cap,
                usage: wgpu::BufferUsages::INDEX | wgpu::BufferUsages::COPY_DST,
                mapped_at_creation: false,
            }));
            self.index_cap = new_cap;
        }
        queue.write_buffer(self.index_buf.as_ref().unwrap(), 0, i_bytes);

        if self.bind_group.is_none() {
            self.bind_group = Some(device.create_bind_group(&wgpu::BindGroupDescriptor {
                label: Some("klyff_msdf bg"),
                layout: &self.bgl,
                entries: &[wgpu::BindGroupEntry {
                    binding: 0,
                    resource: self.segment_buf.as_ref().unwrap().as_entire_binding(),
                }],
            }));
        }

        let vertex_buf = self.vertex_buf.as_ref().unwrap();
        let index_buf = self.index_buf.as_ref().unwrap();
        let bind_group = self.bind_group.as_ref().unwrap();

        for (layer, base_vertex, index_start, index_count) in layer_draws {
            let view = texture.create_view(&wgpu::TextureViewDescriptor {
                label: Some("klyff_msdf atlas layer view"),
                dimension: Some(wgpu::TextureViewDimension::D2),
                base_array_layer: layer,
                array_layer_count: Some(1),
                ..Default::default()
            });
            let mut pass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
                label: Some("klyff_msdf write pass"),
                color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                    view: &view,
                    depth_slice: None,
                    resolve_target: None,
                    ops: wgpu::Operations {
                        load: wgpu::LoadOp::Load,
                        store: wgpu::StoreOp::Store,
                    },
                })],
                depth_stencil_attachment: None,
                timestamp_writes: None,
                occlusion_query_set: None,
            });
            pass.set_pipeline(&self.pipeline);
            pass.set_bind_group(0, bind_group, &[]);
            pass.set_vertex_buffer(0, vertex_buf.slice(..));
            pass.set_index_buffer(index_buf.slice(..), wgpu::IndexFormat::Uint32);
            pass.draw_indexed(index_start..index_start + index_count, base_vertex, 0..1);
        }

        self.segments.clear();
        for l in &mut self.layers {
            l.vertices.clear();
            l.indices.clear();
        }
    }
}

fn push_segments(out: &mut Vec<SegmentGpu>, segs: &[Segment], soa: &SegmentSoa) {
    out.reserve(segs.len());
    for (i, seg) in segs.iter().enumerate() {
        out.push(SegmentGpu {
            from: seg.from.to_array(),
            to: seg.to.to_array(),
            mid: seg.mid.to_array(),
            tangent: soa.tangent[i].to_array(),
            tangent_unit: soa.tangent_unit[i].to_array(),
            inv_tangent_len_sq: soa.inv_tangent_len_sq[i],
            color_mask: soa.color_mask[i] as u32,
        });
    }
}

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

    #[test]
    fn shader_compiles() {
        let instance = wgpu::Instance::default();
        let Some(adapter) =
            pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
                power_preference: wgpu::PowerPreference::HighPerformance,
                force_fallback_adapter: false,
                compatible_surface: None,
            }))
            .ok()
        else {
            eprintln!("no wgpu adapter available; skipping shader_compiles");
            return;
        };
        let (device, _queue) =
            pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor {
                label: Some("klyff_msdf shader compile test device"),
                required_features: wgpu::Features::empty(),
                required_limits: wgpu::Limits::default(),
                memory_hints: wgpu::MemoryHints::Performance,
                trace: wgpu::Trace::Off,
                experimental_features: wgpu::ExperimentalFeatures::disabled(),
            }))
            .expect("request_device");

        device.push_error_scope(wgpu::ErrorFilter::Validation);
        let _writer = MtsdfGpuWriter::new(&device);
        let err = pollster::block_on(device.pop_error_scope());
        assert!(
            err.is_none(),
            "shader/pipeline failed validation: {:?}",
            err
        );
    }
}