[rustc.git] / compiler / rustc_macros / src / newtype.rs

use proc_macro2::{Span, TokenStream};
use quote::quote;
use syn::parse::*;
use syn::*;

// We parse the input and emit the output in a single step.
// This field stores the final macro output
struct Newtype(TokenStream);

impl Parse for Newtype {
    fn parse(input: ParseStream<'_>) -> Result<Self> {
        let mut attrs = input.call(Attribute::parse_outer)?;
        let vis: Visibility = input.parse()?;
        input.parse::<Token![struct]>()?;
        let name: Ident = input.parse()?;

        let body;
        braced!(body in input);

        // Any additional `#[derive]` macro paths to apply
        let mut derive_paths: Vec<Path> = Vec::new();
        let mut debug_format: Option<Lit> = None;
        let mut max = None;
        let mut consts = Vec::new();
        let mut encodable = true;
        let mut ord = true;

        attrs.retain(|attr| match attr.path.get_ident() {
            Some(ident) => match &*ident.to_string() {
                "custom_encodable" => {
                    encodable = false;
                    false
                }
                "no_ord_impl" => {
                    ord = false;
                    false
                }
                "max" => {
                    let Ok(Meta::NameValue(literal) )= attr.parse_meta() else {
                        panic!("#[max = NUMBER] attribute requires max value");
                    };

                    if let Some(old) = max.replace(literal.lit) {
                        panic!("Specified multiple max: {old:?}");
                    }

                    false
                }
                "debug_format" => {
                    let Ok(Meta::NameValue(literal) )= attr.parse_meta() else {
                        panic!("#[debug_format = FMT] attribute requires a format");
                    };

                    if let Some(old) = debug_format.replace(literal.lit) {
                        panic!("Specified multiple debug format options: {old:?}");
                    }

                    false
                }
                _ => true,
            },
            _ => true,
        });

        loop {
            // We've parsed everything that the user provided, so we're done
            if body.is_empty() {
                break;
            }

            // Otherwise, we are parsing a user-defined constant
            let const_attrs = body.call(Attribute::parse_outer)?;
            body.parse::<Token![const]>()?;
            let const_name: Ident = body.parse()?;
            body.parse::<Token![=]>()?;
            let const_val: Expr = body.parse()?;
            body.parse::<Token![;]>()?;
            consts.push(quote! { #(#const_attrs)* #vis const #const_name: #name = #name::from_u32(#const_val); });
        }

        let debug_format =
            debug_format.unwrap_or_else(|| Lit::Str(LitStr::new("{}", Span::call_site())));

        // shave off 256 indices at the end to allow space for packing these indices into enums
        let max = max.unwrap_or_else(|| Lit::Int(LitInt::new("0xFFFF_FF00", Span::call_site())));

        let encodable_impls = if encodable {
            quote! {
                impl<D: ::rustc_serialize::Decoder> ::rustc_serialize::Decodable<D> for #name {
                    fn decode(d: &mut D) -> Self {
                        Self::from_u32(d.read_u32())
                    }
                }
                impl<E: ::rustc_serialize::Encoder> ::rustc_serialize::Encodable<E> for #name {
                    fn encode(&self, e: &mut E) {
                        e.emit_u32(self.private);
                    }
                }
            }
        } else {
            quote! {}
        };

        if ord {
            derive_paths.push(parse_quote!(Ord));
            derive_paths.push(parse_quote!(PartialOrd));
        }

        let step = if ord {
            quote! {
                impl ::std::iter::Step for #name {
                    #[inline]
                    fn steps_between(start: &Self, end: &Self) -> Option<usize> {
                        <usize as ::std::iter::Step>::steps_between(
                            &Self::index(*start),
                            &Self::index(*end),
                        )
                    }

                    #[inline]
                    fn forward_checked(start: Self, u: usize) -> Option<Self> {
                        Self::index(start).checked_add(u).map(Self::from_usize)
                    }

                    #[inline]
                    fn backward_checked(start: Self, u: usize) -> Option<Self> {
                        Self::index(start).checked_sub(u).map(Self::from_usize)
                    }
                }

                // Safety: The implementation of `Step` upholds all invariants.
                unsafe impl ::std::iter::TrustedStep for #name {}
            }
        } else {
            quote! {}
        };

