rustmax-macros 0.0.11

//! Procedural macros for the [`rustmax`] crate.
//!
//! This crate is an implementation detail of `rustmax`.
//! Use it through the `rmx` re-exports,
//! not as a direct dependency.
//!
//! [`rustmax`]: https://docs.rs/rustmax

use proc_macro::TokenStream;
use proc_macro2::{Ident, Literal, Span, TokenStream as TokenStream2, TokenTree};
use proc_macro_crate::{FoundCrate, crate_name};
use quote::{quote, quote_spanned};

/// Derive macros that `rustmax` knows how to route, and their origin crate.
///
/// Names that collide with a standard library derive are deliberately
/// absent so that the standard derive keeps winning: `Debug`, `Eq` and
/// `PartialEq` all exist in `derive_more` but are far more often meant as
/// the built-in ones. Those are still reachable by qualifying the name,
/// as in `#[rmx::derive(derive_more::Debug)]`.
///
/// `Error` is routed to `thiserror` rather than `derive_more` because that
/// is what it usually means, even though `thiserror` is the one crate here
/// whose derive cannot work through a re-export.
const KNOWN_DERIVES: &[(&str, &str)] = &[
    ("Args", "clap"),
    ("Parser", "clap"),
    ("Subcommand", "clap"),
    ("ValueEnum", "clap"),
    ("Add", "derive_more"),
    ("AddAssign", "derive_more"),
    ("AsMut", "derive_more"),
    ("AsRef", "derive_more"),
    ("Binary", "derive_more"),
    ("BitAnd", "derive_more"),
    ("BitAndAssign", "derive_more"),
    ("BitOr", "derive_more"),
    ("BitOrAssign", "derive_more"),
    ("BitXor", "derive_more"),
    ("BitXorAssign", "derive_more"),
    ("Constructor", "derive_more"),
    ("Deref", "derive_more"),
    ("DerefMut", "derive_more"),
    ("Display", "derive_more"),
    ("Div", "derive_more"),
    ("DivAssign", "derive_more"),
    ("From", "derive_more"),
    ("FromStr", "derive_more"),
    ("Index", "derive_more"),
    ("IndexMut", "derive_more"),
    ("Into", "derive_more"),
    ("IntoIterator", "derive_more"),
    ("IsVariant", "derive_more"),
    ("LowerExp", "derive_more"),
    ("LowerHex", "derive_more"),
    ("Mul", "derive_more"),
    ("MulAssign", "derive_more"),
    ("Neg", "derive_more"),
    ("Not", "derive_more"),
    ("Octal", "derive_more"),
    ("Pointer", "derive_more"),
    ("Product", "derive_more"),
    ("Rem", "derive_more"),
    ("RemAssign", "derive_more"),
    ("Shl", "derive_more"),
    ("ShlAssign", "derive_more"),
    ("Shr", "derive_more"),
    ("ShrAssign", "derive_more"),
    ("Sub", "derive_more"),
    ("SubAssign", "derive_more"),
    ("Sum", "derive_more"),
    ("TryFrom", "derive_more"),
    ("TryInto", "derive_more"),
    ("TryUnwrap", "derive_more"),
    ("Unwrap", "derive_more"),
    ("UpperExp", "derive_more"),
    ("UpperHex", "derive_more"),
    ("FromPrimitive", "num_enum"),
    ("IntoPrimitive", "num_enum"),
    ("TryFromPrimitive", "num_enum"),
    ("UnsafeFromPrimitive", "num_enum"),
    ("Deserialize", "serde"),
    ("Serialize", "serde"),
    ("Error", "thiserror"),
];

/// How the code generated by a derive macro names its own crate.
///
/// This is what decides whether a derive can be used through a re-export
/// at all, and what has to be done to it if so.
enum Naming {
    /// Generated code does `extern crate serde as _serde`, overridable
    /// with a `#[serde(crate = "...")]` string attribute.
    CrateAttrStr,
    /// Generated code uses an absolute path, overridable with a
    /// `#[num_enum(crate = ...)]` path attribute.
    CrateAttrPath,
    /// Generated code names the crate relative to the enclosing scope, so
    /// the crate name only has to be bound there.
    Scoped,
    /// Generated code uses an absolute path with no override, so only a
    /// direct dependency can satisfy it.
    ExternOnly,
}

fn naming(krate: &str) -> Naming {
    match krate {
        "serde" => Naming::CrateAttrStr,
        "num_enum" => Naming::CrateAttrPath,
        "clap" | "derive_more" => Naming::Scoped,
        "thiserror" => Naming::ExternOnly,
        _ => unreachable!("no naming strategy for `{krate}`"),
    }
}

