btrfs_explorer/btrfs_parse_derive/src/lib.rs

use quote::{quote, format_ident};
use proc_macro2::Span;
use proc_macro::TokenStream;
use syn::{DeriveInput, Data::Enum, parse_macro_input};

#[proc_macro_derive(AllVariants)]
pub fn derive_all_variants(input: TokenStream) -> TokenStream {
	let syn_item: DeriveInput = parse_macro_input!(input);

	let variants = match syn_item.data {
		Enum(enum_item) => {
			enum_item.variants.into_iter().map(|v|v.ident)
		},
		_ => panic!("AllVariants only works on enums!"),
	};
	let enum_name = syn_item.ident;

	let expanded = quote! {
		impl #enum_name {
			fn all_variants() -> &'static[#enum_name] {
				&[ #(#enum_name::#variants),* ]
			}
		}
	};

	expanded.into()
}

#[proc_macro_derive(ParseBin, attributes(skip_bytes, len))]
pub fn derive_parse_bin(input: TokenStream) -> TokenStream {
	let syn_item: DeriveInput = parse_macro_input!(input);
	let name = syn_item.ident;

	match syn_item.data {
		syn::Data::Struct(struct_item) => {
			match struct_item.fields {
				syn::Fields::Named(fields_named) => {
					derive_parse_bin_struct(&name, &fields_named.named)
				},
				syn::Fields::Unnamed(fields_unnamed) => {
					if fields_unnamed.unnamed.len() != 1 {
						panic!("ParseBin does not support tuple structs!");
					}

					let inner_type = fields_unnamed.unnamed.into_iter().next().unwrap().ty;
					derive_parse_bin_alias(name, inner_type)
				},
				_ => panic!("ParseBin on unit structs makes no sense!"),
			}
		},
		_ => panic!("ParseBin only works on structs so far!"),
	}
}

fn derive_parse_bin_alias(name: syn::Ident, ty: syn::Type) -> TokenStream {
	quote! {
		impl ParseBin for #name {
			fn parse_len(bytes: &[u8]) -> Result<(Self, usize), ParseError> {
				let (result, size) = <#ty>::parse_len(bytes)?;
				Ok((#name(result), size))
			}
		}
	}.into()
}

fn derive_parse_bin_struct<'a, T>(name: &syn::Ident, fields: T) -> TokenStream
	where T: IntoIterator<Item = &'a syn::Field>
{
	let mut parsing_statements = Vec::new();
	let mut combining_expressions = Vec::new();

	for field in fields {
		let field_name = field.ident.as_ref().unwrap();
		let field_type = &field.ty;
		let mut skip: Option<usize> = None;
		let mut veclen: Option<String> = None;

		// look for attributes
		for at in &field.attrs {
			if let syn::Meta::NameValue(nv) = &at.meta {
				if nv.path.segments.len() == 1 {
					let attr_name = nv.path.segments[0].ident.to_string();
					if attr_name == "skip_bytes" {
						if let syn::Expr::Lit(expr) = &nv.value {
							if let syn::Lit::Int(nbytes) = &expr.lit {
//								println!("reserved = {}", nbytes);
								skip = nbytes.base10_parse::<usize>().ok()
							}
						}
					} else if attr_name == "len" {
						if let syn::Expr::Lit(expr) = &nv.value {
							if let syn::Lit::Str(litstr) = &expr.lit {
//								println!("len = {}", litstr.value());
								veclen = Some(litstr.value());
							}
						}
					}
				}
			}
		}

		if let Some(offset) = skip {
			parsing_statements.push(quote!{
				__parse_bin_derive_size += #offset;
			});
		}

		if let Some(varname) = veclen {
			let field_name_item = format_ident!("{}_item", field_name);

			enum FieldType<'a> {
				Vec(&'a syn::Type),
				CString,
			}

			let syn::Type::Path(tp) = &field_type else { panic!() };
			let single = tp.path.segments.iter().next().unwrap();

			let field_type = if &single.ident.to_string() ==  "Vec" {
				let syn::PathArguments::AngleBracketed(args) = &single.arguments else { panic!() };
				let firstarg = args.args.iter().next().unwrap();
				let syn::GenericArgument::Type(ty) = firstarg else { panic!() };
				FieldType::Vec(ty)
			} else if &single.ident.to_string() == "CString" {
				FieldType::CString
			} else {
				panic!("The len attribute is only allowed on Vec<_> or CString")
			};

			let varname_ident = syn::Ident::new(&varname, Span::call_site());


			match field_type {
				FieldType::Vec(field_type_item) => {
					parsing_statements.push(quote!{
						let mut #field_name = Vec::new();
						for i in 0 .. #varname_ident.0 as usize {
							let #field_name_item = <#field_type_item>::parse_len(&bytes[__parse_bin_derive_size..])?;
							__parse_bin_derive_size += #field_name_item.1;
							#field_name.push(#field_name_item.0);
						}
					});
					combining_expressions.push(quote!(#field_name: #field_name));
				},
				FieldType::CString => {
					parsing_statements.push(quote!{
						let #field_name = CString::parse_len(&bytes[__parse_bin_derive_size .. __parse_bin_derive_size + #varname_ident.0 as usize])?;
						__parse_bin_derive_size += #varname_ident.0 as usize;
					});
					combining_expressions.push(quote!(#field_name: #field_name.0));
				},
			}
		} else {
			parsing_statements.push(quote!{
				let #field_name = <#field_type>::parse_len(&bytes[__parse_bin_derive_size..])?;
				__parse_bin_derive_size += #field_name.1;
			});

			combining_expressions.push(quote!(#field_name: #field_name.0));
		}
	}

	quote! {
		impl ParseBin for #name {
			fn parse_len(bytes: &[u8]) -> Result<(Self, usize), ParseError> {
				let mut __parse_bin_derive_size: usize = 0;

				#(#parsing_statements)*

				let result = #name {
					#(#combining_expressions),*
				};

				Ok((result, __parse_bin_derive_size))
			}
		}
	}.into()
}