latkit 0.2.1

//! Links every WESL entry, validates it with naga, and writes `OUT_DIR/shaders.rs`:
//!
//! - every constant and struct of `lib/types.wesl`, as Rust constants and `#[repr(C)]` mirrors whose
//!   size is checked against WGSL;
//! - the shared bind group layouts of `lib/bind.wesl` (group 0 per view, group 1 per mark);
//! - per shader and variant: its declarations, entry points, the layout of its pass group (group 2),
//!   and whether it reads the shared groups at all.
//!
//! Variants share most of what they link, so each distinct top-level declaration is stored once, in
//! `OUT_DIR/shaders.wgsl`, and a variant is the list of its declarations, which join back into exactly
//! the module it linked.
//!
//! It also writes each variant's whole module and `OUT_DIR/shaders.json`, every variant's file, entry
//! points, and the pipeline constants it is made with, from which the browser gate times each
//! pipeline in Chrome.

use std::collections::{BTreeMap, HashMap};
use std::env;
use std::fmt::Write as _;
use std::fs;
use std::path::{Path, PathBuf};

use naga::valid::{Capabilities, ModuleInfo, ValidationFlags, Validator};
use naga::{AddressSpace, ArraySize, ImageClass, ImageDimension, ScalarKind, StorageAccess, TypeInner, VectorSize};

/// A shader module and the variants it is linked in: each variant turns on its WESL features.
struct Entry {
    module: &'static str,
    variants: &'static [(&'static str, &'static [&'static str])],
}

const ENTRIES: &[Entry] = &[
    Entry { module: "compute::evaluate", variants: &[("main", &["writes"])] },
    Entry { module: "compute::freeze", variants: &[("main", &[])] },
    Entry { module: "compute::bounds", variants: &[("main", &[])] },
    Entry { module: "compute::normals", variants: &[("main", &[])] },
    Entry { module: "compute::probe", variants: &[("main", &[]), ("multisampled", &["multisampled"])] },
    Entry { module: "compute::lists", variants: &[("main", &["lists"])] },
    Entry { module: "draw::backdrop", variants: &[("main", &[])] },
    Entry {
        module: "draw::points",
        variants: &[
            ("main", &[]),
            ("cast", &["casts"]),
            ("custom", &["custom"]),
            ("custom_cast", &["custom", "casts"]),
            ("custom_focus", &["custom", "focused"]),
            ("ports", &["ports"]),
            ("focus", &["focused"]),
        ],
    },
    Entry { module: "draw::strokes", variants: &[("main", &[]), ("cast", &["casts"]), ("focus", &["focused"])] },
    Entry { module: "draw::triangles", variants: &[("main", &[]), ("cast", &["casts"])] },
    Entry { module: "draw::areas", variants: &[("main", &[]), ("focus", &["focused"])] },
    Entry { module: "draw::shade", variants: &[("main", &[])] },
    Entry { module: "draw::glyphs", variants: &[("main", &[]), ("ports", &["ports"])] },
    Entry { module: "compute::candidates", variants: &[("main", &[])] },
    Entry { module: "compute::labels", variants: &[("main", &[])] },
    Entry { module: "draw::overlay", variants: &[("main", &[])] },
    Entry { module: "compute::compose", variants: &[("main", &[])] },
    Entry { module: "compute::fields", variants: &[("main", &[])] },
];

const FEATURES: &[&str] = &["writes", "casts", "custom", "multisampled", "lists", "ports", "focused"];

/// What a custom marker supplies, for validating its variant.
const CUSTOM_MARKER: &str = "fn marker(f: MarkerFragment) -> MarkerColor { return MarkerColor(f.color, 0.0); }";

/// Uniforms bound at a dynamic offset: each mark's, each pass's eye, and each probe dispatch's step.
const DYNAMIC: &[&str] = &["mark", "eye", "step"];

/// The values each pipeline constant is made with, as constants of `lib/types.wesl` or booleans: the
/// primitive a probe or a label's anchor reads, the path a stroke reads, and whether channels read
/// joins, samples, and the view.
const OVERRIDES: &[(&str, &[&str])] = &[
    ("PRIMITIVE", &["POINTS", "STROKES", "TRIANGLES"]),
    ("PATH", &["PATH_ENDS", "PATH_TRACE", "PATH_GEOMETRY", "PATH_POLES"]),
    ("JOINED", &["false", "true"]),
    ("TIMED", &["false", "true"]),
    ("VIEWED", &["false", "true"]),
];

