use proc_macro2::{Span, TokenStream};
use quote::ToTokens;
use syn::{Data, DataEnum, DataStruct, DataUnion, Field, Ident, Index, Type};
pub(crate) trait DataExt {
fn fields(&self) -> Vec<(TokenStream, &Type)>;
fn variants(&self) -> Vec<Vec<(TokenStream, &Type)>>;
fn tag(&self) -> Option<Ident>;
}
impl DataExt for Data {
fn fields(&self) -> Vec<(TokenStream, &Type)> {
match self {
Data::Struct(strc) => strc.fields(),
Data::Enum(enm) => enm.fields(),
Data::Union(un) => un.fields(),
}
}
fn variants(&self) -> Vec<Vec<(TokenStream, &Type)>> {
match self {
Data::Struct(strc) => strc.variants(),
Data::Enum(enm) => enm.variants(),
Data::Union(un) => un.variants(),
}
}
fn tag(&self) -> Option<Ident> {
match self {
Data::Struct(strc) => strc.tag(),
Data::Enum(enm) => enm.tag(),
Data::Union(un) => un.tag(),
}
}
}
impl DataExt for DataStruct {
fn fields(&self) -> Vec<(TokenStream, &Type)> {
map_fields(&self.fields)
}
fn variants(&self) -> Vec<Vec<(TokenStream, &Type)>> {
vec![self.fields()]
}
fn tag(&self) -> Option<Ident> {
None
}
}
impl DataExt for DataEnum {
fn fields(&self) -> Vec<(TokenStream, &Type)> {
map_fields(self.variants.iter().flat_map(|var| &var.fields))
}
fn variants(&self) -> Vec<Vec<(TokenStream, &Type)>> {
self.variants.iter().map(|var| map_fields(&var.fields)).collect()
}
fn tag(&self) -> Option<Ident> {
Some(Ident::new("___ZerocopyTag", Span::call_site()))
}
}
impl DataExt for DataUnion {
fn fields(&self) -> Vec<(TokenStream, &Type)> {
map_fields(&self.fields.named)
}
fn variants(&self) -> Vec<Vec<(TokenStream, &Type)>> {
vec![self.fields()]
}
fn tag(&self) -> Option<Ident> {
None
}
}
fn map_fields<'a>(
fields: impl 'a + IntoIterator<Item = &'a Field>,
) -> Vec<(TokenStream, &'a Type)> {
fields
.into_iter()
.enumerate()
.map(|(idx, f)| {
(
f.ident
.as_ref()
.map(ToTokens::to_token_stream)
.unwrap_or_else(|| Index::from(idx).to_token_stream()),
&f.ty,
)
})
.collect()
}