        let debug_impl = quote! {
            impl ::std::fmt::Debug for #name {
                fn fmt(&self, fmt: &mut ::std::fmt::Formatter<'_>) -> ::std::fmt::Result {
                    write!(fmt, #debug_format, self.as_u32())
                }
            }
        };

        let spec_partial_eq_impl = if let Lit::Int(max) = &max {
            if let Ok(max_val) = max.base10_parse::<u32>() {
                quote! {
                    impl core::option::SpecOptionPartialEq for #name {
                        #[inline]
                        fn eq(l: &Option<Self>, r: &Option<Self>) -> bool {
                            if #max_val < u32::MAX {
                                l.map(|i| i.private).unwrap_or(#max_val+1) == r.map(|i| i.private).unwrap_or(#max_val+1)
                            } else {
                                match (l, r) {
                                    (Some(l), Some(r)) => r == l,
                                    (None, None) => true,
                                    _ => false
                                }
                            }
                        }
                    }
                }
            } else {
                quote! {}
            }
        } else {
            quote! {}
        };

        Ok(Self(quote! {
            #(#attrs)*
            #[derive(Clone, Copy, PartialEq, Eq, Hash, #(#derive_paths),*)]
            #[rustc_layout_scalar_valid_range_end(#max)]
            #[rustc_pass_by_value]
            #vis struct #name {
                private: u32,
            }

            #(#consts)*

            impl #name {
                /// Maximum value the index can take, as a `u32`.
                #vis const MAX_AS_U32: u32  = #max;

                /// Maximum value the index can take.
                #vis const MAX: Self = Self::from_u32(#max);

                /// Creates a new index from a given `usize`.
                ///
                /// # Panics
                ///
                /// Will panic if `value` exceeds `MAX`.
                #[inline]
                #vis const fn from_usize(value: usize) -> Self {
                    assert!(value <= (#max as usize));
                    // SAFETY: We just checked that `value <= max`.
                    unsafe {
                        Self::from_u32_unchecked(value as u32)
                    }
                }

                /// Creates a new index from a given `u32`.
                ///
                /// # Panics
                ///
                /// Will panic if `value` exceeds `MAX`.
                #[inline]
                #vis const fn from_u32(value: u32) -> Self {
                    assert!(value <= #max);
                    // SAFETY: We just checked that `value <= max`.
                    unsafe {
                        Self::from_u32_unchecked(value)
                    }
                }

                /// Creates a new index from a given `u32`.
                ///
                /// # Safety
                ///
                /// The provided value must be less than or equal to the maximum value for the newtype.
                /// Providing a value outside this range is undefined due to layout restrictions.
                ///
                /// Prefer using `from_u32`.
                #[inline]
                #vis const unsafe fn from_u32_unchecked(value: u32) -> Self {
                    Self { private: value }
                }

                /// Extracts the value of this index as a `usize`.
                #[inline]
                #vis const fn index(self) -> usize {
                    self.as_usize()
                }

                /// Extracts the value of this index as a `u32`.
                #[inline]
                #vis const fn as_u32(self) -> u32 {
                    self.private
                }

                /// Extracts the value of this index as a `usize`.
                #[inline]
                #vis const fn as_usize(self) -> usize {
                    self.as_u32() as usize
                }
            }

            impl std::ops::Add<usize> for #name {
                type Output = Self;

                fn add(self, other: usize) -> Self {
                    Self::from_usize(self.index() + other)
                }
            }

            impl rustc_index::vec::Idx for #name {
                #[inline]
                fn new(value: usize) -> Self {
                    Self::from_usize(value)
                }

                #[inline]
                fn index(self) -> usize {
                    self.as_usize()
                }
            }

            #step

            #spec_partial_eq_impl

            impl From<#name> for u32 {
                #[inline]
                fn from(v: #name) -> u32 {
                    v.as_u32()
                }
            }

            impl From<#name> for usize {
                #[inline]
                fn from(v: #name) -> usize {
                    v.as_usize()
                }
            }

            impl From<usize> for #name {
                #[inline]
                fn from(value: usize) -> Self {
                    Self::from_usize(value)
                }
            }

            impl From<u32> for #name {
                #[inline]
                fn from(value: u32) -> Self {
                    Self::from_u32(value)
                }
            }

            #encodable_impls
            #debug_impl
        }))
    }
}

