Introduce ruff_index crate (#4597)

crates/ruff_macros/src/lib.rs

@@ -1,6 +1,7 @@
 //! This crate implements internal macros for the `ruff` library.
 
 use crate::cache_key::derive_cache_key;
+use crate::newtype_index::generate_newtype_index;
 use proc_macro::TokenStream;
 use syn::{parse_macro_input, DeriveInput, ItemFn, ItemStruct};
 
@@ -9,6 +10,7 @@ mod combine_options;
 mod config;
 mod derive_message_formats;
 mod map_codes;
+mod newtype_index;
 mod register_rules;
 mod rule_code_prefix;
 mod rule_namespace;
@@ -79,3 +81,33 @@ pub fn derive_message_formats(_attr: TokenStream, item: TokenStream) -> TokenStr
     let func = parse_macro_input!(item as ItemFn);
     derive_message_formats::derive_message_formats(&func).into()
 }
+
+/// Derives a newtype wrapper that can be used as an index.
+/// The wrapper can represent indices up to `u32::MAX - 1`.
+///
+/// The `u32::MAX - 1` is an optimization so that `Option<Index>` has the same size as `Index`.
+///
+/// Can store at most `u32::MAX - 1` values
+///
+/// ## Warning
+///
+/// Additional `derive` attributes must come AFTER this attribute:
+///
+/// Good:
+///
+/// ```rust
+/// #[newtype_index]
+/// #[derive(Ord, PartialOrd)]
+/// struct MyIndex;
+/// ```
+#[proc_macro_attribute]
+pub fn newtype_index(_metadata: TokenStream, input: TokenStream) -> TokenStream {
+    let item = parse_macro_input!(input as ItemStruct);
+
+    let output = match generate_newtype_index(item) {
+        Ok(output) => output,
+        Err(err) => err.to_compile_error(),
+    };
+
+    TokenStream::from(output)
+}

crates/ruff_macros/src/newtype_index.rs

@@ -0,0 +1,139 @@
+use quote::quote;
+use syn::spanned::Spanned;
+use syn::{Error, ItemStruct};
+
+pub(super) fn generate_newtype_index(item: ItemStruct) -> syn::Result<proc_macro2::TokenStream> {
+    if !item.fields.is_empty() {
+        return Err(Error::new(
+            item.span(),
+            "A new type index cannot have any fields.",
+        ));
+    }
+
+    if !item.generics.params.is_empty() {
+        return Err(Error::new(
+            item.span(),
+            "A new type index cannot be generic.",
+        ));
+    }
+
+    let ItemStruct {
+        attrs,
+        vis,
+        struct_token,
+        ident,
+        generics: _,
+        fields: _,
+        semi_token,
+    } = item;
+
+    let debug_name = ident.to_string();
+
+    let semi_token = semi_token.unwrap_or_default();
+    let output = quote! {
+        #(#attrs)*
+        #[derive(Copy, Clone, Eq, PartialEq, Hash)]
+        #vis #struct_token #ident(std::num::NonZeroU32)#semi_token
+
+        impl #ident {
+            const MAX: u32 = u32::MAX - 1;
+
+            #vis const fn from_usize(value: usize) -> Self {
+                assert!(value <= Self::MAX as usize);
+
+                // SAFETY:
+                // * The `value < u32::MAX` guarantees that the add doesn't overflow.
+                // * The `+ 1` guarantees that the index is not zero
+                #[allow(unsafe_code)]
+                Self(unsafe { std::num::NonZeroU32::new_unchecked((value as u32) + 1) })
+            }
+
+            #vis const fn from_u32(value: u32) -> Self {
+                assert!(value <= Self::MAX);
+
+                // SAFETY:
+                // * The `value < u32::MAX` guarantees that the add doesn't overflow.
+                // * The `+ 1` guarantees that the index is larger than zero.
+                #[allow(unsafe_code)]
+                Self(unsafe { std::num::NonZeroU32::new_unchecked(value + 1) })
+            }
+
+            /// Returns the index as a `u32` value
+            #[inline]
+            #vis const fn as_u32(self) -> u32 {
+                self.0.get() - 1
+            }
+
+            /// Returns the index as a `u32` value
+            #[inline]
+            #vis const fn as_usize(self) -> usize {
+                self.as_u32() as usize
+            }
+
+            #[inline]
+            #vis const fn index(self) -> usize {
+                self.as_usize()
+            }
+        }
+
+        impl std::ops::Add<usize> for #ident {
+            type Output = #ident;
+
+            fn add(self, rhs: usize) -> Self::Output {
+                #ident::from_usize(self.index() + rhs)
+            }
+        }
+
+        impl std::ops::Add for #ident {
+            type Output = #ident;
+
+            fn add(self, rhs: Self) -> Self::Output {
+                #ident::from_usize(self.index() + rhs.index())
+            }
+        }
+
+        impl std::fmt::Debug for #ident {
+            fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+                f.debug_tuple(#debug_name).field(&self.index()).finish()
+            }
+        }
+
+        impl ruff_index::Idx for #ident {
+            #[inline]
+            fn new(value: usize) -> Self {
+                #ident::from_usize(value)
+            }
+
+            #[inline]
+            fn index(self) -> usize {
+                self.index()
+            }
+        }
+
+        impl From<usize> for #ident {
+            fn from(value: usize) -> Self {
+                #ident::from_usize(value)
+            }
+        }
+
+        impl From<u32> for #ident {
+            fn from(value: u32) -> Self {
+                #ident::from_u32(value)
+            }
+        }
+
+        impl From<#ident> for usize {
+            fn from(value: #ident) -> Self {
+                value.as_usize()
+            }
+        }
+
+        impl From<#ident> for u32 {
+            fn from(value: #ident) -> Self {
+                value.as_u32()
+            }
+        }
+    };
+
+    Ok(output)
+}