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