luby/crates/luaffi_impl/src/cdef.rs

use crate::utils::{ffi_crate, syn_assert, syn_error};
use darling::FromMeta;
use proc_macro2::TokenStream;
use quote::{format_ident, quote, quote_spanned};
use syn::{ext::IdentExt, spanned::Spanned, *};

#[derive(Debug, FromMeta)]
pub struct Args {
    module: Option<String>,
}

pub fn transform(args: Args, mut item: Item) -> Result<TokenStream> {
    let (name, impl_type, impl_module, impl_cdef) = match item {
        Item::Struct(ref mut str) => (
            str.ident.clone(),
            generate_type(&str.ident)?,
            args.module
                .map(|name| generate_module(&name, &str.ident))
                .transpose()?,
            generate_cdef_structure(str)?,
        ),
        Item::Enum(ref mut enu) => (
            enu.ident.clone(),
            generate_type(&enu.ident)?,
            args.module
                .map(|name| generate_module(&name, &enu.ident))
                .transpose()?,
            generate_cdef_enum(enu)?,
        ),
        _ => syn_error!(item, "expected struct or enum"),
    };

    let mod_name = format_ident!("__{name}_cdef");

    Ok(quote!(
        #[repr(C)]
        #[allow(non_camel_case_types)]
        #item

        #[doc(hidden)]
        #[allow(unused, non_snake_case)]
        /// Automatically generated by luaffi.
        mod #mod_name {
            use super::*;
            #impl_type
            #impl_module
            #impl_cdef
        }
    ))
}

fn generate_type(ty: &Ident) -> Result<TokenStream> {
    let ffi = ffi_crate();
    let name = ty.unraw().to_string();

    Ok(quote!(
        unsafe impl #ffi::Type for #ty {
            fn name() -> impl ::std::fmt::Display { #name }

            fn ty() -> #ffi::TypeType {
                #ffi::TypeType::Aggregate
            }

            fn cdecl(name: impl ::std::fmt::Display) -> impl ::std::fmt::Display {
                ::std::format!("struct {} {name}", #name)
            }

            fn build(b: &mut #ffi::TypeBuilder) {
                b.cdef::<Self>().metatype::<Self>();
            }
        }

        impl #ffi::Annotate for #ty {
            fn annotation() -> impl ::std::fmt::Display { #name }
        }

        // SAFETY: we can always implement `IntoFfi` because it transfers ownership from Rust to Lua
        unsafe impl #ffi::IntoFfi for #ty {
            type Into = Self;
            fn convert(self) -> Self::Into { self }
        }
    ))
}

fn generate_module(name: &str, ty: &Ident) -> Result<TokenStream> {
    let ffi = ffi_crate();

    Ok(quote!(
        impl #ffi::Module for #ty {
            fn name() -> impl ::std::fmt::Display { #name }
        }
    ))
}

fn generate_cdef_structure(str: &mut ItemStruct) -> Result<TokenStream> {
    syn_assert!(
        str.generics.params.is_empty(),
        str.generics,
        "cannot be generic (not yet implemented)"
    );

    let ffi = ffi_crate();
    let ty = &str.ident;
    let build = generate_cdef_build(&parse_cfields(&mut str.fields)?)?;

    Ok(quote!(
        unsafe impl #ffi::Cdef for #ty {
            fn build(b: &mut #ffi::CdefBuilder) { #build }
        }
    ))
}

fn generate_cdef_enum(enu: &mut ItemEnum) -> Result<TokenStream> {
    syn_assert!(
        enu.generics.params.is_empty(),
        enu.generics,
        "cannot be generic (not yet implemented)"
    );

    let ffi = ffi_crate();
    let ty = &enu.ident;
    let build = enu
        .variants
        .iter_mut()
        .map(|variant| {
            let build = generate_cdef_build(&parse_cfields(&mut variant.fields)?)?;
            Ok(quote!(b.inner_struct(|b| { #build })))
        })
        .collect::<Result<Vec<_>>>()?;

    Ok(quote!(
        unsafe impl #ffi::Cdef for #ty {
            fn build(b: &mut #ffi::CdefBuilder) {
                b.field::<::std::ffi::c_int>("__tag").inner_union(|b| { #(#build;)* });
            }
        }
    ))
}

struct CField {
    name: String,
    ty: Type,
    attrs: CFieldAttrs,
}

#[derive(Default)]
struct CFieldAttrs {
    opaque: bool,
}

fn parse_cfields(fields: &mut Fields) -> Result<Vec<CField>> {
    match fields {
        Fields::Named(fields) => fields.named.iter_mut(),
        Fields::Unnamed(fields) => fields.unnamed.iter_mut(),
        Fields::Unit => return Ok(vec![]),
    }
    .enumerate()
    .map(|(i, field)| {
        Ok(CField {
            name: match field.ident {
                Some(ref name) => name.unraw().to_string(),
                None => format!("__{i}"),
            },
            ty: field.ty.clone(),
            attrs: parse_cfield_attrs(&mut field.attrs)?,
        })
    })
    .collect()
}

fn parse_cfield_attrs(attrs: &mut Vec<Attribute>) -> Result<CFieldAttrs> {
    let mut parsed = CFieldAttrs::default();
    let mut i = 0;
    while let Some(attr) = attrs.get(i) {
        let path = attr.path();
        if path.is_ident("opaque") {
            parsed.opaque = true;
            attrs.remove(i);
            continue;
        }
        i += 1;
    }

    Ok(parsed)
}

fn generate_cdef_build(fields: &[CField]) -> Result<TokenStream> {
    let mut body = vec![quote!(
        let mut offset = 0;
        let mut align = 1;
    )];

    fn offset(i: usize) -> Ident {
        format_ident!("offset{i}")
    }

    let max = quote!(::std::cmp::Ord::max);
    let size_of = quote!(::std::mem::size_of);
    let align_of = quote!(::std::mem::align_of);

    for (i, field) in fields.iter().enumerate() {
        let ty = &field.ty;
        let offset = offset(i);
        body.push(quote!(
            // round up current offset to the alignment of field type for field offset
            offset = (offset + #align_of::<#ty>() - 1) & !(#align_of::<#ty>() - 1);
            align = #max(align, #align_of::<#ty>());
            let #offset = offset;
            offset += #size_of::<#ty>();
        ));
    }

    body.push(quote!(
        // round up final offset to the total alignment of struct for struct size
        let size = (offset + align - 1) & !(align - 1);
    ));

    let len = fields.len();
    for (i, field) in fields.iter().enumerate() {
        let name = &field.name;
        let ty = &field.ty;
        body.push(if field.attrs.opaque {
            if i == len - 1 {
                let a = offset(i);
                quote_spanned!(ty.span() => b.field_opaque(size - #a);) // last field
            } else {
                let a = offset(i);
                let b = offset(i + 1);
                quote_spanned!(ty.span() => b.field_opaque(#b - #a);)
            }
        } else {
            quote_spanned!(ty.span() => b.field::<#ty>(#name);)
        });
    }

    Ok(quote!(#(#body)*))
}