use proc_macro2::{Span, TokenStream};
use quote::ToTokens;
use syn::{Data, DataEnum, DataStruct, DataUnion, Field, Ident, Index, Type, Variant, Visibility};
pub(crate) trait DataExt {
fn fields(&self) -> Vec<(&Visibility, TokenStream, &Type)>;
fn variants(&self) -> Vec<(Option<&Variant>, Vec<(&Visibility, TokenStream, &Type)>)>;
fn tag(&self) -> Option<Ident>;
}
impl DataExt for Data {
fn fields(&self) -> Vec<(&Visibility, TokenStream, &Type)> {
match self {
Data::Struct(strc) => strc.fields(),
Data::Enum(enm) => enm.fields(),
Data::Union(un) => un.fields(),
}
}
fn variants(&self) -> Vec<(Option<&Variant>, Vec<(&Visibility, 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<(&Visibility, TokenStream, &Type)> {
map_fields(&self.fields)
}
fn variants(&self) -> Vec<(Option<&Variant>, Vec<(&Visibility, TokenStream, &Type)>)> {
vec![(None, self.fields())]
}
fn tag(&self) -> Option<Ident> {
None
}
}
impl DataExt for DataEnum {
fn fields(&self) -> Vec<(&Visibility, TokenStream, &Type)> {
map_fields(self.variants.iter().flat_map(|var| &var.fields))
}
fn variants(&self) -> Vec<(Option<&Variant>, Vec<(&Visibility, TokenStream, &Type)>)> {
self.variants.iter().map(|var| (Some(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<(&Visibility, TokenStream, &Type)> {
map_fields(&self.fields.named)
}
fn variants(&self) -> Vec<(Option<&Variant>, Vec<(&Visibility, TokenStream, &Type)>)> {
vec![(None, self.fields())]
}
fn tag(&self) -> Option<Ident> {
None
}
}
fn map_fields<'a>(
fields: impl 'a + IntoIterator<Item = &'a Field>,
) -> Vec<(&'a Visibility, TokenStream, &'a Type)> {
fields
.into_iter()
.enumerate()
.map(|(idx, f)| {
(
&f.vis,
f.ident
.as_ref()
.map(ToTokens::to_token_stream)
.unwrap_or_else(|| Index::from(idx).to_token_stream()),
&f.ty,
)
})
.collect()
}
pub(crate) fn to_ident_str(t: &impl ToString) -> String {
let s = t.to_string();
if let Some(stripped) = s.strip_prefix("r#") {
stripped.to_string()
} else {
s
}
}