lib.rs

//! # omitty
//!
//! A Rust procedural macro for generating type variants with omitted fields.
//!
//! ## Overview
//!
//! `omitty` provides a simple way to create variants of your structs with specific fields omitted,
//! similar to TypeScript's `Omit<>` utility type. This is useful for creating public-facing types
//! from internal types, API response types, or different permission-level views of your data.
//!
//! ## Features
//!
//! - **On-demand generation** - Types are created only when actually used via `Omit!` macro
//! - **Module-level attribute** - Use `#[omitty]` on modules to enable on-demand type generation
//! - **Automatic derive** - Generates Debug, Clone, Serialize, Deserialize, PartialEq, Eq
//! - **Serde compatible** - Works seamlessly with serde for JSON/serialization
//! - **Type-safe** - Compile-time guarantees for omitted fields
//! - **Automatic conversions** - `impl From<Original>` for easy value creation
//! - **Generic support** - Works with generic types and lifetimes
//!
//! ## Usage
//!
//! ### Basic Example
//!
//! ```rust,ignore
//! use omitty::omitty;
//!
//! #[omitty]
//! mod api {
//!     #[derive(omitty::Omittable, Debug, Clone, serde::Serialize, serde::Deserialize)]
//!     pub struct User {
//!         pub id: u64,
//!         pub name: String,
//!         pub secret_token: String,
//!     }
//!
//!     pub struct Member {
//!         pub user: Omit!(User, secret_token),
//!     }
//! }
//!
//! // Use automatic conversion with .into()
//! let user = api::User { id: 1, name: "Alice".into(), secret_token: "secret".into() };
//! let member = api::Member { user: user.into() };
//! ```
//!
//! ### Multiple Omitted Fields
//!
//! ```rust,ignore
//! use omitty::omitty;
//!
//! #[omitty]
//! mod api {
//!     #[derive(omitty::Omittable, Debug, Clone, serde::Serialize, serde::Deserialize)]
//!     pub struct User {
//!         pub id: u64,
//!         pub name: String,
//!         pub email: String,
//!         pub password_hash: String,
//!         pub secret_token: String,
//!     }
//!
//!     pub struct PublicProfile {
//!         pub user: Omit!(User, password_hash, secret_token, email),
//!     }
//!
//!     pub struct AdminView {
//!         pub user: Omit!(User, secret_token),
//!     }
//! }
//!
//! // Two types are generated on-demand:
//! // 1. __omitty__User::UserWithout_password_hashWithout_secret_tokenWithout_email
//! // 2. __omitty__User::UserWithout_secret_token
//! ```
//!
//! ### Automatic Conversions
//!
//! Each generated type includes an automatic `From<Original>` implementation:
//!
//! ```rust,ignore
//! let user = api::User {
//!     id: 42,
//!     name: "Alice".to_string(),
//!     secret_token: "secret".to_string(),
//! };
//!
//! // Use .into() for automatic conversion
//! let public_user: api::__omitty__User::UserWithout_secret_token = user.into();
//!
//! // Serialization works seamlessly
//! let json = serde_json::to_string(&public_user).unwrap();
//! // json: {"id":42,"name":"Alice"}
//! ```
//!
//! ## Generated Type Naming Convention
//!
//! - Generated types are placed in a submodule: `__omitty__<OriginalType>`
//! - Type names follow the pattern: `<OriginalType>Without_<field1>Without_<field2>...`
//! - Example: `Omit!(User, secret_token)` → `__omitty__User::UserWithout_secret_token`
//! - Generic parameters are preserved: `Omit!(Data<T>, field)` → `__omitty__Data::DataWithout_field<T>`
//!
//! ## Error Handling
//!
//! - **Type not found**: If `Omit!(NonExistent, field)` references a type not in scope,
//!   you'll get a clear compile error: "cannot find type NonExistent"
//! - **Module-only attribute**: `#[omitty]` can only be applied to inline modules with content
//! - **Struct-only derive**: `Omittable` can only be derived for structs with named fields

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,
    }
}