orio_derive/

lib.rs

use proc_macro::TokenStream;
use proc_macro2::{Literal, Span, TokenStream as TokenStream2};
use quote::{quote, ToTokens};
use syn::*;

#[proc_macro_derive(Io, attributes(io))]
pub fn derive_io(input: TokenStream) -> TokenStream {
	let DeriveInput { ident, data, .. } = parse_macro_input!(input);
	derive_io_inner(ident, data).into()
}

fn derive_io_inner(ident: Ident, data: Data) -> TokenStream2 {
	let (read, write, size) = match data {
		Data::Struct(data) => (
			read_fields(quote! { Self }, &data.fields),
			write_struct_fields(&data.fields),
			struct_field_sizes(&data.fields)
		),
		Data::Enum(data) => {
			// TODO: support #[repr(u8)] annotations and stuff
			let variant_type = match data.variants.len() {
				0 => return Error::new(ident.span(), "Empty enums are not supported, use () instead.").to_compile_error(),
				1..=255 => quote! { u8 },
				256..=65535 => quote! { u16 },
				_ => quote! { u32 }
			};
			// TODO: support enum discriminants
			let read_variants = data.variants.iter()
				.enumerate()
				.map(|(i, variant)| {
					let i = Literal::usize_unsuffixed(i);
					let ident = &variant.ident;
					let read_fields = read_fields(quote! { Self::#ident }, &variant.fields);
					quote! {
						#i => #read_fields,
					}
				});
			// writing variant is done first since every enum has it, save the write calls and just make it a single one
			let match_variants = data.variants.iter()
				.enumerate()
				.map(|(i, variant)| {
					// need to correctly discard fields when getting variant
					let i = Literal::usize_unsuffixed(i);
					let ident = &variant.ident;
					let matcher = match &variant.fields {
						Fields::Named(_) => quote! { {..} },
						Fields::Unnamed(_) => quote! { (..) },
						Fields::Unit => quote! {}
					};
					quote! {
						Self::#ident #matcher => #i
					}
				});
			let write_variants = data.variants.iter()
				.filter_map(|variant| {
					let ident = &variant.ident;
					let matcher = enum_matcher(&variant.fields)?;
					let writes = enum_field_idents(&variant.fields, "write", |func, ident| quote! {
						#func(#ident, w)?;
					});
					Some(quote! {
						Self::#ident #matcher => { #(#writes)* }
					})
				});
			let size_variants = data.variants.iter()
				.filter_map(|variant| {
					let ident = &variant.ident;
					let matcher = enum_matcher(&variant.fields)?;
					let sizes = enum_field_idents(&variant.fields, "size", |func, ident| quote! {
						#func(#ident)
					});
					Some(quote! {
						Self::#ident #matcher => #(#sizes)+*
					})
				});
			(quote! {
				// block needed since its inside Ok()
				{
					let variant = #variant_type::read(r)?;
					match variant {
						#(#read_variants)*
						_ => return Err(derive_enum_variant_error(variant.into(), stringify!(#ident)))
					}
				}
			}, quote! {
				let variant: #variant_type = match self {
					#(#match_variants),*
				};
				variant.write(w)?;
				match self {
					#(#write_variants,)*
					_ => {} // for fields that are ignored or have no data
				}
			}, quote! {
				1 + match self {
					#(#size_variants,)*
					_ => 0
				}
			})
		},
		Data::Union(_) => return Error::new(ident.span(), "Unions are not supported").to_compile_error()
	};

	quote! {
		#[automatically_derived]
		impl Io for #ident {
			fn read(r: &mut Reader) -> ::std::io::Result<Self> {
				Ok(#read)
			}

			fn write(&self, w: &mut Writer) -> ::std::io::Result<()> {
				#write
				Ok(())
			}

			fn size(&self) -> usize {
				#size
			}
		}
	}
}

enum FieldOptions {
	None,
	Ignored,
	Len(Type)
}

impl FieldOptions {
	fn is_some(&self) -> bool {
		!matches!(self, Self::None)
	}

	fn is_ignored(&self) -> bool {
		matches!(self, Self::Ignored)
	}
}

fn field_options(field: &Field) -> FieldOptions {
	let mut options = FieldOptions::None;
	for attr in field.attrs.iter() {
		// check for #[io] attribute
		if !attr.path().is_ident("io") {
			continue;
		}

		// find #[io(ignored)]
		if let Err(e) = attr.parse_nested_meta(|meta| {
			if options.is_some() {
				return Err(meta.error("duplicate io attributes used"));
			}

			if meta.path.is_ident("ignore") {
				options = FieldOptions::Ignored;
				return Ok(());
			}

			if meta.path.is_ident("len") {
				let content;
				parenthesized!(content in meta.input);
				let len = content.parse()?;
				options = FieldOptions::Len(len);
				return Ok(());
			}

			Err(meta.error(format!("unknown io attribute '{}' used", meta.path.to_token_stream())))
		}) {
			panic!("Failed to parse attribute for field '{}': {e}", field.ident.as_ref().unwrap());
		}
	}

	options
}


fn is_ignored(field: &Field) -> bool {
	field_options(field).is_ignored()
}

fn read_fields(base: TokenStream2, fields: &Fields) -> TokenStream2 {
	match fields {
		Fields::Named(fields) => {
			let read_fields = fields.named.iter()
				.map(|field| {
					let ident = &field.ident;
					let read = read_field(&field);
					quote! {
						#ident: #read
					}
				});
			quote! {
				#base {
					#(#read_fields),*
				}
			}
		},
		Fields::Unnamed(fields) => {
			let read_fields = fields.unnamed.iter()
				.map(|field| read_field(&field));
			quote! {
				#base(#(#read_fields),*)
			}
		},
		Fields::Unit => base
	}
}

// part of reading a field common to structs and enums
fn read_field(field: &Field) -> TokenStream2 {
	let ty = &field.ty;
	let options = field_options(field);
	match options {
		FieldOptions::Ignored => {
			quote! {
				<#ty as ::std::default::Default>::default()
			}
		},
		FieldOptions::None => {
			quote! {
				<#ty as Io>::read(r)?
			}
		},
		FieldOptions::Len(len) => {
			quote! {
				<#ty as LenIo>::read::<#len>(r)?
			}
		}
	}
}

fn write_struct_fields(fields: &Fields) -> TokenStream2 {
	match fields {
		Fields::Named(fields) => {
			let fields = fields.named.iter()
				.filter_map(|field| {
					let options = field_options(field);
					let ident = &field.ident;
					Some(match options {
						FieldOptions::Ignored => return None,
						FieldOptions::None => quote! {
							Io::write(&self.#ident, w)?;
						},
						FieldOptions::Len(len) => quote! {
							LenIo::write::<#len>(&self.#ident, w)?;
						}
					})
				});
			quote! {
				#(#fields)*
			}
		},
		Fields::Unnamed(fields) => {
			// this is why this cant be reused for enums
			// named fields are ok since "&self." can be removed, but this would need #i to
			// be replaced with a name generated from a-z, which is a bit much.
			let fields = fields.unnamed.iter()
				.enumerate()
				.filter_map(|(i, field)| {
					let options = field_options(field);
					let i: Index = i.into();
					Some(match options {
						FieldOptions::Ignored => return None,
						FieldOptions::None => quote! {
							Io::write(&self.#i, w)?;
						},
						FieldOptions::Len(len) => quote! {
							LenIo::write::<#len>(&self.#i, w)?;
						}
					})
				});
			quote! {
				#(#fields)*
			}
		},
		Fields::Unit => quote! {}
	}
}

fn struct_field_sizes(fields: &Fields) -> TokenStream2 {
	match fields {
		Fields::Named(fields) => {
			let fields = fields.named.iter()
				.filter_map(|field| {
					let options = field_options(field);
					let ident = &field.ident;
					Some(match options {
						FieldOptions::Ignored => return None,
						FieldOptions::None => quote! {
							Io::size(&self.#ident)
						},
						FieldOptions::Len(len) => quote! {
							LenIo::size::<#len>(&self.#ident)
						}
					})
				});
			quote! {
				#(#fields)+*
			}
		},
		Fields::Unnamed(fields) => {
			let fields = fields.unnamed.iter()
				.enumerate()
				.filter_map(|(i, field)| {
					let options = field_options(field);
					let i: Index = i.into();
					Some(match options {
						FieldOptions::Ignored => return None,
						FieldOptions::None => quote! {
							Io::size(&self.#i)
						},
						FieldOptions::Len(len) => quote! {
							LenIo::size::<#len>(&self.#i)
						}
					})
				});
			quote! {
				#(#fields)+*
			}
		},
		Fields::Unit => quote! {}
	}
}

fn enum_matcher(fields: &Fields) -> Option<TokenStream2> {
	Some(match fields {
		Fields::Named(fields) => {
			let fields = fields.named.iter()
				.filter(|field| !is_ignored(field))
				.filter_map(|field| field.ident.clone());
			quote! {
				{ #(#fields),* }
			}
		},
		Fields::Unnamed(fields) => {
			let fields = fields.unnamed.iter()
				.enumerate()
				.filter(|(_, field)| !is_ignored(field))
				.map(|(i, _)| name_unnamed_field(i));
			quote! {
				(#(#fields),*)
			}
		},
		Fields::Unit => return None
	})
}

fn enum_field_idents(fields: &Fields, name: &'static str, f: impl Fn(TokenStream2, Ident) -> TokenStream2) -> Vec<TokenStream2> {
	let name = Ident::new(name, Span::call_site());
	match fields {
		Fields::Named(fields) => fields.named.iter()
			.filter_map(|field| {
				let ident = field.ident.clone().unwrap();
				map_trait(&name, field, ident)
			})
			.map(|(a, b)| f(a, b))
			.collect(),
		Fields::Unnamed(fields) => fields.unnamed.iter()
			.enumerate()
			.filter_map(|(i, field)| {
				let ident = name_unnamed_field(i);
				map_trait(&name, field, ident)
			})
			.map(|(a, b)| f(a, b))
			.collect(),
		Fields::Unit => vec![]
	}
}

fn map_trait(name: &Ident, field: &Field, ident: Ident) -> Option<(TokenStream2, Ident)> {
	let options = field_options(field);
	Some((match options {
		FieldOptions::Ignored => return None,
		FieldOptions::None => quote! {
			Io::#name
		},
		FieldOptions::Len(len) => quote! {
			LenIo::#name::<#len>
		}
	}, ident))
}

// not allowed to match Variant(1, 2, 3) so have to assign them letters
fn name_unnamed_field(i: usize) -> Ident {
	let c = (b'a' + (i as u8)) as char;
	let s = c.to_string();
	Ident::new(&s, Span::call_site())
}