mirage-engine-derive 0.2.0

//! The derive macros behind Mirage's vocabulary traits, which the engine
//! makes public; nothing here depends on the engine itself.

use proc_macro::TokenStream;

use quote::{format_ident, quote};
use syn::punctuated::Punctuated;
use syn::{
    Attribute, Data, DeriveInput, Expr, Fields, Ident, LitStr, Path, Token, parse_macro_input,
    parse_quote,
};

/// The block a shader reads a set of values in: what each of them starts
/// on, and what the whole of them fills.
const BLOCK: usize = 16;

/// The parts of a mesh, named by the materials a source states.
const PART: Named = Named {
    trait_name: "Part",
    what: "part",
};

/// The clips of a mesh, named by the animations a source states.
const CLIP: Named = Named {
    trait_name: "Clip",
    what: "clip",
};

/// Derives `Catalog` for a mesh vocabulary: every value the engine builds
/// before the game's startup closure runs, to prove the assets it names
/// are loaded.
///
/// Fieldless variants catalog themselves. A variant with fields needs one
/// `#[catalog(...)]` naming the values to prove, each written as a value of
/// the type, and omitting it fails to compile:
/// `#[catalog(Self::Asteroid { seed: 1 }, Self::Asteroid { seed: 7 })]`.
/// Wrap a primitive in a variant's field only when each of its values is its
/// own mesh; a variant that draws one fixed primitive stays fieldless and
/// builds it in `build`.
#[proc_macro_derive(Catalog, attributes(catalog))]
pub fn derive_catalog(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    match catalog_impl(&input) {
        Ok(implementation) => implementation,
        Err(error) => error.to_compile_error().into(),
    }
}

