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};
use syn::{ext::IdentExt, *};

#[derive(Debug, FromMeta)]
pub struct Args {}

pub fn transform(_args: Args, mut item: Item) -> Result<TokenStream> {
    let (name, impl_type, impl_cdef) = match item {
        Item::Struct(ref mut str) => (
            str.ident.clone(),
            generate_type(&str.ident)?,
            generate_cdef_structure(str)?,
        ),
        Item::Enum(ref mut enu) => (
            enu.ident.clone(),
            generate_type(&enu.ident)?,
            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)]
        mod #mod_name {
            use super::*;
            #impl_type
            #impl_cdef
        }
    })
}

fn generate_type(ty: &Ident) -> Result<TokenStream> {
    let ffi = ffi_crate();
    let fmt = quote!(::std::format!);
    let name = LitStr::new(&format!("{}", ty.unraw()), ty.span());
    let cdecl_fmt = LitStr::new(&format!("struct {} {{}}", ty.unraw()), ty.span());

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

            fn cdecl(name: impl ::std::fmt::Display) -> impl ::std::fmt::Display {
                #fmt(#cdecl_fmt, name)
            }

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

        unsafe impl #ffi::ToFfi for #ty {
            type To = Self;
            fn convert(self) -> Self::To { self }
        }
    })
}

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_build_cdef(&to_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_build_cdef(&to_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 to_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) => format!("{}", name.unraw()),
                None => format!("__{i}"),
            },
            ty: field.ty.clone(),
            attrs: parse_attrs(&mut field.attrs)?,
        })
    })
    .collect()
}

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

    Ok(parsed)
}

fn generate_build_cdef(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;
        if field.attrs.opaque {
            body.push(if i == len - 1 {
                let a = offset(i);
                quote! { b.field_opaque(size - #a); } // last field
            } else {
                let a = offset(i);
                let b = offset(i + 1);
                quote! { b.field_opaque(#b - #a); }
            });
        } else {
            body.push(quote! { b.field::<#ty>(#name); });
        }
    }

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