pub(crate) fn generate() {
    let root = Path::new(&env::var("CARGO_MANIFEST_DIR").unwrap()).join("shaders");
    let out = PathBuf::from(env::var("OUT_DIR").unwrap());
    let mut rust = String::from("// Generated by build/shaders.rs from the WESL shaders. Do not edit.\n\n");

    let types = link(&root, "lib::types", &[], true);
    let (types_module, _) = validate("lib::types", &types);
    mirror_consts(&types_module, &mut rust);
    mirror_structs(&types_module, &mut rust);

    let bind = link(&root, "lib::bind", &[], true);
    let (bind_module, _) = validate("lib::bind", &bind);
    for (group, name) in [(0, "VIEW_GROUP"), (1, "MARK_GROUP")] {
        let bindings = bindings(&bind_module, None, group);
        writeln!(rust, "pub(crate) static {name}: &[Binding] = &[{}];\n", bindings.join(", ")).unwrap();
    }

    let mut shaders = String::new();
    let mut sources = String::new();
    let mut manifest = Vec::new();
    let mut pieces = Pieces::default();
    for entry in ENTRIES {
        for &(variant, features) in entry.variants {
            let source = link(&root, entry.module, features, false);
            // A custom marker's variant lacks the marker an application supplies; a stand-in validates it.
            let custom = features.contains(&"custom");
            let checked = if custom { format!("{source}\n{CUSTOM_MARKER}") } else { source.clone() };
            let (module, info) = validate(&format!("{}.{variant}", entry.module), &checked);
            check_shared(&module, &bind_module, entry.module);
            let file = format!("{}.{variant}.wgsl", entry.module.replace("::", "."));
            fs::write(out.join(&file), &source).unwrap();
            manifest.push(manifest_entry(entry.module, &file, &module, &types_module, custom));
            let name = shader_name(entry.module, variant);
            writeln!(shaders, "    {name},").unwrap();
            let pass = bindings(&module, Some(&info), 2);
            let overrides: Vec<String> = module.overrides.iter().filter_map(|(_, o)| o.name.clone()).collect();
            let shared = module.global_variables.iter().any(|(_, g)| g.binding.as_ref().is_some_and(|b| b.group < 2));
            writeln!(
                sources,
                "        Shader::{name} => &ShaderSource {{ label: {label:?}, pieces: &{refs:?}, pass: &[{pass}], overrides: &{overrides:?}, shared: {shared} }},",
                label = format!("latkit {}.{variant}", entry.module),
                refs = pieces.add(&source),
                pass = pass.join(", "),
            )
            .unwrap();
        }
    }
    fs::write(out.join("shaders.wgsl"), &pieces.text).unwrap();
    writeln!(
        rust,
        "/// Where each declaration lies in `shaders.wgsl`.\npub(crate) static PIECES: &[(u32, u32)] = &{:?};\n",
        pieces.ranges
    )
    .unwrap();
    writeln!(rust, "/// Every linked shader and variant.\n#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]\npub(crate) enum Shader {{\n{shaders}}}\n").unwrap();
    writeln!(rust, "impl Shader {{\n    pub(crate) fn source(self) -> &'static ShaderSource {{\n        match self {{\n{sources}        }}\n    }}\n}}").unwrap();
    fs::write(out.join("shaders.rs"), rust).unwrap();
    fs::write(out.join("shaders.json"), format!("[\n{}\n]\n", manifest.join(",\n"))).unwrap();
}

/// Top-level declarations, each stored once: their text back to back, and where each lies.
#[derive(Default)]
struct Pieces {
    text: String,
    ranges: Vec<(u32, u32)>,
    index: HashMap<String, u16>,
}