fn catalog_impl(input: &DeriveInput) -> syn::Result<TokenStream> {
    let name = &input.ident;
    let values = match &input.data {
        Data::Enum(data) => data
            .variants
            .iter()
            .map(|variant| {
                let variant_name = &variant.ident;
                cataloged(
                    &parse_quote!(#name::#variant_name),
                    &variant.fields,
                    &variant.attrs,
                    variant_name,
                )
            })
            .collect::<syn::Result<Vec<_>>>()?
            .concat(),
        Data::Struct(data) => cataloged(&parse_quote!(#name), &data.fields, &input.attrs, name)?,
        Data::Union(_) => {
            return Err(syn::Error::new_spanned(
                name,
                "Catalog covers enums and structs; a union needs the impl written by hand",
            ));
        }
    };

    let (impl_generics, type_generics, where_clause) = input.generics.split_for_impl();
    Ok(quote! {
        impl #impl_generics ::mirage_engine::Catalog for #name #type_generics #where_clause {
            fn catalog() -> ::std::vec::Vec<Self> {
                ::std::vec![#(#values),*]
            }
        }
    }
    .into())
}

/// The values `#[catalog(…)]` names for one variant or struct, or the value
/// it is on its own when it has no fields.
fn cataloged(
    path: &Path,
    fields: &Fields,
    attributes: &[Attribute],
    name: &Ident,
) -> syn::Result<Vec<Expr>> {
    let mut declared = attributes
        .iter()
        .filter(|attribute| attribute.path().is_ident("catalog"));
    let Some(attribute) = declared.next() else {
        return match fields {
            Fields::Unit => Ok(vec![parse_quote!(#path)]),
            _ => Err(syn::Error::new_spanned(
                name,
                format!(
                    "`{name}` has fields, so each value of it is a mesh of its own: name a value \
                     of the type in `#[catalog({}, …)]`",
                    representative(path, fields)
                ),
            )),
        };
    };
    if let Some(extra) = declared.next() {
        return Err(syn::Error::new_spanned(
            extra,
            format!(
                "`{name}` names its values in one `#[catalog(…)]`; drop the attribute past the \
                 first"
            ),
        ));
    }
    if matches!(fields, Fields::Unit) {
        return Err(syn::Error::new_spanned(
            attribute,
            format!("`{name}` has no fields, so it catalogs itself; drop the attribute"),
        ));
    }

    let values: Vec<Expr> = attribute
        .parse_args_with(Punctuated::<Expr, Token![,]>::parse_terminated)?
        .into_iter()
        .collect();
    match values.is_empty() {
        true => Err(syn::Error::new_spanned(
            attribute,
            format!(
                "`{name}` names no value; give the attribute a value of the type, `{}`",
                representative(path, fields)
            ),
        )),
        false => Ok(values),
    }
}

/// One value of the item, written the way the attribute would name it, for
/// the error that requests the attribute.
fn representative(path: &Path, fields: &Fields) -> String {
    let spelled = path
        .segments
        .iter()
        .map(|segment| segment.ident.to_string())
        .collect::<Vec<_>>()
        .join("::");
    let each: Vec<String> = fields
        .iter()
        .map(|field| match &field.ident {
            Some(field) => format!("{field}: …"),
            None => "…".to_owned(),
        })
        .collect();
    match fields {
        Fields::Named(_) => format!("{spelled} {{ {} }}", each.join(", ")),
        _ => format!("{spelled}({})", each.join(", ")),
    }
}

/// Derives `Part` for a vocabulary of mesh parts: each fieldless variant
/// (and a unit struct) is named by what it is called in code, or by
/// `#[part("...")]` when the loaded material name cannot be written as one
/// — repeat the attribute for other spellings. The variants are numbered
/// in order, which is the position a draw's override of one is stored at.
///
/// A variant with fields fails to compile: a part is a plain name.
#[proc_macro_derive(Part, attributes(part))]
pub fn derive_part(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    match named_impl(&input, &PART) {
        Ok(implementation) => implementation,
        Err(error) => error.to_compile_error().into(),
    }
}

/// Derives `Clip` for a vocabulary of the clips a mesh is posed by: each
/// fieldless variant (and a unit struct) is named by what it is called in
/// code, or by `#[clip("...")]` when the loaded animation name cannot be
/// written as one — repeat the attribute for other spellings. The variants
/// are numbered in order.
///
/// A variant with fields fails to compile: a clip is a plain name.
#[proc_macro_derive(Clip, attributes(clip))]
pub fn derive_clip(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    match named_impl(&input, &CLIP) {
        Ok(implementation) => implementation,
        Err(error) => error.to_compile_error().into(),
    }
}

/// A vocabulary whose every value is one name a loaded source states: the
/// trait to write, and the attribute and the word its errors spell it with.
struct Named {
    trait_name: &'static str,
    what: &'static str,
}

fn named_impl(input: &DeriveInput, named: &Named) -> syn::Result<TokenStream> {
    let name = &input.ident;
    let parts = match &input.data {
        Data::Enum(data) => data
            .variants
            .iter()
            .map(|variant| {
                let variant_name = &variant.ident;
                spellings(
                    &parse_quote!(#name::#variant_name),
                    &variant.fields,
                    &variant.attrs,
                    variant_name,
                    named,
                )
            })
            .collect::<syn::Result<Vec<_>>>()?,
        Data::Struct(data) => vec![spellings(
            &parse_quote!(#name),
            &data.fields,
            &input.attrs,
            name,
            named,
        )?],
        Data::Union(_) => {
            let trait_name = named.trait_name;
            return Err(syn::Error::new_spanned(
                name,
                format!(
                    "{trait_name} covers enums and unit structs; a union needs the impl written \
                     by hand"
                ),
            ));
        }
    };

    let every = parts.iter().map(|(key, _)| key);
    let indices = 0u32..parts.len() as u32;
    let indexed = parts.iter().map(|(key, _)| key);
    let (modelled, keys): (Vec<&LitStr>, Vec<&Expr>) = parts
        .iter()
        .flat_map(|(key, names)| names.iter().map(move |spelling| (spelling, key)))
        .unzip();

    let trait_name = format_ident!("{}", named.trait_name);
    let (impl_generics, type_generics, where_clause) = input.generics.split_for_impl();
    Ok(quote! {
        impl #impl_generics ::mirage_engine::#trait_name for #name #type_generics #where_clause {
            fn from_name(name: &str) -> ::core::option::Option<Self> {
                match name {
                    #(#modelled => ::core::option::Option::Some(#keys),)*
                    _ => ::core::option::Option::None,
                }
            }

            fn all() -> ::std::vec::Vec<Self> {
                ::std::vec![#(#every),*]
            }

            fn index(&self) -> u32 {
                match self {
                    #(#indexed => #indices,)*
                }
            }
        }
    }
    .into())
}

/// One variant or struct: the value it is, and the names that resolve to
/// it — whatever `#[part("...")]` or `#[clip("...")]` spells, or what it is
/// called in code.
fn spellings(
    path: &Path,
    fields: &Fields,
    attributes: &[Attribute],
    name: &Ident,
    named: &Named,
) -> syn::Result<(Expr, Vec<LitStr>)> {
    let what = named.what;
    if !matches!(fields, Fields::Unit) {
        return Err(syn::Error::new_spanned(
            name,
            format!("`{name}` has fields, but a {what} is a plain name; give it none"),
        ));
    }

    let declared: Vec<LitStr> = attributes
        .iter()
        .filter(|attribute| attribute.path().is_ident(what))
        .map(Attribute::parse_args)
        .collect::<syn::Result<_>>()?;

    let names = match declared.is_empty() {
        true => vec![LitStr::new(&name.to_string(), name.span())],
        false => declared,
    };
    Ok((parse_quote!(#path), names))
}

/// Derives the engine's view of a vocabulary of actions whose value is held
/// or not: its values, and the names a rebind of one is kept under.
///
/// The bindings themselves are declared by hand in the `InputButtonAction`
/// trait, which the derived code calls; a variant with fields fails to compile,
/// because an action is a verb and its parameters belong to the game's own
/// state.
#[proc_macro_derive(InputButtonAction)]
pub fn derive_input_button_action(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    emit(
        &input,
        &parse_quote!(::mirage_engine::ButtonBinding),
        &parse_quote!(::mirage_engine::InputButtonAction),
        "NoInputButtons",
    )
}

/// Derives the engine's view of a vocabulary of actions whose value is a
/// number, whose bindings are declared by hand in the `InputAxisAction` trait; see
/// [`InputButtonAction`](macro@InputButtonAction).
#[proc_macro_derive(InputAxisAction)]
pub fn derive_input_axis_action(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    emit(
        &input,
        &parse_quote!(::mirage_engine::AxisBinding),
        &parse_quote!(::mirage_engine::InputAxisAction),
        "NoInputAxes",
    )
}

/// Derives the engine's view of a vocabulary of actions whose value is a
/// vector, whose bindings are declared by hand in the `InputAxis2Action`
/// trait; see
/// [`InputButtonAction`](macro@InputButtonAction).
#[proc_macro_derive(InputAxis2Action)]
pub fn derive_input_axis2_action(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    emit(
        &input,
        &parse_quote!(::mirage_engine::Axis2Binding),
        &parse_quote!(::mirage_engine::InputAxis2Action),
        "NoInputAxes2",
    )
}

fn emit(input: &DeriveInput, binding: &Path, kind: &Path, empty: &str) -> TokenStream {
    match actions(input, binding, kind, empty) {
        Ok(implementation) => implementation,
        Err(error) => error.to_compile_error().into(),
    }
}

fn actions(
    input: &DeriveInput,
    binding: &Path,
    kind: &Path,
    empty: &str,
) -> syn::Result<TokenStream> {
    let name = &input.ident;
    let verbs: Vec<(Path, &Ident)> = match &input.data {
        Data::Enum(data) => data
            .variants
            .iter()
            .map(|variant| {
                let variant_name = &variant.ident;
                verb(
                    parse_quote!(#name::#variant_name),
                    &variant.fields,
                    variant_name,
                )
            })
            .collect::<syn::Result<_>>()?,
        Data::Struct(data) => vec![verb(parse_quote!(#name), &data.fields, name)?],
        Data::Union(_) => {
            return Err(syn::Error::new_spanned(
                name,
                "an action vocabulary is an enum or a unit struct; a union needs the impl \
                 written by hand",
            ));
        }
    };
    if verbs.is_empty() {
        return Err(syn::Error::new_spanned(
            name,
            format!("`{name}` names no action; the vocabulary of none is `mirage_engine::{empty}`"),
        ));
    }

    let (paths, idents): (Vec<&Path>, Vec<&Ident>) =
        verbs.iter().map(|(path, ident)| (path, *ident)).unzip();
    let (impl_generics, type_generics, where_clause) = input.generics.split_for_impl();
    Ok(quote! {
        impl #impl_generics ::mirage_engine::InputAction for #name #type_generics #where_clause {
            type Binding = #binding;

            fn defaults(&self) -> ::std::vec::Vec<#binding> {
                <Self as #kind>::bindings(self)
            }

            fn all() -> ::std::vec::Vec<Self> {
                ::std::vec![#(#paths),*]
            }

            fn name(&self) -> &'static str {
                match self {
                    #(#paths => ::core::stringify!(#idents),)*
                }
            }

            fn from_name(name: &str) -> ::core::option::Option<Self> {
                match name {
                    #(::core::stringify!(#idents) => ::core::option::Option::Some(#paths),)*
                    _ => ::core::option::Option::None,
                }
            }
        }
    }
    .into())
}

/// One action: the value it is, and what it is called in code.
fn verb<'a>(path: Path, fields: &Fields, name: &'a Ident) -> syn::Result<(Path, &'a Ident)> {
    fieldless(
        fields,
        name,
        "an action is a plain verb; move what varies into the game's own state",
    )?;
    Ok((path, name))
}

/// Derives the engine's view of a vocabulary of keys a game keeps between
/// runs: its values, and the name the store keeps each of them under — the
/// vocabulary's own name, `.`, and the value's name.
///
/// The value a key keeps, and its fallback before any run has saved one,
/// are declared by hand in `SaveKey`; a variant with fields fails to
/// compile, because a key is a plain name.
#[proc_macro_derive(Saves)]
pub fn derive_saves(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    match saves(&input) {
        Ok(implementation) => implementation,
        Err(error) => error.to_compile_error().into(),
    }
}

fn saves(input: &DeriveInput) -> syn::Result<TokenStream> {
    let name = &input.ident;
    let keys: Vec<(Path, LitStr)> = match &input.data {
        Data::Enum(data) => data
            .variants
            .iter()
            .map(|variant| {
                let variant_name = &variant.ident;
                kept(
                    parse_quote!(#name::#variant_name),
                    &variant.fields,
                    variant_name,
                    &format!("{name}.{variant_name}"),
                )
            })
            .collect::<syn::Result<_>>()?,
        Data::Struct(data) => vec![kept(
            parse_quote!(#name),
            &data.fields,
            name,
            &name.to_string(),
        )?],
        Data::Union(_) => {
            return Err(syn::Error::new_spanned(
                name,
                "a save vocabulary is an enum or a unit struct; a union needs the impl written \
                 by hand",
            ));
        }
    };

    let (paths, names): (Vec<&Path>, Vec<&LitStr>) =
        keys.iter().map(|(path, name)| (path, name)).unzip();
    let (impl_generics, type_generics, where_clause) = input.generics.split_for_impl();
    Ok(quote! {
        impl #impl_generics ::mirage_engine::Saves for #name #type_generics #where_clause {
            fn name(&self) -> &'static str {
                match *self {
                    #(#paths => #names,)*
                }
            }
        }
    }
    .into())
}

/// One key: the value it is, and the name `under` the store keeps it.
fn kept(path: Path, fields: &Fields, name: &Ident, under: &str) -> syn::Result<(Path, LitStr)> {
    fieldless(
        fields,
        name,
        "a save key is a plain name; move what varies into what it keeps",
    )?;
    Ok((path, LitStr::new(under, name.span())))
}

/// Errors on the fields a vocabulary of plain names has none of.
fn fieldless(fields: &Fields, name: &Ident, complaint: &str) -> syn::Result<()> {
    match fields {
        Fields::Unit => Ok(()),
        _ => Err(syn::Error::new_spanned(
            name,
            format!("`{name}` has fields, but {complaint}"),
        )),
    }
}

/// Derives `ShaderValues` for the values a style's or an effect's WGSL
/// reads: the WGSL struct declaring them, and the layout writing each of
/// them where that struct reads it.
///
/// The WGSL struct takes the Rust type's own name and each field its own,
/// in the order they were written; a field of any other type than these
/// fails to compile, and a type with no fields reads no values and binds
/// none.
///
/// | Rust   | WGSL                                          |
/// |--------|-----------------------------------------------|
/// | `f32`  | `f32`                                         |
/// | `u32`  | `u32`                                         |
/// | `Vec2` | `vec2<f32>`                                   |
/// | `Vec3` | `vec3<f32>`                                   |
/// | `Vec4` | `vec4<f32>`                                   |
/// | `Mat4` | `mat4x4<f32>`                                 |
/// | `Color` | `vec4<f32>`, linear red, green, blue and alpha |
#[proc_macro_derive(ShaderValues)]
pub fn derive_shader_values(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as DeriveInput);
    match shader_values(&input) {
        Ok(implementation) => implementation,
        Err(error) => error.to_compile_error().into(),
    }
}

fn shader_values(input: &DeriveInput) -> syn::Result<TokenStream> {
    let name = &input.ident;
    let Data::Struct(data) = &input.data else {
        return Err(syn::Error::new_spanned(
            name,
            "values a shader reads are a struct of the fields it reads",
        ));
    };
    let read = match &data.fields {
        Fields::Named(fields) => fields.named.iter().collect(),
        Fields::Unit => Vec::new(),
        Fields::Unnamed(_) => {
            return Err(syn::Error::new_spanned(
                name,
                "values a shader reads are named, so that it reads them by name",
            ));
        }
    };

    let layout = Layout::of(name, &read)?;
    let declaration = &layout.declaration;
    let size = layout.size;
    let written = layout
        .placed
        .iter()
        .map(|placed| placed.lane.written(&placed.field, placed.offset));
    let (impl_generics, type_generics, where_clause) = input.generics.split_for_impl();
    Ok(quote! {
        impl #impl_generics ::mirage_engine::Sealed for #name #type_generics #where_clause {}

        impl #impl_generics ::mirage_engine::ShaderValues for #name #type_generics #where_clause {
            const TYPE: &'static str = ::core::stringify!(#name);
            const DECLARATION: &'static str = #declaration;

            fn write(&self, into: &mut ::std::vec::Vec<u8>) {
                let start = into.len();
                #(#written)*
                into.resize(start + #size, 0);
            }
        }
    }
    .into())
}

/// The values a shader reads as the derive lays them out: the WGSL
/// declaring them, where each field is written, and the bytes the whole of
/// them fills.
struct Layout {
    declaration: String,
    placed: Vec<Placed>,
    size: usize,
}

impl Layout {
    /// How a shader reads the fields `read` of the struct `name`.
    fn of(name: &Ident, read: &[&syn::Field]) -> syn::Result<Self> {
        if read.is_empty() {
            return Ok(Self {
                declaration: String::new(),
                placed: Vec::new(),
                size: 0,
            });
        }

        let mut declaration = format!("struct {name} {{\n");
        let mut placed = Vec::new();
        let mut offset = 0usize;
        for read in read {
            let Some(field) = read.ident.clone() else {
                continue;
            };
            let lane = Lane::of(&read.ty)?;
            offset = lane.aligned(offset);
            declaration.push_str(&format!("    {field}: {},\n", lane.wgsl()));
            placed.push(Placed {
                lane,
                field,
                offset,
            });
            offset += lane.size();
        }
        declaration.push('}');
        Ok(Self {
            declaration,
            placed,
            size: offset.next_multiple_of(BLOCK),
        })
    }
}

/// One field as the shader reads it: through which lane, under which name,
/// at which offset.
struct Placed {
    lane: Lane,
    field: Ident,
    offset: usize,
}

/// One kind of value a shader reads, of the kinds it has a lane for.
#[derive(Clone, Copy)]
enum Lane {
    Number,
    Count,
    Vec2,
    Vec3,
    Vec4,
    Mat4,
    Color,
}

impl Lane {
    /// The lane `ty` is read through, or an error naming the kinds there
    /// are.
    fn of(ty: &syn::Type) -> syn::Result<Self> {
        let syn::Type::Path(path) = ty else {
            return Err(Self::unread(ty));
        };
        match path.path.segments.last() {
            Some(segment) => match segment.ident.to_string().as_str() {
                "f32" => Ok(Self::Number),
                "u32" => Ok(Self::Count),
                "Vec2" => Ok(Self::Vec2),
                "Vec3" => Ok(Self::Vec3),
                "Vec4" => Ok(Self::Vec4),
                "Mat4" => Ok(Self::Mat4),
                "Color" => Ok(Self::Color),
                _ => Err(Self::unread(ty)),
            },
            None => Err(Self::unread(ty)),
        }
    }

    fn unread(ty: &syn::Type) -> syn::Error {
        syn::Error::new_spanned(
            ty,
            "a shader reads `f32`, `u32`, `Vec2`, `Vec3`, `Vec4`, `Mat4` and `Color`, and \
             nothing else",
        )
    }

    /// Type the shader declares this lane as.
    fn wgsl(self) -> &'static str {
        match self {
            Self::Number => "f32",
            Self::Count => "u32",
            Self::Vec2 => "vec2<f32>",
            Self::Vec3 => "vec3<f32>",
            Self::Vec4 | Self::Color => "vec4<f32>",
            Self::Mat4 => "mat4x4<f32>",
        }
    }

    /// Bytes of it the shader reads.
    fn size(self) -> usize {
        match self {
            Self::Number | Self::Count => 4,
            Self::Vec2 => 8,
            Self::Vec3 => 12,
            Self::Vec4 | Self::Color => 16,
            Self::Mat4 => 64,
        }
    }

    /// The next offset a value of this lane may start at, from `offset`.
    fn aligned(self, offset: usize) -> usize {
        let align = match self {
            Self::Number | Self::Count => 4,
            Self::Vec2 => 8,
            Self::Vec3 | Self::Vec4 | Self::Color | Self::Mat4 => BLOCK,
        };
        offset.next_multiple_of(align)
    }

    /// Writing one field of this lane where the shader reads it.
    fn written(self, field: &Ident, offset: usize) -> impl quote::ToTokens {
        let numbers = match self {
            Self::Number | Self::Count => vec![quote!(self.#field)],
            Self::Vec2 => vec![quote!(self.#field.x), quote!(self.#field.y)],
            Self::Vec3 => vec![
                quote!(self.#field.x),
                quote!(self.#field.y),
                quote!(self.#field.z),
            ],
            Self::Vec4 => vec![
                quote!(self.#field.x),
                quote!(self.#field.y),
                quote!(self.#field.z),
                quote!(self.#field.w),
            ],
            Self::Color => vec![
                quote!(self.#field.red),
                quote!(self.#field.green),
                quote!(self.#field.blue),
                quote!(self.#field.alpha),
            ],
            Self::Mat4 => {
                return quote! {
                    into.resize(start + #offset, 0);
                    for number in self.#field.to_cols_array() {
                        into.extend_from_slice(&number.to_le_bytes());
                    }
                };
            }
        };

        quote! {
            into.resize(start + #offset, 0);
            #(into.extend_from_slice(&#numbers.to_le_bytes());)*
        }
    }
}