enum-path-derive 0.1.0

#![allow(
    clippy::needless_continue,
    reason = "originates in darling::FromAttributes derive macro"
)]

use darling::{FromAttributes, FromMeta};
use proc_macro2::TokenStream;
use quote::quote;
use syn::{DeriveInput, Fields};

// ---------------------------------------------------------------------------
// Attribute types
// ---------------------------------------------------------------------------

#[derive(Clone, Copy, Debug, PartialEq, Eq, FromMeta)]
enum RenameAll
{
    #[darling(rename = "lowercase")]
    Lower,
    #[darling(rename = "UPPERCASE")]
    Upper,
    #[darling(rename = "PascalCase")]
    Pascal,
    #[darling(rename = "camelCase")]
    Camel,
    #[darling(rename = "snake_case")]
    Snake,
    #[darling(rename = "SCREAMING_SNAKE_CASE")]
    ScreamingSnake,
    #[darling(rename = "kebab-case")]
    Kebab,
    #[darling(rename = "SCREAMING-KEBAB-CASE")]
    ScreamingKebab,
}

/// Enum-level attributes: `#[enum_path(FromStr, Display, ...)]`.
#[derive(Debug, FromAttributes)]
#[darling(attributes(enum_path))]
struct EnumAttrs
{
    rename_all: Option<RenameAll>,
    #[darling(default = "default_delimiter")]
    delimiter: String,
    #[darling(rename = "FromStr", default)]
    from_str: darling::util::Flag,
    #[darling(rename = "Display", default)]
    display: darling::util::Flag,
    #[darling(default)]
    case_insensitive: bool,
    #[darling(default)]
    error: Option<syn::Path>,
    /// Overrides the path of the runtime crate; defaults to `::enum_path`
    /// when absent. Useful when the user re-exports `enum_path` under a
    /// different name.
    #[darling(rename = "crate", default)]
    crate_path: Option<syn::Path>,
}

/// Variant-level attributes: `#[enum_path(rename = "...")]`.
#[derive(Debug, Default, FromAttributes)]
#[darling(attributes(enum_path), default)]
struct VariantAttrs
{
    rename: Option<String>,
}

fn default_delimiter() -> String
{
    ".".to_owned()
}

// ---------------------------------------------------------------------------
// Word splitting
// ---------------------------------------------------------------------------

/// Split a string into its constituent words.
///
/// `"SendMessage"` -> `["Send", "Message"]`.
fn split_words(input: impl AsRef<str>) -> Vec<String>
{
    let mut words = Vec::new();
    let mut current = String::with_capacity(input.as_ref().len());
    let mut seen_lowercase = false;
    let mut seen_digit = false;

    for ch in input.as_ref().chars() {
        if ch.is_lowercase() && !seen_digit {
            seen_lowercase = true;
        } else if ch.is_ascii_digit() {
            seen_digit = true;
        } else if seen_lowercase || !ch.is_alphanumeric() || (seen_digit && !ch.is_ascii_digit()) {
            words.push(capitalize(&current));
            seen_lowercase = false;
            seen_digit = false;
            current.clear();
        } else if !current.is_empty() {
            current.extend(ch.to_lowercase());
            continue;
        }

        if ch.is_alphanumeric() {
            current.push(ch);
        }
    }
    if !current.is_empty() {
        words.push(capitalize(current));
    }
    words
}

impl RenameAll
{
    fn apply(self, ident: impl AsRef<str>) -> String
    {
        let words = split_words(&ident);
        match self {
            Self::Lower => ident.as_ref().to_lowercase(),
            Self::Upper => ident.as_ref().to_uppercase(),
            Self::Pascal => words.join(""),
            Self::Camel => uncapitalize(words.join("")),
            Self::Snake => words.join("_").to_lowercase(),
            Self::ScreamingSnake => words.join("_").to_uppercase(),
            Self::Kebab => words.join("-").to_lowercase(),
            Self::ScreamingKebab => words.join("-").to_uppercase(),
        }
    }
}