pub fn newtype(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let input = parse_macro_input!(input as Newtype);
    input.0.into()
}
Commit	Line	Data
5e7ed085 FG	1	use proc_macro2::{Span, TokenStream};
	2	use quote::quote;
	3	use syn::parse::*;
5e7ed085 FG	4	use syn::*;
5e7ed085 FG	5
5e7ed085 FG	6	// We parse the input and emit the output in a single step.
	7	// This field stores the final macro output
	8	struct Newtype(TokenStream);
	9
	10	impl Parse for Newtype {
	11	fn parse(input: ParseStream<'_>) -> Result<Self> {
9c376795	12	let mut attrs = input.call(Attribute::parse_outer)?;
5e7ed085 FG	13	let vis: Visibility = input.parse()?;
	14	input.parse::<Token![struct]>()?;
	15	let name: Ident = input.parse()?;
	16
	17	let body;
	18	braced!(body in input);
	19
	20	// Any additional `#[derive]` macro paths to apply
	21	let mut derive_paths: Vec<Path> = Vec::new();
9c376795	22	let mut debug_format: Option<Lit> = None;
5e7ed085 FG	23	let mut max = None;
	24	let mut consts = Vec::new();
	25	let mut encodable = true;
	26	let mut ord = true;
	27
9c376795 FG	28	attrs.retain(\|attr\| match attr.path.get_ident() {
	29	Some(ident) => match &*ident.to_string() {
	30	"custom_encodable" => {
5e7ed085	31	encodable = false;
9c376795	32	false
5e7ed085	33	}
9c376795	34	"no_ord_impl" => {
5e7ed085	35	ord = false;
9c376795	36	false
5e7ed085	37	}
9c376795 FG	38	"max" => {
	39	let Ok(Meta::NameValue(literal) )= attr.parse_meta() else {
	40	panic!("#[max = NUMBER] attribute requires max value");
	41	};
5e7ed085	42
9c376795 FG	43	if let Some(old) = max.replace(literal.lit) {
	44	panic!("Specified multiple max: {old:?}");
	45	}
	46
	47	false
5e7ed085	48	}
9c376795 FG	49	"debug_format" => {
	50	let Ok(Meta::NameValue(literal) )= attr.parse_meta() else {
	51	panic!("#[debug_format = FMT] attribute requires a format");
	52	};
5e7ed085	53
9c376795 FG	54	if let Some(old) = debug_format.replace(literal.lit) {
	55	panic!("Specified multiple debug format options: {old:?}");
	56	}
	57
	58	false
	59	}
	60	_ => true,
	61	},
	62	_ => true,
	63	});
	64
	65	loop {
	66	// We've parsed everything that the user provided, so we're done
	67	if body.is_empty() {
	68	break;
5e7ed085	69	}
9c376795 FG	70
	71	// Otherwise, we are parsing a user-defined constant
	72	let const_attrs = body.call(Attribute::parse_outer)?;
	73	body.parse::<Token![const]>()?;
	74	let const_name: Ident = body.parse()?;
	75	body.parse::<Token![=]>()?;
	76	let const_val: Expr = body.parse()?;
	77	body.parse::<Token![;]>()?;
	78	consts.push(quote! { #(#const_attrs)* #vis const #const_name: #name = #name::from_u32(#const_val); });
5e7ed085 FG	79	}
5e7ed085 FG	80
9c376795 FG	81	let debug_format =
	82	debug_format.unwrap_or_else(\|\| Lit::Str(LitStr::new("{}", Span::call_site())));
	83
5e7ed085 FG	84	// shave off 256 indices at the end to allow space for packing these indices into enums
	85	let max = max.unwrap_or_else(\|\| Lit::Int(LitInt::new("0xFFFF_FF00", Span::call_site())));
	86
	87	let encodable_impls = if encodable {
	88	quote! {
	89	impl<D: ::rustc_serialize::Decoder> ::rustc_serialize::Decodable<D> for #name {
	90	fn decode(d: &mut D) -> Self {
	91	Self::from_u32(d.read_u32())
	92	}
	93	}
	94	impl<E: ::rustc_serialize::Encoder> ::rustc_serialize::Encodable<E> for #name {
923072b8 FG	95	fn encode(&self, e: &mut E) {
923072b8 FG	96	e.emit_u32(self.private);
5e7ed085 FG	97	}
	98	}
	99	}
	100	} else {
	101	quote! {}
	102	};
	103
	104	if ord {
	105	derive_paths.push(parse_quote!(Ord));
	106	derive_paths.push(parse_quote!(PartialOrd));
	107	}
	108
	109	let step = if ord {
	110	quote! {
	111	impl ::std::iter::Step for #name {
	112	#[inline]
	113	fn steps_between(start: &Self, end: &Self) -> Option<usize> {
	114	<usize as ::std::iter::Step>::steps_between(
	115	&Self::index(*start),
	116	&Self::index(*end),
	117	)
	118	}
	119
	120	#[inline]
	121	fn forward_checked(start: Self, u: usize) -> Option<Self> {
	122	Self::index(start).checked_add(u).map(Self::from_usize)
	123	}
	124
	125	#[inline]
	126	fn backward_checked(start: Self, u: usize) -> Option<Self> {
	127	Self::index(start).checked_sub(u).map(Self::from_usize)
	128	}
	129	}
	130
	131	// Safety: The implementation of `Step` upholds all invariants.
	132	unsafe impl ::std::iter::TrustedStep for #name {}
	133	}
	134	} else {
	135	quote! {}
	136	};
	137
9c376795 FG	138	let debug_impl = quote! {
	139	impl ::std::fmt::Debug for #name {
	140	fn fmt(&self, fmt: &mut ::std::fmt::Formatter<'_>) -> ::std::fmt::Result {
	141	write!(fmt, #debug_format, self.as_u32())
5e7ed085 FG	142	}
	143	}
	144	};
9c376795	145
487cf647 FG	146	let spec_partial_eq_impl = if let Lit::Int(max) = &max {
	147	if let Ok(max_val) = max.base10_parse::<u32>() {
	148	quote! {
	149	impl core::option::SpecOptionPartialEq for #name {
	150	#[inline]
	151	fn eq(l: &Option<Self>, r: &Option<Self>) -> bool {
	152	if #max_val < u32::MAX {
	153	l.map(\|i\| i.private).unwrap_or(#max_val+1) == r.map(\|i\| i.private).unwrap_or(#max_val+1)
	154	} else {
	155	match (l, r) {
	156	(Some(l), Some(r)) => r == l,
	157	(None, None) => true,
	158	_ => false
	159	}
	160	}
	161	}
	162	}
	163	}
	164	} else {
	165	quote! {}
	166	}
	167	} else {
	168	quote! {}
	169	};