/// Apply derive macros re-exported by `rustmax`.
///
/// Most derive macros generate code that names their own crate, which
/// fails when the crate is reached through `rustmax` instead of being a
/// direct dependency. This attribute rewrites the derive paths and injects
/// whatever each derive crate needs to find its runtime.
///
/// ```ignore
/// #[rmx::derive(Serialize, Deserialize, Clone, Debug)]
/// #[serde(rename_all = "camelCase")]
/// struct Config {
///     some_field: u32,
/// }
/// ```
///
/// Names it does not recognize are passed through to `#[derive(..)]`
/// unchanged, so built-in and third-party derives can be listed alongside.
#[proc_macro_attribute]
pub fn derive(attr: TokenStream, item: TokenStream) -> TokenStream {
    let item = TokenStream2::from(item);
    let requests = match parse_requests(TokenStream2::from(attr)) {
        Ok(requests) => requests,
        Err(err) => return prepend(err, item).into(),
    };

    let rustmax = match rustmax_path() {
        Ok(path) => path,
        Err(err) => return prepend(err, item).into(),
    };

    let mut paths = Vec::new();
    let mut scoped = Vec::new();
    let mut crate_attrs = Vec::new();
    let mut errors = TokenStream2::new();

    for request in requests {
        let (name, krate, span) = match request {
            Request::Passthrough(tokens) => {
                paths.push(tokens);
                continue;
            }
            Request::Bare(name) => match lookup(&name.to_string()) {
                Some(krate) => {
                    let span = name.span();
                    (name, krate, span)
                }
                None => {
                    paths.push(TokenTree::Ident(name).into());
                    continue;
                }
            },
            Request::Qualified(krate, name) => match lookup_crate(&krate.to_string()) {
                Some(krate) => {
                    let span = name.span();
                    (name, krate, span)
                }
                None => {
                    paths.push(quote!(#krate::#name));
                    continue;
                }
            },
        };

        match naming(krate) {
            // A crate-path attribute works whether or not the crate is also
            // a direct dependency, so always route these through `rustmax`.
            Naming::CrateAttrStr | Naming::CrateAttrPath => {
                if !crate_attrs.iter().any(|(k, _)| *k == krate) {
                    crate_attrs.push((krate, crate_attr(krate, &rustmax)));
                }
            }
            // Binding the crate's name in scope conflicts with a direct
            // dependency of the same name, so only do it when there is none.
            Naming::Scoped => match direct_path(krate) {
                Some(direct) => {
                    paths.push(quote!(#direct::#name));
                    continue;
                }
                None => {
                    if !scoped.contains(&krate) {
                        scoped.push(krate);
                    }
                }
            },
            Naming::ExternOnly => match direct_path(krate) {
                Some(direct) => {
                    paths.push(quote!(#direct::#name));
                    continue;
                }
                None => {
                    errors.extend(extern_only_error(span, krate));
                    continue;
                }
            },
        }

        let root = &rustmax.tokens;
        let module = Ident::new(krate, span);
        paths.push(quote!(#root::#module::#name));
    }

    let root = &rustmax.tokens;
    let scope_imports = scoped.iter().map(|krate| {
        let module = Ident::new(krate, Span::call_site());
        quote! {
            #[allow(unused_imports)]
            use #root::__macro_scope::#module::*;
        }
    });
    let crate_attrs = crate_attrs.into_iter().map(|(_, attr)| attr);
    let derive_attr = (!paths.is_empty()).then(|| quote!(#[derive(#(#paths),*)]));

    quote! {
        #errors
        #(#scope_imports)*
        #derive_attr
        #(#crate_attrs)*
        #item
    }
    .into()
}

/// The crate-path attribute that tells a derive crate where its runtime is.
fn crate_attr(krate: &str, rustmax: &RustmaxPath) -> TokenStream2 {
    let module = Ident::new(krate, Span::call_site());
    match naming(krate) {
        Naming::CrateAttrStr => {
            let path = Literal::string(&format!("{}::{}", rustmax.text, krate));
            quote!(#[#module(crate = #path)])
        }
        Naming::CrateAttrPath => {
            let root = &rustmax.tokens;
            quote!(#[#module(crate = #root::#module)])
        }
        Naming::Scoped | Naming::ExternOnly => {
            unreachable!("`{krate}` has no crate-path attribute")
        }
    }
}

fn extern_only_error(span: Span, krate: &str) -> TokenStream2 {
    let message = format!(
        "`{krate}` derives only work when `{krate}` is a direct dependency: \
         the code they generate names `::{krate}` absolutely, which resolves \
         only against direct dependencies. Add `{krate}` to your Cargo.toml \
         alongside `rustmax`, or use `derive_more` instead, as in \
         `#[rmx::derive(Display, derive_more::Error)]`."
    );
    error(span, &message)
}

/// The path by which the crate being compiled can name `rustmax`.
struct RustmaxPath {
    tokens: TokenStream2,
    text: String,
}

fn rustmax_path() -> Result<RustmaxPath, TokenStream2> {
    match crate_name("rustmax") {
        Ok(FoundCrate::Itself) if in_doctest() => Ok(RustmaxPath {
            tokens: quote!(::rustmax),
            text: "::rustmax".to_owned(),
        }),
        Ok(FoundCrate::Itself) => Ok(RustmaxPath {
            tokens: quote!(crate),
            text: "crate".to_owned(),
        }),
        Ok(FoundCrate::Name(name)) => {
            let ident = Ident::new(&name, Span::call_site());
            Ok(RustmaxPath {
                tokens: quote!(::#ident),
                text: format!("::{name}"),
            })
        }
        Err(_) => Err(error(
            Span::call_site(),
            "`#[rmx::derive]` requires `rustmax` to be a direct dependency \
             of this crate",
        )),
    }
}

/// Whether the current compilation is a doctest.
///
/// `proc-macro-crate` reports `FoundCrate::Itself` whenever the manifest
/// being compiled is `rustmax`'s own, but a doctest is a separate crate
/// that reaches `rustmax` through the extern prelude, so `crate` does not
/// name it there. `rustdoc` sets this variable for the doctests it builds.
fn in_doctest() -> bool {
    std::env::var_os("UNSTABLE_RUSTDOC_TEST_PATH").is_some()
}

/// The path by which the crate being compiled can name `krate`, if it is a
/// direct dependency.
///
/// This reads the manifest, which does not distinguish `[dependencies]`
/// from `[dev-dependencies]`, while the extern prelude of a given target
/// contains only one of the two. A crate that is a normal dependency but
/// not a dev-dependency is reported here for an integration test that
/// cannot actually name it.
fn direct_path(krate: &str) -> Option<TokenStream2> {
    match crate_name(krate) {
        Ok(FoundCrate::Itself) => Some(quote!(crate)),
        Ok(FoundCrate::Name(name)) => {
            let ident = Ident::new(&name, Span::call_site());
            Some(quote!(::#ident))
        }
        Err(_) => None,
    }
}

fn lookup(name: &str) -> Option<&'static str> {
    KNOWN_DERIVES
        .iter()
        .find(|(known, _)| *known == name)
        .map(|(_, krate)| *krate)
}

fn lookup_crate(name: &str) -> Option<&'static str> {
    KNOWN_DERIVES
        .iter()
        .find(|(_, krate)| *krate == name)
        .map(|(_, krate)| *krate)
}

/// One entry in the attribute's comma-separated list.
enum Request {
    /// A bare derive name, as in `Serialize`.
    Bare(Ident),
    /// A crate-qualified derive name, as in `derive_more::Debug`.
    Qualified(Ident, Ident),
    /// Anything else, passed to `#[derive(..)]` unchanged.
    Passthrough(TokenStream2),
}

fn parse_requests(attr: TokenStream2) -> Result<Vec<Request>, TokenStream2> {
    let mut requests = Vec::new();

    for segment in split_commas(attr) {
        let idents: Vec<&TokenTree> = segment.iter().collect();
        match idents.as_slice() {
            [] => continue,
            [TokenTree::Ident(name)] => requests.push(Request::Bare(name.clone())),
            [
                TokenTree::Ident(krate),
                TokenTree::Punct(first),
                TokenTree::Punct(second),
                TokenTree::Ident(name),
            ] if first.as_char() == ':' && second.as_char() == ':' => {
                requests.push(Request::Qualified(krate.clone(), name.clone()));
            }
            _ => {
                let span = segment
                    .first()
                    .map(|tree| tree.span())
                    .unwrap_or_else(Span::call_site);
                if segment.iter().any(|tree| matches!(tree, TokenTree::Ident(_))) {
                    requests.push(Request::Passthrough(segment.into_iter().collect()));
                } else {
                    return Err(error(
                        span,
                        "expected a derive macro name, as in \
                         `#[rmx::derive(Serialize, Clone)]`",
                    ));
                }
            }
        }
    }

    Ok(requests)
}

fn split_commas(input: TokenStream2) -> Vec<Vec<TokenTree>> {
    let mut segments = vec![Vec::new()];

    for tree in input {
        match &tree {
            TokenTree::Punct(punct) if punct.as_char() == ',' => segments.push(Vec::new()),
            _ => segments
                .last_mut()
                .expect("`segments` is never empty")
                .push(tree),
        }
    }

    segments
}

fn error(span: Span, message: &str) -> TokenStream2 {
    quote_spanned!(span=> ::core::compile_error!{#message})
}

/// Emit `errors` ahead of `item` so that the item still parses and only the
/// intended diagnostic is reported.
fn prepend(errors: TokenStream2, item: TokenStream2) -> TokenStream2 {
    quote! {
        #errors
        #item
    }
}