use proc_macro::TokenStream;
use quote::quote;
#[proc_macro_derive(TryFromSlice, attributes(grib_template))]
pub fn derive_try_from_slice(input: TokenStream) -> TokenStream {
let input = syn::parse_macro_input!(input as syn::DeriveInput);
match &input.data {
syn::Data::Struct(data) => impl_try_from_slice_for_struct(&input, data),
syn::Data::Enum(data) => impl_try_from_slice_for_enum(&input, data),
_ => unimplemented!("`TryFromSlice` can only be derived for structs/enums"),
}
.into()
}
fn impl_try_from_slice_for_struct(
input: &syn::DeriveInput,
data: &syn::DataStruct,
) -> proc_macro2::TokenStream {
let name = &input.ident;
let (impl_generics, type_generics, where_clause) = input.generics.split_for_impl();
let Some((kind, fields)) = extract_struct_info(data) else {
unimplemented!(
"`TryFromSlice` can only be derived for structs with named fields or with a single unnamed `u8` field"
)
};
if kind == StructKind::TupleStruct {
let field_reads = fields.iter().map(|field| {
let ty = &field.ty;
quote! {
<#ty as grib_template_helpers::TryFromSlice>::try_from_slice(slice, pos)?
}
});
return quote! {
impl #impl_generics grib_template_helpers::TryFromSlice for #name #type_generics #where_clause {
fn try_from_slice(
slice: &[u8],
pos: &mut usize,
) -> grib_template_helpers::TryFromSliceResult<Self> {
Ok(Self(#(#field_reads),*))
}
}
};
}
let mut field_reads = Vec::new();
let mut idents = Vec::new();
for field in fields {
let ident = field.ident.as_ref().unwrap();
let ty = &field.ty;
let len_attr = field
.attrs
.iter()
.find_map(|attr| attr_value(attr, "len").map(|v| parse_len_attr(&v)));
if let Some(len) = len_attr {
if let syn::Type::Path(type_path) = ty
&& let Some((inner_ty, has_option)) = extract_vec_inner(type_path)
{
let tokens = quote! {
let mut #ident = Vec::with_capacity(#len);
for _ in 0..#len {
let item =
<#inner_ty as grib_template_helpers::TryFromSlice>::try_from_slice(
slice,
pos,
)?;
#ident.push(item);
}
};
let tokens = if has_option {
quote! {
let #ident = if *pos == slice.len() {
None
} else {
#tokens
Some(#ident)
};
}
} else {
tokens
};
field_reads.push(tokens);
idents.push(ident);
continue;
}
unimplemented!(
"`#[grib_template(len = N)]` is only available for `Vec<T>` and `Option<Vec<T>>`"
);
}
let disc_attr = field
.attrs
.iter()
.find_map(|attr| attr_value(attr, "variant").map(|v| parse_variant_attr(&v)));
if let Some(disc_ident) = disc_attr {
field_reads.push(quote! {
let #ident = <#ty as grib_template_helpers::TryEnumFromSlice>::try_enum_from_slice(
#disc_ident,
slice,
pos,
)?;
});
idents.push(ident);
continue;
}
field_reads.push(quote! {
let #ident = <#ty as grib_template_helpers::TryFromSlice>::try_from_slice(slice, pos)?;
});
idents.push(ident);
}
quote! {
impl #impl_generics grib_template_helpers::TryFromSlice for #name #type_generics #where_clause {
fn try_from_slice(
slice: &[u8],
pos: &mut usize,
) -> grib_template_helpers::TryFromSliceResult<Self> {
#(#field_reads)*
Ok(Self { #(#idents),* })
}
}
}
}
fn impl_try_from_slice_for_enum(
input: &syn::DeriveInput,
data: &syn::DataEnum,
) -> proc_macro2::TokenStream {
let name = &input.ident;
let (impl_generics, type_generics, where_clause) = input.generics.split_for_impl();
let mut arms = Vec::new();
for variant in &data.variants {
let variant_ident = &variant.ident;
let disc_expr = variant
.discriminant
.as_ref()
.expect("`TryFromSlice` requires the enum to have explicit discriminant")
.1
.clone();
if let syn::Fields::Unnamed(fields) = &variant.fields
&& fields.unnamed.len() == 1
{
let inner_ty = &fields.unnamed.first().unwrap().ty;
arms.push(quote! {
#disc_expr => {
let inner = <#inner_ty as grib_template_helpers::TryFromSlice>::try_from_slice(
slice,
pos
)?;
Ok(#name::#variant_ident(inner))
}
});
} else {
unimplemented!("`TryFromSlice` only supports single-field tuple variants");
}
}
quote! {
impl #impl_generics grib_template_helpers::TryEnumFromSlice for #name #type_generics #where_clause {
fn try_enum_from_slice(
discriminant: impl Into<u64>,
slice: &[u8],
pos: &mut usize,
) -> grib_template_helpers::TryFromSliceResult<Self> {
match discriminant.into() {
#(#arms),*,
_ => panic!("unknown variant for {}", stringify!(#name)),
}
}
}
}
}
enum LenKind {
Literal(usize),
Ident(syn::Ident),
}
impl quote::ToTokens for LenKind {
fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
match self {
LenKind::Literal(n) => {
tokens.extend(quote! { #n });
}
LenKind::Ident(ident) => {
tokens.extend(quote! { #ident as usize });
}
}
}
}
fn attr_value(attr: &syn::Attribute, ident: &str) -> Option<syn::Expr> {
if !attr.path().is_ident("grib_template") {
return None;
}
let meta = attr.parse_args::<syn::Meta>().ok()?;
if let syn::Meta::NameValue(nv) = meta {
if !nv.path.is_ident(ident) {
return None;
}
Some(nv.value)
} else {
None
}
}
fn parse_len_attr(attr_value: &syn::Expr) -> Option<LenKind> {
match attr_value {
syn::Expr::Lit(syn::ExprLit {
lit: syn::Lit::Int(lit_int),
..
}) => Some(LenKind::Literal(lit_int.base10_parse::<usize>().unwrap())),
syn::Expr::Lit(syn::ExprLit {
lit: syn::Lit::Str(lit_str),
..
}) => Some(LenKind::Ident(syn::Ident::new(
&lit_str.value(),
lit_str.span(),
))),
_ => None,
}
}
fn parse_variant_attr(attr_value: &syn::Expr) -> Option<syn::Ident> {
match attr_value {
syn::Expr::Lit(syn::ExprLit {
lit: syn::Lit::Str(lit_str),
..
}) => Some(syn::Ident::new(&lit_str.value(), lit_str.span())),
_ => None,
}
}
fn extract_vec_inner(type_path: &syn::TypePath) -> Option<(syn::Type, bool)> {
if type_path.path.segments.len() == 1 {
let (type_path, has_option) = if type_path.path.segments[0].ident == "Option"
&& let syn::PathArguments::AngleBracketed(ref args) =
type_path.path.segments[0].arguments
&& let Some(syn::GenericArgument::Type(inner_ty)) = args.args.first()
&& let syn::Type::Path(type_path) = inner_ty
{
(type_path, true)
} else {
(type_path, false)
};
if type_path.path.segments[0].ident == "Vec"
&& let syn::PathArguments::AngleBracketed(ref args) =
type_path.path.segments[0].arguments
&& let Some(syn::GenericArgument::Type(inner_ty)) = args.args.first()
{
return Some((inner_ty.clone(), has_option));
}
}
None
}
#[proc_macro_derive(Dump)]
pub fn derive_dump(input: TokenStream) -> TokenStream {
let input = syn::parse_macro_input!(input as syn::DeriveInput);
match &input.data {
syn::Data::Struct(data) => impl_dump_for_struct(&input, data),
syn::Data::Enum(data) => impl_dump_for_enum(&input, data),
_ => unimplemented!("`Dump` can only be derived for structs/enums"),
}
.into()
}
fn impl_dump_for_struct(
input: &syn::DeriveInput,
data: &syn::DataStruct,
) -> proc_macro2::TokenStream {
let name = &input.ident;
let (impl_generics, type_generics, where_clause) = input.generics.split_for_impl();
let Some((kind, fields)) = extract_struct_info(data) else {
unimplemented!(
"`Dump` can only be derived for structs with named fields or with a single unnamed `u8` field"
)
};
if kind == StructKind::TupleStruct {
let doc = get_doc(&fields[0].attrs)
.map(|s| format!(" // {}", s.trim()))
.unwrap_or_default();
return quote! {
impl #impl_generics grib_template_helpers::Dump for #name #type_generics #where_clause {
fn dump<W: std::io::Write>(
&self,
parent: Option<&std::borrow::Cow<str>>,
pos: &mut usize,
output: &mut W,
) -> Result<(), std::io::Error> {
let size = 1;
grib_template_helpers::write_position_column(output, pos, size)?;
if let Some(parent) = parent {
write!(output, "{}", parent)?;
}
writeln!(output, " = {:#010b}{}",
self.0,
#doc,
)
}
}
};
}
let mut dumps = Vec::new();
for field in fields {
let ident = field.ident.as_ref().unwrap();
let ty = &field.ty;
let doc = get_doc(&field.attrs)
.map(|s| format!(" // {}", s.trim()))
.unwrap_or_default();
dumps.push(quote! {
<#ty as grib_template_helpers::DumpField>::dump_field(
&self.#ident,
stringify!(#ident),
parent,
#doc,
pos,
output,
)?;
});
}
quote! {
impl #impl_generics grib_template_helpers::Dump for #name #type_generics #where_clause {
fn dump<W: std::io::Write>(
&self,
parent: Option<&std::borrow::Cow<str>>,
pos: &mut usize,
output: &mut W,
) -> Result<(), std::io::Error> {
#(#dumps)*;
Ok(())
}
}
}
}
fn impl_dump_for_enum(input: &syn::DeriveInput, data: &syn::DataEnum) -> proc_macro2::TokenStream {
let name = &input.ident;
let (impl_generics, type_generics, where_clause) = input.generics.split_for_impl();
let mut arms = Vec::new();
for variant in &data.variants {
let variant_ident = &variant.ident;
if let syn::Fields::Unnamed(fields) = &variant.fields
&& fields.unnamed.len() == 1
{
let inner_ty = &fields.unnamed.first().unwrap().ty;
arms.push(quote! {
#name::#variant_ident(inner) => <#inner_ty as grib_template_helpers::Dump>::dump(
inner,
parent,
pos,
output
)
});
} else {
unimplemented!("`Dump` only supports single-field tuple variants");
}
}
quote! {
impl #impl_generics grib_template_helpers::Dump for #name #type_generics #where_clause {
fn dump<W: std::io::Write>(
&self,
parent: Option<&std::borrow::Cow<str>>,
pos: &mut usize,
output: &mut W,
) -> Result<(), std::io::Error> {
match self {
#(#arms),*,
}
}
}
}
}
fn get_doc(attrs: &[syn::Attribute]) -> Option<String> {
let mut doc = String::new();
for attr in attrs.iter() {
match attr.meta {
syn::Meta::NameValue(ref value) if value.path.is_ident("doc") => {
if let syn::Expr::Lit(lit) = &value.value
&& let syn::Lit::Str(s) = &lit.lit
{
doc.push_str(&s.value());
}
}
_ => {}
}
}
if doc.is_empty() { None } else { Some(doc) }
}
fn extract_struct_info(
data: &syn::DataStruct,
) -> Option<(
StructKind,
&syn::punctuated::Punctuated<syn::Field, syn::token::Comma>,
)> {
match &data.fields {
syn::Fields::Named(fields) => Some((StructKind::NamedStruct, &fields.named)),
syn::Fields::Unnamed(fields) => {
let fields = &fields.unnamed;
if fields.len() == 1 && is_type_u8(&fields.first().unwrap().ty) {
Some((StructKind::TupleStruct, fields))
} else {
None
}
}
_ => None,
}
}
#[derive(PartialEq)]
enum StructKind {
TupleStruct,
NamedStruct,
}
fn is_type_u8(ty: &syn::Type) -> bool {
if let syn::Type::Path(syn::TypePath { path, .. }) = ty
&& let Some(segment) = path.segments.last()
{
return segment.ident == "u8";
}
false
}