impl Pieces {
    /// The declarations of `source`, as indices, storing those not stored yet; they join back into
    /// `source` exactly.
    fn add(&mut self, source: &str) -> Vec<u16> {
        let refs: Vec<u16> = source
            .split(SEPARATOR)
            .map(|piece| {
                *self.index.entry(piece.to_owned()).or_insert_with(|| {
                    let start = self.text.len() as u32;
                    self.text.push_str(piece);
                    self.ranges.push((start, self.text.len() as u32));
                    u16::try_from(self.ranges.len() - 1).expect("fewer than 65536 declarations")
                })
            })
            .collect();
        let joined: Vec<&str> = refs
            .iter()
            .map(|&r| &self.text[self.ranges[r as usize].0 as usize..self.ranges[r as usize].1 as usize])
            .collect();
        assert_eq!(joined.join(SEPARATOR), source, "a variant's declarations join back into the module it linked");
        refs
    }
}

/// What separates top-level declarations in a linked module, and joins them back.
const SEPARATOR: &str = "\n\n";

/// A variant as JSON: its file, whether it draws, its compute entry points, every combination of the
/// pipeline constants it declares, and the WGSL appended to make it whole (a custom marker's stand-in).
fn manifest_entry(shader: &str, file: &str, module: &naga::Module, types: &naga::Module, custom: bool) -> String {
    let compute: Vec<String> = module
        .entry_points
        .iter()
        .filter(|e| e.stage == naga::ShaderStage::Compute)
        .map(|e| format!("{:?}", e.name))
        .collect();
    let mut combinations: Vec<Vec<(&str, u32)>> = vec![Vec::new()];
    for (_, o) in module.overrides.iter() {
        let name = o.name.as_deref().expect("a named override");
        let (_, values) = OVERRIDES.iter().find(|(n, _)| *n == name).unwrap_or_else(|| panic!("no values for {name}"));
        combinations = combinations
            .iter()
            .flat_map(|before| {
                values.iter().map(move |value| {
                    let mut next = before.clone();
                    let value = match *value {
                        "false" => 0,
                        "true" => 1,
                        name => type_constant(types, name),
                    };
                    next.push((name, value));
                    next
                })
            })
            .filter(|c| made(shader, c, types))
            .collect();
    }
    let constants: Vec<String> = combinations
        .iter()
        .map(|c| format!("{{{}}}", c.iter().map(|(n, v)| format!("{n:?}: {v}")).collect::<Vec<_>>().join(", ")))
        .collect();
    format!(
        "  {{ \"file\": {file:?}, \"draws\": {draws}, \"compute\": [{}], \"constants\": [{}], \"append\": {:?} }}",
        compute.join(", "),
        constants.join(", "),
        if custom { CUSTOM_MARKER } else { "" },
        draws = compute.is_empty(),
    )
}

/// Whether a pipeline of `module` is ever made with `constants`: only channels go without joins or
/// samples; faces read their ends and areas their geometry; a pipeline made for points reads no path,
/// so it keeps the default one, and one made for triangles reads a face's ends or an area's geometry.
fn made(module: &str, constants: &[(&str, u32)], types: &naga::Module) -> bool {
    let value = |name: &str| constants.iter().find(|(n, _)| *n == name).map(|&(_, v)| v);
    if module != "compute::evaluate" && (value("JOINED") == Some(0) || value("TIMED") == Some(0)) {
        return false;
    }
    let path = value("PATH");
    match module {
        "draw::triangles" => return path == Some(type_constant(types, "PATH_ENDS")),
        "draw::areas" => return path == Some(type_constant(types, "PATH_GEOMETRY")),
        _ => {}
    }
    match (value("PRIMITIVE"), path) {
        (Some(primitive), Some(path)) => {
            primitive == type_constant(types, "STROKES")
                || path == type_constant(types, "PATH_ENDS")
                || (primitive == type_constant(types, "TRIANGLES") && path == type_constant(types, "PATH_GEOMETRY"))
        }
        _ => true,
    }
}

/// The value of `lib/types.wesl`'s `u32` constant `name`.
fn type_constant(types: &naga::Module, name: &str) -> u32 {
    let (_, constant) = types.constants.iter().find(|(_, c)| c.name.as_deref() == Some(name)).expect("a constant");
    match types.global_expressions[constant.init] {
        naga::Expression::Literal(naga::Literal::U32(v)) => v,
        _ => panic!("{name} is not a u32"),
    }
}