fn capitalize(s: impl AsRef<str>) -> String
{
    let mut chars = s.as_ref().chars();
    match chars.next() {
        None => String::new(),
        Some(c) => {
            let mut ret = String::with_capacity(s.as_ref().len());
            let mut start_chars = c.to_uppercase();
            ret.extend(start_chars.by_ref().take(1));
            ret.extend(start_chars.flat_map(char::to_lowercase));
            for ch in chars {
                ret.extend(ch.to_lowercase());
            }
            ret
        }
    }
}

fn uncapitalize(s: impl AsRef<str>) -> String
{
    let mut chars = s.as_ref().chars();
    match chars.next() {
        None => String::new(),
        Some(c) => {
            let mut ret = String::with_capacity(s.as_ref().len());
            ret.extend(c.to_lowercase());
            ret.extend(chars);
            ret
        }
    }
}

/// Compute the serialized name for a variant, respecting per-variant `rename`
/// and enum-level `rename_all`.
fn variant_name(
    ident: &syn::Ident,
    variant_attrs: &VariantAttrs,
    rename_all: Option<RenameAll>,
) -> String
{
    if let Some(explicit) = &variant_attrs.rename {
        return explicit.clone();
    }
    match rename_all {
        Some(r) => r.apply(ident.to_string()),
        None => ident.to_string(),
    }
}

// ---------------------------------------------------------------------------
// Variant classification
// ---------------------------------------------------------------------------

