ambit_macros/

lib.rs

use proc_macro2::{Span, TokenStream};
use quote::{ToTokens, quote, quote_spanned};
use syn::{
    Attribute, Error, Fields, Ident, ItemStruct, LitInt, LitStr, Path, Result, Token, Visibility,
    parse::{Parse, ParseStream},
    parse_macro_input,
    spanned::Spanned,
};

#[proc_macro_attribute]
pub fn ambit(
    attr: proc_macro::TokenStream,
    item: proc_macro::TokenStream,
) -> proc_macro::TokenStream {
    ambit_impl(
        parse_macro_input!(attr as Args),
        parse_macro_input!(item as Decl),
    )
    .into()
}

fn ambit_impl(args: Args, decl: Decl) -> TokenStream {
    let (decl_attrs, decl_vis, decl_ident) = (decl.attrs, decl.vis, decl.ident);
    let name_lower = decl_ident.to_string().to_lowercase();

    let Args { range } = args;
    let (start, end) = (range.start, range.end);
    let (range_kind, range_repr) = (&range.kind, &range.repr);
    let is_nonzero = range_repr.raw.contains("NonZero");

    let error_ident = Ident::new(&format!("{decl_ident}Error"), decl_ident.span());
    let error_fstring = format!("invalid {name_lower}: {{}}");

    let iter_ident = Ident::new(&format!("{decl_ident}Iter"), decl_ident.span());

    let macro_ident = Ident::new(&name_lower, decl_ident.span());
    let macro_arms = (start.value..=end.value)
        .map(proc_macro2::Literal::u64_unsuffixed)
        .map(|n| quote! { (#n) => { unsafe { #decl_ident::new_unchecked(#n) } }; });

    let new_unchecked_method_body = if is_nonzero {
        quote! { Self(unsafe { <#range_repr>::new_unchecked(n) }) }
    } else {
        quote! { Self(n) }
    };

    let value_method_body = if is_nonzero {
        quote! { self.0.get() }
    } else {
        quote! { self.0 }
    };

    quote! {
        #(#decl_attrs)*
        #[repr(transparent)]
        #decl_vis struct #decl_ident(#range_repr);

        impl #decl_ident {
            pub const MIN_VALUE: #range_kind = #start;
            pub const MAX_VALUE: #range_kind = #end;

            pub const MIN: Self = unsafe { Self::new_unchecked(#start) };
            pub const MAX: Self = unsafe { Self::new_unchecked(#end) };

            pub fn iter() -> #iter_ident {
                #iter_ident(Some(Self::MIN))
            }

            pub const fn new(n: #range_kind) -> Result<Self, #error_ident> {
                if n < #start || n > #end {
                    return Err(#error_ident(n));
                }
                Ok(unsafe { Self::new_unchecked(n) })
            }

            pub const unsafe fn new_unchecked(n: #range_kind) -> Self {
                #new_unchecked_method_body
            }

            pub const fn value(&self) -> #range_kind {
                #value_method_body
            }

            pub const fn is_min(&self) -> bool {
                self.value() == Self::MIN.value()
            }

            pub const fn is_max(&self) -> bool {
                self.value() == Self::MAX.value()
            }

            pub fn pred(&self) -> Option<Self> {
                (!self.is_min()).then(|| unsafe { Self::new_unchecked(self.value() - 1 ) })
            }

            pub fn succ(&self) -> Option<Self> {
                (!self.is_max()).then(|| unsafe { Self::new_unchecked(self.value() + 1 ) })
            }
        }

        impl PartialEq<#range_kind> for #decl_ident {
            fn eq(&self, &other: &#range_kind) -> bool {
                self.value() == other
            }
        }

        impl TryFrom<#range_kind> for #decl_ident {
            type Error = #error_ident;

            fn try_from(n: #range_kind) -> Result<Self, #error_ident> {
                Self::new(n)
            }
        }

        #decl_vis struct #iter_ident(Option<#decl_ident>);

        impl Iterator for #iter_ident {
            type Item = #decl_ident;

            fn next(&mut self) -> Option<#decl_ident> {
                if let Some(next) = self.0.as_ref()?.succ() {
                    self.0.replace(next)
                } else {
                    self.0.take()
                }
            }
        }

        #[non_exhaustive]
        #[derive(Debug)]
        #decl_vis struct #error_ident(#decl_vis #range_kind);

        impl std::error::Error for #error_ident {}
        impl std::fmt::Display for #error_ident {
            fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
                write!(f, #error_fstring, self.0)
            }
        }

        #[macro_export]
        macro_rules! #macro_ident {
            #(#macro_arms)*
        }
    }
}

struct Decl {
    attrs: Vec<Attribute>,
    vis: Visibility,
    ident: Ident,
}

impl Parse for Decl {
    fn parse(input: ParseStream) -> Result<Self> {
        let item: ItemStruct = input.parse()?;
        let (attrs, vis, ident) = (item.attrs, item.vis, item.ident);

        let None = item.generics.lt_token else {
            return Err(Error::new(item.generics.span(), "Expected concrete struct"));
        };

        let Fields::Unit = item.fields else {
            return Err(Error::new(item.fields.span(), "Expected unit struct"));
        };

        Ok(Self { attrs, vis, ident })
    }
}

struct Args {
    range: Range,
}

impl Parse for Args {
    fn parse(input: ParseStream) -> Result<Self> {
        mod kw {
            syn::custom_keyword!(range);
        }

        let _: kw::range = input.parse()?;
        let _: Token![=] = input.parse()?;
        let s: LitStr = input.parse()?;
        let range = syn::parse_str(&s.value())?;

        Ok(Self { range })
    }
}

struct Range {
    start: Bound,
    end: Bound, // inclusive
    kind: ParsedStr<Ident>,
    repr: ParsedStr<Path>,
}

impl Parse for Range {
    fn parse(input: ParseStream) -> Result<Self> {
        let start: Bound = LitInt::parse(input)?.try_into()?;
        let end = if input.peek(Token![..=]) {
            let _ = <Token![..=]>::parse(input)?;
            LitInt::parse(input)?.try_into()?
        } else {
            let _ = <Token![..]>::parse(input)?;
            let mut bound: Bound = LitInt::parse(input)?.try_into()?;
            bound.value -= 1;
            bound
        };

        if end.value <= start.value {
            return Err(Error::new(input.span(), "expected ascending range"));
        }

        macro_rules! fits {
            ($ty:ident) => {
                $ty::try_from(end.value).is_ok()
            };
        }

        let kind = ParsedStr::new(match () {
            () if fits!(u8) => "u8",
            () if fits!(u16) => "u16",
            () if fits!(u32) => "u32",
            _ => "u64",
        })?;

        let repr = ParsedStr::new(match (kind.raw, start.value == 0) {
            (s, true) => s,
            ("u8", false) => "::std::num::NonZero<u8>",
            ("u16", false) => "::std::num::NonZero<u16>",
            ("u32", false) => "::std::num::NonZero<u32>",
            ("u64", false) => "::std::num::NonZero<u64>",
            (_, false) => unreachable!(),
        })?;

        Ok(Self {
            start,
            end,
            kind,
            repr,
        })
    }
}

struct Bound {
    value: u64,
    span: Span,
}

impl TryFrom<LitInt> for Bound {
    type Error = Error;

    fn try_from(int: LitInt) -> Result<Self> {
        Ok(Self {
            value: int.base10_parse()?,
            span: int.span(),
        })
    }
}

impl ToTokens for Bound {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        let lit = proc_macro2::Literal::u64_unsuffixed(self.value);
        tokens.extend(quote_spanned!(self.span => #lit));
    }
}

struct ParsedStr<T> {
    raw: &'static str,
    parsed: T,
}

impl<T: Parse> ParsedStr<T> {
    fn new(s: &'static str) -> Result<Self> {
        let parsed = syn::parse_str(s)?;
        Ok(Self { raw: s, parsed })
    }
}

impl<T: ToTokens> ToTokens for ParsedStr<T> {
    fn to_tokens(&self, tokens: &mut TokenStream) {
        self.parsed.to_tokens(tokens);
    }
}