fn shader_name(module: &str, variant: &str) -> String {
    let words = module.split("::").skip(1).chain((variant != "main").then_some(variant));
    words
        .flat_map(|w| w.split('_'))
        .map(|w| {
            let mut chars = w.chars();
            chars.next().map(|c| c.to_ascii_uppercase().to_string() + chars.as_str()).unwrap_or_default()
        })
        .collect()
}

fn link(root: &Path, module: &str, features: &[&str], keep_main: bool) -> String {
    let mut options =
        wesl::CompileOptions { mangler: wesl::ManglerKind::None, validate: false, keep_main, ..Default::default() };
    for feature in FEATURES {
        options.features.set(*feature, features.contains(feature));
    }
    let path = format!("package::{module}").parse().unwrap();
    wesl::Compiler::new(options)
        .compile_module(root, &path)
        .unwrap_or_else(|e| panic!("WESL {module} {features:?}: {e}"))
        .to_string()
}

fn validate(name: &str, source: &str) -> (naga::Module, ModuleInfo) {
    let module = naga::front::wgsl::parse_str(source)
        .unwrap_or_else(|e| panic!("WGSL {name}:\n{}\n{source}", e.emit_to_string(source)));
    let info = Validator::new(ValidationFlags::all(), Capabilities::empty())
        .validate(&module)
        .unwrap_or_else(|e| panic!("WGSL {name}:\n{}", e.emit_to_string(source)));
    (module, info)
}

/// Every `const` of a scalar type, as a Rust constant.
fn mirror_consts(module: &naga::Module, rust: &mut String) {
    for (_, constant) in module.constants.iter() {
        let Some(name) = &constant.name else { continue };
        let value = match module.global_expressions[constant.init] {
            naga::Expression::Literal(naga::Literal::U32(v)) => format!("u32 = {v:#x}"),
            naga::Expression::Literal(naga::Literal::I32(v)) => format!("i32 = {v}"),
            naga::Expression::Literal(naga::Literal::F32(v)) => format!("f32 = {v:?}"),
            _ => continue,
        };
        writeln!(rust, "#[allow(dead_code)]\npub(crate) const {name}: {value};").unwrap();
    }
    rust.push('\n');
}

/// Every struct as a `#[repr(C)]` mirror with explicit padding.
fn mirror_structs(module: &naga::Module, rust: &mut String) {
    for (_, ty) in module.types.iter() {
        let (Some(name), TypeInner::Struct { members, span }) = (&ty.name, &ty.inner) else { continue };
        writeln!(
            rust,
            "#[repr(C)]\n#[derive(Clone, Copy, Debug, bytemuck::Pod, bytemuck::Zeroable)]\npub(crate) struct {name} {{"
        )
        .unwrap();
        let mut at = 0;
        let mut pads = 0;
        for member in members {
            pad(rust, &mut at, member.offset, &mut pads);
            let (field_type, size) = rust_type(module, member.ty);
            writeln!(rust, "    pub(crate) {}: {field_type},", member.name.as_deref().unwrap()).unwrap();
            at = member.offset + size;
        }
        pad(rust, &mut at, *span, &mut pads);
        writeln!(rust, "}}\n\nconst _: () = assert!(size_of::<{name}>() == {span});\n").unwrap();
    }
}

fn pad(rust: &mut String, at: &mut u32, to: u32, pads: &mut u32) {
    if to > *at {
        assert_eq!((to - *at) % 4, 0, "padding is whole words");
        writeln!(rust, "    pub(crate) _pad{pads}: [u32; {}],", (to - *at) / 4).unwrap();
        *pads += 1;
        *at = to;
    }
}