5e7ed085 FG	170
	171	Ok(Self(quote! {
	172	#(#attrs)*
	173	#[derive(Clone, Copy, PartialEq, Eq, Hash, #(#derive_paths),*)]
	174	#[rustc_layout_scalar_valid_range_end(#max)]
	175	#[rustc_pass_by_value]
	176	#vis struct #name {
	177	private: u32,
	178	}
	179
	180	#(#consts)*
	181
	182	impl #name {
	183	/// Maximum value the index can take, as a `u32`.
	184	#vis const MAX_AS_U32: u32 = #max;
	185
	186	/// Maximum value the index can take.
	187	#vis const MAX: Self = Self::from_u32(#max);
	188
	189	/// Creates a new index from a given `usize`.
	190	///
	191	/// # Panics
	192	///
	193	/// Will panic if `value` exceeds `MAX`.
	194	#[inline]
	195	#vis const fn from_usize(value: usize) -> Self {
	196	assert!(value <= (#max as usize));
	197	// SAFETY: We just checked that `value <= max`.
	198	unsafe {
	199	Self::from_u32_unchecked(value as u32)
	200	}
	201	}
	202
	203	/// Creates a new index from a given `u32`.
	204	///
	205	/// # Panics
	206	///
	207	/// Will panic if `value` exceeds `MAX`.
	208	#[inline]
	209	#vis const fn from_u32(value: u32) -> Self {
	210	assert!(value <= #max);
	211	// SAFETY: We just checked that `value <= max`.
	212	unsafe {
	213	Self::from_u32_unchecked(value)
	214	}
	215	}
	216
	217	/// Creates a new index from a given `u32`.
	218	///
	219	/// # Safety
	220	///
	221	/// The provided value must be less than or equal to the maximum value for the newtype.
	222	/// Providing a value outside this range is undefined due to layout restrictions.
	223	///
	224	/// Prefer using `from_u32`.
	225	#[inline]
	226	#vis const unsafe fn from_u32_unchecked(value: u32) -> Self {
	227	Self { private: value }
	228	}
	229
	230	/// Extracts the value of this index as a `usize`.
	231	#[inline]
	232	#vis const fn index(self) -> usize {
	233	self.as_usize()
234	}
235
236	/// Extracts the value of this index as a `u32`.
237	#[inline]
238	#vis const fn as_u32(self) -> u32 {
239	self.private
240	}
241
242	/// Extracts the value of this index as a `usize`.
243	#[inline]
244	#vis const fn as_usize(self) -> usize {
245	self.as_u32() as usize
246	}
247	}
248
249	impl std::ops::Add<usize> for #name {
250	type Output = Self;
251
252	fn add(self, other: usize) -> Self {
253	Self::from_usize(self.index() + other)
254	}
255	}
256
257	impl rustc_index::vec::Idx for #name {
258	#[inline]
259	fn new(value: usize) -> Self {
260	Self::from_usize(value)
261	}
262
263	#[inline]
264	fn index(self) -> usize {
265	self.as_usize()
266	}
267	}
268
269	#step
270
487cf647 FG	271	#spec_partial_eq_impl
487cf647 FG	272
5e7ed085 FG	273	impl From<#name> for u32 {
	274	#[inline]
	275	fn from(v: #name) -> u32 {
	276	v.as_u32()
	277	}
	278	}
	279
	280	impl From<#name> for usize {
	281	#[inline]
	282	fn from(v: #name) -> usize {
	283	v.as_usize()
	284	}
	285	}
	286
	287	impl From<usize> for #name {
	288	#[inline]
	289	fn from(value: usize) -> Self {
	290	Self::from_usize(value)
	291	}
	292	}
	293
	294	impl From<u32> for #name {
	295	#[inline]
	296	fn from(value: u32) -> Self {
	297	Self::from_u32(value)
	298	}
	299	}
	300
	301	#encodable_impls
	302	#debug_impl
	303	}))
	304	}
	305	}
	306
	307	pub fn newtype(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
	308	let input = parse_macro_input!(input as Newtype);
	309	input.0.into()
	310	}