enum VariantKind<'a>
{
    /// A unit variant like `Exit`.
    Unit,
    /// A single unnamed field like `SendMessage(String)`.
    Single(&'a syn::Type),
}

fn classify_variant(variant: &syn::Variant) -> Result<VariantKind<'_>, syn::Error>
{
    match &variant.fields {
        Fields::Unit => Ok(VariantKind::Unit),
        Fields::Unnamed(fields) if fields.unnamed.len() == 1 => {
            Ok(VariantKind::Single(&fields.unnamed[0].ty))
        }
        _ => Err(syn::Error::new_spanned(
            &variant.ident,
            "EnumPath variants must be unit or single-field tuple variants",
        )),
    }
}

/// Parse variant attrs and classify in one step.
fn resolve_variant(
    v: &syn::Variant,
    rename_all: Option<RenameAll>,
) -> Result<(String, VariantKind<'_>), TokenStream>
{
    let vattrs = VariantAttrs::from_attributes(&v.attrs).map_err(darling::Error::write_errors)?;
    let name = variant_name(&v.ident, &vattrs, rename_all);
    let kind = classify_variant(v).map_err(|e| e.to_compile_error())?;
    Ok((name, kind))
}

/// Resolve the error type. If the user specified a custom type, use it;
/// otherwise default to `::enum_path::Error`.
fn error_type(attrs: &EnumAttrs, crate_path: &syn::Path) -> TokenStream
{
    if let Some(path) = &attrs.error {
        quote! { #path }
    } else {
        quote! { #crate_path::Error }
    }
}

/// Resolve the runtime crate path; defaults to `::enum_path` when the
/// user did not specify `#[enum_path(crate = ...)]`.
fn resolve_crate_path(attrs: &EnumAttrs) -> syn::Path
{
    attrs
        .crate_path
        .clone()
        .unwrap_or_else(|| syn::parse_quote!(::enum_path))
}

// ---------------------------------------------------------------------------
// String comparison helper
// ---------------------------------------------------------------------------

/// Produce the comparison expression for matching a variant name.
///
/// When `case_insensitive` is set, the generated code uses
/// `str::eq_ignore_ascii_case`; otherwise a plain `==`.
fn eq_expr(case_insensitive: bool) -> TokenStream
{
    if case_insensitive {
        quote! { <::core::primitive::str>::eq_ignore_ascii_case }
    } else {
        quote! { <::core::primitive::str as ::core::cmp::PartialEq>::eq }
    }
}

/// Produce an expression that tests whether `haystack` starts with `prefix`
/// and yields the remainder.
///
/// Case-sensitive mode uses `strip_prefix`; case-insensitive mode checks
/// length and `eq_ignore_ascii_case` on the prefix slice.
fn strip_prefix_expr(case_insensitive: bool, prefix: &str) -> TokenStream
{
    if case_insensitive {
        let len = prefix.len();
        quote! {
            if s.len() >= #len
                && <::core::primitive::str>::eq_ignore_ascii_case(&s[..#len], #prefix)
            {
                ::core::option::Option::Some(&s[#len..])
            } else {
                ::core::option::Option::None
            }
        }
    } else {
        quote! { <::core::primitive::str>::strip_prefix(s, #prefix) }
    }
}

// ---------------------------------------------------------------------------
// FromStr code generation
// ---------------------------------------------------------------------------

/// Per-enum context shared by every `FromStr` variant arm.
struct FromStrCtx<'a>
{
    type_name: &'a syn::Ident,
    delimiter: &'a str,
    case_insensitive: bool,
    err_ty: TokenStream,
    crate_path: syn::Path,
}

fn gen_from_str(type_name: &syn::Ident, data: &syn::DataEnum, attrs: &EnumAttrs) -> TokenStream
{
    let crate_path = resolve_crate_path(attrs);
    let ctx = FromStrCtx {
        type_name,
        delimiter: &attrs.delimiter,
        case_insensitive: attrs.case_insensitive,
        err_ty: error_type(attrs, &crate_path),
        crate_path,
    };
    let type_str = type_name.to_string();

    let arms = data.variants.iter().map(|v| {
        let ident = &v.ident;
        match resolve_variant(v, attrs.rename_all) {
            Ok((name, VariantKind::Unit)) => ctx.gen_unit(ident, &name),
            Ok((name, VariantKind::Single(ty))) => ctx.gen_single(ident, &name, ty),
            Err(e) => e,
        }
    });

    let err_ty = &ctx.err_ty;
    let crate_path = &ctx.crate_path;
    quote! {
        impl ::core::str::FromStr for #type_name {
            type Err = #err_ty;

            fn from_str(s: &::core::primitive::str)
                -> ::core::result::Result<Self, Self::Err>
            {
                #(#arms)*

                ::core::result::Result::Err(::core::convert::Into::into(
                    #crate_path::Error {
                        input: #crate_path::__private::ToOwned::to_owned(s),
                        expected: #type_str,
                    }
                ))
            }
        }
    }
}
impl FromStrCtx<'_>
{
    fn gen_unit(&self, ident: &syn::Ident, name: &str) -> TokenStream
    {
        let eq_fn = eq_expr(self.case_insensitive);
        let type_name = self.type_name;
        quote! {
            if #eq_fn(s, #name) {
                return ::core::result::Result::Ok(#type_name::#ident);
            }
        }
    }

    fn gen_single(&self, ident: &syn::Ident, name: &str, field_ty: &syn::Type) -> TokenStream
    {
        let prefix = format!("{name}{}", self.delimiter);
        let strip = strip_prefix_expr(self.case_insensitive, &prefix);
        let type_name = self.type_name;
        let err_ty = &self.err_ty;
        let crate_path = &self.crate_path;
        let type_str = type_name.to_string();

        quote! {
            if let ::core::option::Option::Some(rest) = #strip {
                let inner: #field_ty =
                    ::core::result::Result::map_err(
                        <#field_ty as ::core::str::FromStr>::from_str(rest),
                        |_: <#field_ty as ::core::str::FromStr>::Err| {
                            ::core::convert::Into::<#err_ty>::into(#crate_path::Error {
                                input: #crate_path::__private::ToOwned::to_owned(s),
                                expected: #type_str,
                            })
                        },
                    )?;
                return ::core::result::Result::Ok(#type_name::#ident(inner));
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Display code generation
// ---------------------------------------------------------------------------

/// Per-enum context shared by every `Display` variant arm.
struct DisplayCtx<'a>
{
    type_name: &'a syn::Ident,
    delimiter: &'a str,
}

fn gen_display(type_name: &syn::Ident, data: &syn::DataEnum, attrs: &EnumAttrs) -> TokenStream
{
    let ctx = DisplayCtx {
        type_name,
        delimiter: &attrs.delimiter,
    };
    let arms = data.variants.iter().map(|v| {
        let ident = &v.ident;
        match resolve_variant(v, attrs.rename_all) {
            Ok((name, VariantKind::Unit)) => ctx.gen_unit(ident, &name),
            Ok((name, VariantKind::Single(_))) => ctx.gen_single(ident, &name),
            Err(e) => e,
        }
    });

    quote! {
        impl ::core::fmt::Display for #type_name {
            fn fmt(&self, f: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result {
                match self {
                    #(#arms)*
                }
            }
        }
    }
}

impl DisplayCtx<'_>
{
    fn gen_unit(&self, ident: &syn::Ident, name: &str) -> TokenStream
    {
        let type_name = self.type_name;
        quote! {
            #type_name::#ident => ::core::fmt::Write::write_str(f, #name),
        }
    }

    fn gen_single(&self, ident: &syn::Ident, name: &str) -> TokenStream
    {
        let type_name = self.type_name;
        let delimiter = self.delimiter;
        quote! {
            #type_name::#ident(inner) => {
                ::core::fmt::Write::write_str(f, #name)?;
                ::core::fmt::Write::write_str(f, #delimiter)?;
                ::core::fmt::Display::fmt(inner, f)
            }
        }
    }
}

// ---------------------------------------------------------------------------
// Entry point
// ---------------------------------------------------------------------------

pub fn derive_enum_path(input: &DeriveInput) -> TokenStream
{
    let type_name = &input.ident;

    let syn::Data::Enum(data) = &input.data else {
        return syn::Error::new_spanned(&input.ident, "EnumPath can only be derived for enums")
            .to_compile_error();
    };

    let attrs = match EnumAttrs::from_attributes(&input.attrs) {
        Ok(a) => a,
        Err(e) => return e.write_errors(),
    };

    if let Err(e) = validate_variant_names(data, &attrs) {
        return e.to_compile_error();
    }

    let mut output = TokenStream::new();

    if attrs.from_str.is_present() {
        output.extend(gen_from_str(type_name, data, &attrs));
    }
    if attrs.display.is_present() {
        output.extend(gen_display(type_name, data, &attrs));
    }

    output
}

/// Reject duplicate variant names, plus any pair where one name is a
/// prefix of another at the configured delimiter boundary. Both cases
/// would make `FromStr` order-dependent on the source enum.
fn validate_variant_names(data: &syn::DataEnum, attrs: &EnumAttrs) -> syn::Result<()>
{
    let mut resolved: Vec<(&syn::Variant, String)> = Vec::with_capacity(data.variants.len());
    for variant in &data.variants {
        let vattrs = VariantAttrs::from_attributes(&variant.attrs)
            .map_err(|e| syn::Error::new_spanned(variant, e.to_string()))?;
        let name = variant_name(&variant.ident, &vattrs, attrs.rename_all);
        let canonical = if attrs.case_insensitive {
            name.to_lowercase()
        } else {
            name
        };
        resolved.push((variant, canonical));
    }

    for (i, (v_i, name_i)) in resolved.iter().enumerate() {
        for (v_j, name_j) in &resolved[..i] {
            if name_i == name_j {
                return Err(syn::Error::new_spanned(
                    &v_i.ident,
                    format!(
                        "duplicate EnumPath variant name {name_i:?} (also produced by `{}`)",
                        v_j.ident,
                    ),
                ));
            }
            let lhs = Resolved {
                variant: v_i,
                name:    name_i,
            };
            let rhs = Resolved {
                variant: v_j,
                name:    name_j,
            };
            check_prefix_conflict(lhs, rhs, &attrs.delimiter)?;
            check_prefix_conflict(rhs, lhs, &attrs.delimiter)?;
        }
    }
    Ok(())
}

/// A variant paired with its resolved serialized name.
#[derive(Clone, Copy)]
struct Resolved<'a>
{
    variant: &'a syn::Variant,
    name:    &'a str,
}

fn check_prefix_conflict(
    short: Resolved<'_>,
    long: Resolved<'_>,
    delimiter: &str,
) -> syn::Result<()>
{
    if long.name.len() <= short.name.len() {
        return Ok(());
    }
    if !long.name.starts_with(short.name) {
        return Ok(());
    }
    let suffix = &long.name[short.name.len()..];
    if suffix.starts_with(delimiter) {
        let short_name = short.name;
        let long_name = long.name;
        return Err(syn::Error::new_spanned(
            &long.variant.ident,
            format!(
                "EnumPath variant {long_name:?} would alias with {short_name:?} \
                 under delimiter {delimiter:?} (also produced by `{}`)",
                short.variant.ident,
            ),
        ));
    }
    Ok(())
}

#[cfg(test)]
mod tests
{
    use super::*;

    #[test]
    fn rename_all()
    {
        assert_eq!(RenameAll::Lower.apply("FooBAR"), "foobar");
        assert_eq!(RenameAll::Lower.apply("FooBAR_baz"), "foobar_baz");
        assert_eq!(RenameAll::Lower.apply("ƑōőƂɑρ2_βǎζ^qǚχ"), "ƒōőƃɑρ2_βǎζ^qǚχ");

        assert_eq!(RenameAll::Upper.apply("FooBAR"), "FOOBAR");
        assert_eq!(RenameAll::Upper.apply("FooBAR_baz"), "FOOBAR_BAZ");
        assert_eq!(RenameAll::Upper.apply("ƑōőƂɑρ2_βǎζ^qǚχ"), "ƑŌŐƂⱭΡ2_ΒǍΖ^QǙΧ");

        assert_eq!(RenameAll::Camel.apply("FooBAR"), "fooBar");
        assert_eq!(RenameAll::Camel.apply("FooBAR_baz"), "fooBarBaz");
        assert_eq!(RenameAll::Camel.apply("foo23bar"), "foo23Bar");
        assert_eq!(RenameAll::Camel.apply("foo23B5er7ry"), "foo23B5Er7Ry");
        assert_eq!(RenameAll::Camel.apply("ƑōőƂɑρ2_βǎζ^qǚχ"), "ƒōőƂɑρ2ΒǎζQǚχ");

        assert_eq!(RenameAll::Pascal.apply("FooBAR"), "FooBar");
        assert_eq!(RenameAll::Pascal.apply("FooBAR_baz"), "FooBarBaz");
        assert_eq!(RenameAll::Pascal.apply("ƑōőƂɑρ2_βǎζ^qǚχ"), "ƑōőƂɑρ2ΒǎζQǚχ");
        assert_eq!(RenameAll::Pascal.apply("fflOo"), "FflOo");

        assert_eq!(RenameAll::Snake.apply("FooBAR"), "foo_bar");
        assert_eq!(RenameAll::Snake.apply("FooBAR_baz"), "foo_bar_baz");
        assert_eq!(RenameAll::Snake.apply("foo23bar"), "foo23_bar");
        assert_eq!(RenameAll::Snake.apply("foo23B5er7ry"), "foo23_b5_er7_ry");
        assert_eq!(
            RenameAll::Snake.apply("ƑōőƂɑρ2_βǎζ^qǚχ"),
            "ƒōő_ƃɑρ2_βǎζ_qǚχ"
        );

        assert_eq!(RenameAll::ScreamingSnake.apply("FooBAR"), "FOO_BAR");
        assert_eq!(RenameAll::ScreamingSnake.apply("FooBAR_baz"), "FOO_BAR_BAZ");
        assert_eq!(
            RenameAll::ScreamingSnake.apply("ƑōőƂɑρ2_βǎζ^qǚχ"),
            "ƑŌŐ_ƂⱭΡ2_ΒǍΖ_QǙΧ"
        );

        assert_eq!(RenameAll::Kebab.apply("FooBAR"), "foo-bar");
        assert_eq!(RenameAll::Kebab.apply("FooBAR_baz"), "foo-bar-baz");
        assert_eq!(
            RenameAll::Kebab.apply("ƑōőƂɑρ2_βǎζ^qǚχ"),
            "ƒōő-ƃɑρ2-βǎζ-qǚχ"
        );

        assert_eq!(RenameAll::ScreamingKebab.apply("FooBAR"), "FOO-BAR");
        assert_eq!(RenameAll::ScreamingKebab.apply("FooBAR_baz"), "FOO-BAR-BAZ");
        assert_eq!(
            RenameAll::ScreamingKebab.apply("ƑōőƂɑρ2_βǎζ^qǚχ"),
            "ƑŌŐ-ƂⱭΡ2-ΒǍΖ-QǙΧ"
        );
    }
}