fn rust_type(module: &naga::Module, ty: naga::Handle<naga::Type>) -> (String, u32) {
    let scalar = |s: naga::Scalar| match (s.kind, s.width) {
        (ScalarKind::Float, 4) => "f32",
        (ScalarKind::Uint, 4) => "u32",
        (ScalarKind::Sint, 4) => "i32",
        other => panic!("no mirror for scalar {other:?}"),
    };
    let size = |v: VectorSize| v as u32;
    match &module.types[ty].inner {
        TypeInner::Scalar(s) => (scalar(*s).into(), 4),
        TypeInner::Vector { size: n, scalar: s } => (format!("[{}; {}]", scalar(*s), size(*n)), 4 * size(*n)),
        TypeInner::Matrix { columns, rows, scalar: s } => {
            assert_eq!(*rows, VectorSize::Quad, "matrix columns are padded to four rows");
            (format!("[[{}; 4]; {}]", scalar(*s), size(*columns)), 16 * size(*columns))
        }
        TypeInner::Array { base, size: ArraySize::Constant(n), stride } => {
            let (element, element_size) = rust_type(module, *base);
            assert_eq!(element_size, *stride, "array elements are not padded");
            (format!("[{element}; {n}]"), stride * n.get())
        }
        TypeInner::Struct { span, .. } => (module.types[ty].name.clone().expect("a named struct"), *span),
        other => panic!("no mirror for {other:?}"),
    }
}

/// The bindings of `group` as `Binding` literals; with `info`, visibility is where entry points use them.
fn bindings(module: &naga::Module, info: Option<&ModuleInfo>, group: u32) -> Vec<String> {
    let mut found = BTreeMap::new();
    for (handle, global) in module.global_variables.iter() {
        let Some(binding) = &global.binding else { continue };
        if binding.group != group {
            continue;
        }
        let stages = match info {
            Some(info) => module
                .entry_points
                .iter()
                .enumerate()
                .filter(|(i, _)| !info.get_entry_point(*i)[handle].is_empty())
                .map(|(_, e)| match e.stage {
                    naga::ShaderStage::Vertex => "VERTEX",
                    naga::ShaderStage::Fragment => "FRAGMENT",
                    naga::ShaderStage::Compute => "COMPUTE",
                    other => panic!("unexpected stage {other:?}"),
                })
                .collect::<Vec<_>>(),
            None => vec!["VERTEX", "FRAGMENT", "COMPUTE"],
        };
        let kind = match (global.space, &module.types[global.ty].inner) {
            (AddressSpace::Uniform, _) => "Kind::Uniform".to_owned(),
            (AddressSpace::Storage { access }, _) => {
                format!("Kind::Storage {{ read_only: {} }}", !access.contains(StorageAccess::STORE))
            }
            (AddressSpace::Handle, TypeInner::Sampler { .. }) => "Kind::Sampler".to_owned(),
            (
                AddressSpace::Handle,
                TypeInner::Image { dim: ImageDimension::D2, class: ImageClass::Depth { multi }, .. },
            ) => {
                format!("Kind::Depth {{ multisampled: {multi} }}")
            }
            (AddressSpace::Handle, TypeInner::Image { dim: ImageDimension::D2, arrayed, class }) => {
                let filterable = matches!(class, ImageClass::Sampled { kind: ScalarKind::Float, multi: false });
                format!("Kind::Texture {{ array: {arrayed}, filterable: {filterable} }}")
            }
            other => panic!("unsupported binding {:?}: {other:?}", global.name),
        };
        let mut stages = stages;
        stages.sort_unstable();
        stages.dedup();
        let stages = if stages.is_empty() { "0".to_owned() } else { stages.join(" | ") };
        let dynamic = global.name.as_deref().is_some_and(|name| DYNAMIC.contains(&name));
        let entry =
            format!("Binding {{ slot: {}, kind: {kind}, stages: {stages}, dynamic: {dynamic} }}", binding.binding);
        assert!(found.insert(binding.binding, entry).is_none(), "binding {group}.{} twice", binding.binding);
    }
    found.into_values().collect()
}

/// Every group 0 and 1 binding a shader declares must match the shared declaration.
fn check_shared(module: &naga::Module, shared: &naga::Module, name: &str) {
    for (_, global) in module.global_variables.iter() {
        let Some(binding) = &global.binding else { continue };
        if binding.group > 1 {
            continue;
        }
        let same = shared
            .global_variables
            .iter()
            .any(|(_, s)| s.binding.as_ref() == Some(binding) && s.name == global.name && s.space == global.space);
        assert!(
            same,
            "{name}: {:?} at group {} binding {} is not the shared declaration",
            global.name, binding.group, binding.binding
        );
    }
}