Enable skipping of derived traversals, with reason

eggyal · eggyal · commit 3020cbd21b51 · 2023-11-08T21:21:35.000Z
diff --git a/compiler/rustc_macros/src/lib.rs b/compiler/rustc_macros/src/lib.rs
@@ -74,7 +74,7 @@ decl_derive!([TyEncodable] => serialize::type_encodable_derive);
 decl_derive!([MetadataDecodable] => serialize::meta_decodable_derive);
 decl_derive!([MetadataEncodable] => serialize::meta_encodable_derive);
 decl_derive!(
-    [TypeFoldable] =>
+    [TypeFoldable, attributes(skip_traversal)] =>
     /// Derives `TypeFoldable` for the annotated `struct` or `enum` (`union` is not supported).
     ///
     /// Folds will produce a value of the same struct or enum variant as the input, with trivial
@@ -83,7 +83,18 @@ decl_derive!(
     /// both does not reference any generic type parameters and either
     /// - does not reference any `'tcx` lifetime parameter; or
     /// - does not contain anything that could be of interest to folders.
-    /// 
+    ///
+    /// Non-trivial fields (e.g. of type `T`) can instead be left unchanged by applying
+    /// `#[skip_traversal(because_trivial)]` to the field definition (or even to a variant
+    /// definition if it should apply to all fields therein), but the derived implementation will
+    /// only be applicable if `T` does not contain anything that may be of interest to folders
+    /// (thus preventing fields from being so-skipped erroneously).
+    ///
+    /// In some rare situations, it may be desirable to skip folding of an item or field (or
+    /// variant) that might otherwise be of interest to folders: **this is dangerous and could lead
+    /// to miscompilation if user expectations are not met!**  Nevertheless, such can be achieved
+    /// via a `#[skip_traversal(despite_potential_miscompilation_because = "<reason>"]` attribute.
+    ///
     /// If the annotated type has a `'tcx` lifetime parameter, then that will be used as the
     /// lifetime for the type context/interner; otherwise the lifetime of the type context/interner
     /// will be unrelated to the annotated type. It therefore matters how any lifetime parameters of
@@ -98,7 +109,7 @@ decl_derive!(
     traversable::traversable_derive::<traversable::Foldable>
 );
 decl_derive!(
-    [TypeVisitable] =>
+    [TypeVisitable, attributes(skip_traversal)] =>
     /// Derives `TypeVisitable` for the annotated `struct` or `enum` (`union` is not supported).
     ///
     /// Each non-trivial field of the struct or enum variant will be visited (in definition order)
@@ -108,6 +119,17 @@ decl_derive!(
     /// - does not reference any `'tcx` lifetime parameter; or
     /// - does not contain anything that could be of interest to visitors.
     ///
+    /// Non-trivial fields (e.g. of type `T`) can instead be ignored by applying
+    /// `#[skip_traversal(because_trivial)]` to the field definition (or even to a variant
+    /// definition if it should apply to all fields therein), but the derived implementation will
+    /// only be applicable if `T` does not contain anything that may be of interest to visitors
+    /// (thus preventing fields from being so-skipped erroneously).
+    ///
+    /// In some rare situations, it may be desirable to skip visiting of an item or field (or
+    /// variant) that might otherwise be of interest to visitors: **this is dangerous and could lead
+    /// to miscompilation if user expectations are not met!**  Nevertheless, such can be achieved
+    /// via a `#[skip_traversal(despite_potential_miscompilation_because = "<reason>"]` attribute.
+    ///
     /// If the annotated type has a `'tcx` lifetime parameter, then that will be used as the
     /// lifetime for the type context/interner; otherwise the lifetime of the type context/interner
     /// will be unrelated to the annotated type. It therefore matters how any lifetime parameters of
diff --git a/compiler/rustc_macros/src/traversable.rs b/compiler/rustc_macros/src/traversable.rs
@@ -1,7 +1,10 @@
 use proc_macro2::{Ident, Span, TokenStream};
 use quote::{quote, ToTokens};
-use std::collections::HashSet;
-use syn::{parse_quote, visit, Field, Generics, Lifetime};
+use std::collections::{hash_map::Entry, HashMap};
+use syn::{
+    meta::ParseNestedMeta, parse::Error, parse_quote, visit, Attribute, Field, Generics, Lifetime,
+    LitStr, Token,
+};
 
 #[cfg(test)]
 mod tests;
@@ -31,6 +34,114 @@ enum Type {
 }
 use Type::*;
 
+#[derive(Clone, Copy)]
+enum WhenToSkip {
+    /// No skip_traversal annotation requires the annotated item to be skipped
+    Never,
+
+    /// A skip_traversal annotation requires the annotated item to be skipped, with its type
+    /// constrained to TriviallyTraversable
+    Always(Span),
+
+    /// A `despite_potential_miscompilation_because` annotation is present, thus requiring the
+    /// annotated item to be forcibly skipped without its type being constrained to
+    /// TriviallyTraversable
+    Forced,
+}
+
+impl Default for WhenToSkip {
+    fn default() -> Self {
+        Self::Never
+    }
+}
+
+impl WhenToSkip {
+    fn is_skipped(&self) -> bool {
+        !matches!(self, WhenToSkip::Never)
+    }
+}
+
+impl std::ops::BitOrAssign for WhenToSkip {
+    fn bitor_assign(&mut self, rhs: Self) {
+        match self {
+            Self::Forced => (),
+            Self::Always(_) => {
+                if matches!(rhs, Self::Forced) {
+                    *self = Self::Forced;
+                }
+            }
+            Self::Never => *self = rhs,
+        }
+    }
+}
+
+impl WhenToSkip {
+    fn find<const IS_TYPE: bool>(&mut self, attrs: &[Attribute], ty: Type) -> Result<(), Error> {
+        fn parse_reason(meta: &ParseNestedMeta<'_>) -> Result<(), Error> {
+            if !meta.value()?.parse::<LitStr>()?.value().trim().is_empty() {
+                Ok(())
+            } else {
+                Err(meta.error("skip reason must be a non-empty string"))
+            }
+        }
+
+        let mut found = None;
+        for attr in attrs {
+            if attr.path().is_ident("skip_traversal") {
+                found = Some(attr);
+                attr.parse_nested_meta(|meta| {
+                    if meta.path.is_ident("despite_potential_miscompilation_because") {
+                        parse_reason(&meta)?;
+                        *self |= Self::Forced;
+                        return Ok(());
+                    }
+
+                    if !IS_TYPE && ty == Generic && meta.path.is_ident("because_trivial") {
+                        *self |= Self::Always(meta.error("").span());
+                        return Ok(());
+                    }
+
+                    Err(meta.error("unsupported skip reason"))
+                })?;
+            }
+        }
+
+        if let (Self::Never, Some(attr)) = (self, found) {
+            Err(Error::new_spanned(
+                attr,
+                if IS_TYPE {
+                    match ty {
+                        Trivial => {
+                            "trivially traversable types are always skipped, so this attribute is superfluous"
+                        }
+                        _ => {
+                            "\
+                            Justification must be provided for skipping this potentially interesting type, by specifying\n\
+                            `despite_potential_miscompilation_because = \"<reason>\"`\
+                        "
+                        }
+                    }
+                } else {
+                    match ty {
+                        Trivial => {
+                            "trivially traversable fields are always skipped, so this attribute is superfluous"
+                        }
+                        _ => {
+                            "\
+                            Justification must be provided for skipping potentially interesting fields, by specifying EITHER:\n\
+                            `because_trivial` if concrete instances do not actually contain anything of interest (enforced by the compiler); OR\n\
+                            `despite_potential_miscompilation_because = \"<reason>\"` in the rare case that a field should always be skipped regardless\
+                        "
+                        }
+                    }
+                },
+            ))
+        } else {
+            Ok(())
+        }
+    }
+}
+
 pub struct Interner<'a>(Option<&'a Lifetime>);
 
 impl<'a> Interner<'a> {
@@ -73,7 +184,7 @@ pub trait Traversable {
     /// A `match` arm for `variant`, where `f` generates the tokens for each binding.
     fn arm(
         variant: &synstructure::VariantInfo<'_>,
-        f: impl FnMut(&synstructure::BindingInfo<'_>) -> TokenStream,
+        f: impl FnMut(&synstructure::BindingInfo<'_>) -> Result<TokenStream, Error>,
     ) -> TokenStream;
 
     /// The body of an implementation given the `interner`, `traverser` and match expression `body`.
@@ -100,10 +211,10 @@ impl Traversable for Foldable {
     }
     fn arm(
         variant: &synstructure::VariantInfo<'_>,
-        mut f: impl FnMut(&synstructure::BindingInfo<'_>) -> TokenStream,
+        mut f: impl FnMut(&synstructure::BindingInfo<'_>) -> Result<TokenStream, Error>,
     ) -> TokenStream {
         let bindings = variant.bindings();
-        variant.construct(|_, index| f(&bindings[index]))
+        variant.construct(|_, index| f(&bindings[index]).unwrap_or_else(Error::into_compile_error))
     }
     fn impl_body(
         interner: Interner<'_>,
@@ -137,9 +248,14 @@ impl Traversable for Visitable {
     }
     fn arm(
         variant: &synstructure::VariantInfo<'_>,
-        f: impl FnMut(&synstructure::BindingInfo<'_>) -> TokenStream,
+        f: impl FnMut(&synstructure::BindingInfo<'_>) -> Result<TokenStream, Error>,
     ) -> TokenStream {
-        variant.bindings().iter().map(f).collect()
+        variant
+            .bindings()
+            .iter()
+            .map(f)
+            .collect::<Result<_, _>>()
+            .unwrap_or_else(Error::into_compile_error)
     }
     fn impl_body(
         interner: Interner<'_>,
@@ -199,49 +315,101 @@ impl Interner<'_> {
 
 pub fn traversable_derive<T: Traversable>(
     mut structure: synstructure::Structure<'_>,
-) -> TokenStream {
+) -> Result<TokenStream, Error> {
+    use WhenToSkip::*;
+
     let ast = structure.ast();
 
     let interner = Interner::resolve(&ast.generics);
     let traverser = gen_param("T", &ast.generics);
     let traversable = T::traversable(&interner);
 
+    let skip_traversal =
+        |t: &dyn ToTokens| parse_quote! { #interner: ::rustc_middle::ty::TriviallyTraverses<#t> };
+
     structure.underscore_const(true);
     structure.add_bounds(synstructure::AddBounds::None);
     structure.bind_with(|_| synstructure::BindStyle::Move);
 
-    if interner.0.is_none() {
+    let not_generic = if interner.0.is_none() {
         structure.add_impl_generic(parse_quote! { 'tcx });
-    }
+        Trivial
+    } else {
+        NotGeneric
+    };
 
     // If our derived implementation will be generic over the traversable type, then we must
     // constrain it to only those generic combinations that satisfy the traversable trait's
     // supertraits.
-    if ast.generics.type_params().next().is_some() {
+    let ty = if ast.generics.type_params().next().is_some() {
         let supertraits = T::supertraits(&interner);
         structure.add_where_predicate(parse_quote! { Self: #supertraits });
-    }
+        Generic
+    } else {
+        not_generic
+    };
 
-    // We add predicates to each generic field type, rather than to our generic type parameters.
-    // This results in a "perfect derive", but it can result in trait solver cycles if any type
-    // parameters are involved in recursive type definitions; fortunately that is not the case (yet).
-    let mut predicates = HashSet::new();
-    let arms = structure.each_variant(|variant| {
-        T::arm(variant, |bind| {
-            let ast = bind.ast();
-            let field_ty = interner.type_of(&bind.referenced_ty_params(), [ast]);
-            if field_ty == Generic {
-                predicates.insert(ast.ty.clone());
+    let mut when_to_skip = WhenToSkip::default();
+    when_to_skip.find::<true>(&ast.attrs, ty)?;
+    let body = if when_to_skip.is_skipped() {
+        if let Always(_) = when_to_skip {
+            structure.add_where_predicate(skip_traversal(&<Token![Self]>::default()));
+        }
+        T::traverse(quote! { self }, true, &interner)
+    } else {
+        // We add predicates to each generic field type, rather than to our generic type parameters.
+        // This results in a "perfect derive" that avoids having to propagate `#[skip_traversal]` annotations
+        // into wrapping types, but it can result in trait solver cycles if any type parameters are involved
+        // in recursive type definitions; fortunately that is not the case (yet).
+        let mut predicates = HashMap::<_, (_, _)>::new();
+
+        let arms = structure.each_variant(|variant| {
+            let variant_ty = interner.type_of(&variant.referenced_ty_params(), variant.ast().fields);
+            let mut skipped_variant = WhenToSkip::default();
+            if let Err(err) = skipped_variant.find::<false>(variant.ast().attrs, variant_ty) {
+                return err.into_compile_error();
             }
-            T::traverse(bind.into_token_stream(), field_ty == Trivial, &interner)
-        })
-    });
-    // the order in which `where` predicates appear in rust source is irrelevant
-    #[allow(rustc::potential_query_instability)]
-    for ty in predicates {
-        structure.add_where_predicate(parse_quote! { #ty: #traversable });
-    }
-    let body = quote! { match self { #arms } };
+            T::arm(variant, |bind| {
+                let ast = bind.ast();
+                let is_skipped = variant_ty == Trivial || {
+                    let field_ty = interner.type_of(&bind.referenced_ty_params(), [ast]);
+                    field_ty == Trivial || {
+                        let mut skipped_field = skipped_variant;
+                        skipped_field.find::<false>(&ast.attrs, field_ty)?;
+
+                        match predicates.entry(ast.ty.clone()) {
+                            Entry::Occupied(existing) => match (&mut existing.into_mut().0, skipped_field) {
+                                (Never, Never) | (Never, Forced) | (Forced, Forced) | (Always(_), Always(_)) => (),
+                                (existing @ Forced, Never) => *existing = Never,
+                                (&mut Always(span), _) | (_, Always(span)) => return Err(Error::new(span, format!("\
+                                    This annotation only makes sense if all fields of type `{0}` are annotated identically.\n\
+                                    In particular, the derived impl will only be applicable when `{0}: TriviallyTraversable` and therefore all traversals of `{0}` will be no-ops;\n\
+                                    accordingly it makes no sense for other fields of type `{0}` to omit `#[skip_traversal]` or to include `despite_potential_miscompilation_because`.\
+                                ", ast.ty.to_token_stream()))),
+                            },
+                            Entry::Vacant(entry) => { entry.insert((skipped_field, field_ty)); }
+                        }
+
+                        skipped_field.is_skipped()
+                    }
+                };
+
+                Ok(T::traverse(bind.into_token_stream(), is_skipped, &interner))
+            })
+        });
+
+        // the order in which `where` predicates appear in rust source is irrelevant
+        #[allow(rustc::potential_query_instability)]
+        for (ty, (when_to_skip, field_ty)) in predicates {
+            structure.add_where_predicate(match when_to_skip {
+                Always(_) => skip_traversal(&ty),
+                // we only need to add traversable predicate for generic types
+                Never if field_ty == Generic => parse_quote! { #ty: #traversable },
+                _ => continue,
+            });
+        }
+        quote! { match self { #arms } }
+    };
 
-    structure.bound_impl(traversable, T::impl_body(interner, traverser, body))
+    Ok(structure.bound_impl(traversable, T::impl_body(interner, traverser, body)))
 }
diff --git a/compiler/rustc_macros/src/traversable/tests.rs b/compiler/rustc_macros/src/traversable/tests.rs
