extern crate proc_macro;
use proc_macro::TokenStream;
use proc_macro2::{Ident, Span, TokenStream as TokenStream2};
use quote::{ToTokens, format_ident, quote};
use std::collections::{HashMap, HashSet};
use syn::{
Attribute, DeriveInput, Fields, Item, ItemMod, Token, Type,
parse::{Parse, ParseStream},
parse_macro_input,
punctuated::Punctuated,
spanned::Spanned,
};
#[proc_macro_derive(Omittable)]
pub fn derive_omittable(input: TokenStream) -> TokenStream {
let input = parse_macro_input!(input as DeriveInput);
match &input.data {
syn::Data::Struct(s) => match &s.fields {
Fields::Named(_) => {}
_ => {
let err = syn::Error::new(
input.ident.span(),
"omitty: Omittable can only be derived for structs with named fields (not tuple structs or unit structs)",
);
return err.to_compile_error().into();
}
},
_ => {
let err = syn::Error::new(
input.ident.span(),
"omitty: Omittable can only be derived for structs (not enums or unions)",
);
return err.to_compile_error().into();
}
};
TokenStream::new()
}
#[proc_macro_attribute]
pub fn omitty(_attr: TokenStream, item: TokenStream) -> TokenStream {
let mut input = parse_macro_input!(item as ItemMod);
let (brace, items) = match input.content.take() {
Some((brace, items)) => (brace, items),
None => {
return syn::Error::new(
input.span(),
"omitty: attribute can only be applied to inline modules with content",
)
.to_compile_error()
.into();
}
};
let mut omit_usages: HashMap<String, HashSet<Vec<String>>> = HashMap::new();
let mut type_definitions: HashMap<String, Item> = HashMap::new();
for item in &items {
if let Item::Struct(s) = item {
let has_omittable = s.attrs.iter().any(|attr| {
attr.path().segments.iter().any(|seg| seg.ident == "derive")
&& attr.to_token_stream().to_string().contains("Omittable")
});
if has_omittable {
type_definitions.insert(s.ident.to_string(), item.clone());
}
}
collect_omit_calls(item, &mut omit_usages);
}
let mut generated_mods = Vec::new();
for (type_name, field_combinations) in &omit_usages {
let original_item = match type_definitions.get(type_name) {
Some(item) => item,
None => continue,
};
let original_struct = match original_item {
Item::Struct(s) => s,
_ => continue,
};
let original_fields = match &original_struct.fields {
Fields::Named(f) => &f.named,
_ => continue,
};
let mod_name = format_ident!("__omitty__{}", type_name);
let mod_name_ident =
syn::Ident::new(&format!("__omitty__{}", type_name), Span::call_site());
let mut mod_items = Vec::new();
for omitted_fields in field_combinations {
let mut suffix = String::new();
for field in omitted_fields {
suffix.push_str(&format!("Without_{}", field));
}
let variant_name = format_ident!("{}{}", type_name, suffix);
let kept_fields: Vec<_> = original_fields
.iter()
.filter(|f| {
let fname = f.ident.as_ref().unwrap().to_string();
!omitted_fields.contains(&fname)
})
.collect();
let kept_field_names: Vec<_> = kept_fields
.iter()
.map(|f| f.ident.as_ref().unwrap())
.collect();
let vis = &original_struct.vis;
let type_name_ident = format_ident!("{}", type_name);
let omitted_list = omitted_fields.join(", ");
let generics = &original_struct.generics;
let (impl_generics, ty_generics, where_clause) = generics.split_for_impl();
let struct_tokens = quote! {
#[doc = concat!("Generated by `omitty` for `Omit!(", stringify!(#type_name_ident), ", ", #omitted_list, ")`.")]
#[doc = ""]
#[doc = "This type omits the following fields from the original:"]
#[doc = concat!("- ", #omitted_list)]
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize, PartialEq, Eq)]
#vis struct #variant_name #impl_generics #where_clause {
#(#kept_fields),*
}
};
let impl_tokens = quote! {
#[doc = concat!("Automatic conversion from `", stringify!(#type_name_ident), "` to `", stringify!(#variant_name), "`.")]
#[doc = ""]
#[doc = "Generated by `omitty`. Copies all non-omitted fields."]
impl #impl_generics From<super::#type_name_ident #ty_generics> for #variant_name #ty_generics #where_clause {
fn from(value: super::#type_name_ident #ty_generics) -> Self {
Self {
#(#kept_field_names: value.#kept_field_names),*
}
}
}
};
if let Ok(struct_item) = syn::parse2::<Item>(struct_tokens) {
mod_items.push(struct_item);
}
if let Ok(impl_item) = syn::parse2::<Item>(impl_tokens) {
mod_items.push(impl_item);
}
}
let use_super = syn::parse2::<Item>(quote! { use super::*; }).unwrap();
mod_items.insert(0, use_super);
let mod_doc = format!(
"Generated by `omitty` — do not edit.\n\nContains omitted variants of `{}`.",
type_name
);
let generated_mod = Item::Mod(ItemMod {
attrs: vec![syn::parse_quote! { #[doc = #mod_doc] }],
vis: syn::Visibility::Public(syn::token::Pub::default()),
unsafety: None,
mod_token: syn::token::Mod::default(),
ident: mod_name_ident,
content: Some((syn::token::Brace::default(), mod_items)),
semi: None,
});
generated_mods.push(generated_mod);
}
let mut new_items = Vec::new();
for item in items {
let replaced_item = replace_omit_macros(item, &omit_usages);
new_items.push(replaced_item);
}
for gen_mod in generated_mods {
new_items.push(gen_mod);
}
input.content = Some((brace, new_items));
let output = quote! {
#input
};
output.into()
}
fn collect_omit_calls(item: &Item, usages: &mut HashMap<String, HashSet<Vec<String>>>) {
match item {
Item::Struct(s) => {
for field in &s.fields {
if let Type::Macro(tm) = &field.ty {
if let Some((type_name, omitted_fields, _)) = parse_omit_macro(&tm.mac) {
usages
.entry(type_name)
.or_insert_with(HashSet::new)
.insert(omitted_fields);
}
}
}
}
Item::Enum(e) => {
for variant in &e.variants {
for field in &variant.fields {
if let Type::Macro(tm) = &field.ty {
if let Some((type_name, omitted_fields, _)) = parse_omit_macro(&tm.mac) {
usages
.entry(type_name)
.or_insert_with(HashSet::new)
.insert(omitted_fields);
}
}
}
}
}
Item::Mod(m) => {
if let Some((_, items)) = &m.content {
for item in items {
collect_omit_calls(item, usages);
}
}
}
_ => {}
}
}
fn parse_omit_macro(mac: &syn::Macro) -> Option<(String, Vec<String>, TokenStream2)> {
if !mac.path.is_ident("Omit") {
return None;
}
let tokens = mac.tokens.clone();
let tokens_str = tokens.to_string();
let parts: Vec<&str> = tokens_str.split(',').map(|s| s.trim()).collect();
if parts.is_empty() {
return None;
}
let type_part = parts[0];
let full_type_tokens: TokenStream2 = type_part.parse().ok()?;
let type_name = if let Some(pos) = type_part.find('<') {
type_part[..pos].trim().to_string()
} else {
type_part.to_string()
};
let field_names: Vec<String> = parts[1..].iter().map(|s| s.to_string()).collect();
Some((type_name, field_names, full_type_tokens))
}
fn replace_omit_macros(item: Item, usages: &HashMap<String, HashSet<Vec<String>>>) -> Item {
match item {
Item::Struct(mut s) => {
for field in s.fields.iter_mut() {
if let Type::Macro(tm) = &field.ty {
if let Some((type_name, omitted_fields, full_type_tokens)) =
parse_omit_macro(&tm.mac)
{
let mut suffix = String::new();
for field in &omitted_fields {
suffix.push_str(&format!("Without_{}", field));
}
let variant_name = format_ident!("{}{}", type_name, suffix);
let mod_name = format_ident!("__omitty__{}", type_name);
let type_str = full_type_tokens.to_string();
let has_generics = type_str.contains('<');
let new_type: Type = if has_generics {
let generics_part = &type_str[type_str.find('<').unwrap()..];
let generics_tokens: TokenStream2 = generics_part.parse().unwrap();
syn::parse_quote! {
#mod_name::#variant_name #generics_tokens
}
} else {
syn::parse_quote! {
#mod_name::#variant_name
}
};
field.ty = new_type;
}
}
}
Item::Struct(s)
}
Item::Enum(mut e) => {
for variant in e.variants.iter_mut() {
for field in variant.fields.iter_mut() {
if let Type::Macro(tm) = &field.ty {
if let Some((type_name, omitted_fields, full_type_tokens)) =
parse_omit_macro(&tm.mac)
{
let mut suffix = String::new();
for field in &omitted_fields {
suffix.push_str(&format!("Without_{}", field));
}
let variant_name = format_ident!("{}{}", type_name, suffix);
let mod_name = format_ident!("__omitty__{}", type_name);
let type_str = full_type_tokens.to_string();
let has_generics = type_str.contains('<');
let new_type: Type = if has_generics {
let generics_part = &type_str[type_str.find('<').unwrap()..];
let generics_tokens: TokenStream2 = generics_part.parse().unwrap();
syn::parse_quote! {
#mod_name::#variant_name #generics_tokens
}
} else {
syn::parse_quote! {
#mod_name::#variant_name
}
};
field.ty = new_type;
}
}
}
}
Item::Enum(e)
}
Item::Mod(mut m) => {
if let Some((brace, items)) = m.content {
let new_items = items
.into_iter()
.map(|item| replace_omit_macros(item, usages))
.collect();
m.content = Some((brace, new_items));
}
Item::Mod(m)
}
other => other,
}
}