feat: nicer skip context for macro/attribute

src/attr.rs

@@ -337,7 +337,7 @@ impl Rewrite for ast::Attribute {
         } else {
             let should_skip = self
                 .ident()
-                .map(|s| context.skip_context.skip_attribute(s.name.as_str()))
+                .map(|s| context.skip_context.attributes.skip(s.name.as_str()))
                 .unwrap_or(false);
             let prefix = attr_prefix(self);
 
@@ -391,7 +391,7 @@ impl Rewrite for [ast::Attribute] {
 
         // Determine if the source text is annotated with `#[rustfmt::skip::attributes(derive)]`
         // or `#![rustfmt::skip::attributes(derive)]`
-        let skip_derives = context.skip_context.skip_attribute("derive");
+        let skip_derives = context.skip_context.attributes.skip("derive");
 
         // This is not just a simple map because we need to handle doc comments
         // (where we take as many doc comment attributes as possible) and possibly

src/macros.rs

@@ -157,7 +157,8 @@ pub(crate) fn rewrite_macro(
 ) -> Option<String> {
     let should_skip = context
         .skip_context
-        .skip_macro(context.snippet(mac.path.span));
+        .macros
+        .skip(context.snippet(mac.path.span));
     if should_skip {
         None
     } else {

src/skip.rs

@@ -2,41 +2,84 @@
 
 use rustc_ast::ast;
 use rustc_ast_pretty::pprust;
+use std::collections::HashSet;
 
-/// Take care of skip name stack. You can update it by attributes slice or
-/// by other context. Query this context to know if you need skip a block.
+/// Track which blocks of code are to be skipped when formatting.
+///
+/// You can update it by:
+///
+/// - attributes slice
+/// - manually feeding values into the underlying contexts
+///
+/// Query this context to know if you need to skip a block.
 #[derive(Default, Clone)]
 pub(crate) struct SkipContext {
-    pub(crate) all_macros: bool,
-    macros: Vec<String>,
-    attributes: Vec<String>,
+    pub(crate) macros: SkipNameContext,
+    pub(crate) attributes: SkipNameContext,
 }
 
 impl SkipContext {
     pub(crate) fn update_with_attrs(&mut self, attrs: &[ast::Attribute]) {
-        self.macros.append(&mut get_skip_names("macros", attrs));
-        self.attributes
-            .append(&mut get_skip_names("attributes", attrs));
+        self.macros.extend(get_skip_names("macros", attrs));
+        self.attributes.extend(get_skip_names("attributes", attrs));
     }
 
-    pub(crate) fn update(&mut self, mut other: SkipContext) {
-        self.macros.append(&mut other.macros);
-        self.attributes.append(&mut other.attributes);
+    pub(crate) fn update(&mut self, other: SkipContext) {
+        let SkipContext { macros, attributes } = other;
+        self.macros.update(macros);
+        self.attributes.update(attributes);
     }
+}
+
+/// Track which names to skip.
+///
+/// Query this context with a string to know whether to skip it.
+#[derive(Clone)]
+pub(crate) enum SkipNameContext {
+    All,
+    Values(HashSet<String>),
+}
+
+impl Default for SkipNameContext {
+    fn default() -> Self {
+        Self::Values(Default::default())
+    }
+}
+
+impl Extend<String> for SkipNameContext {
+    fn extend<T: IntoIterator<Item = String>>(&mut self, iter: T) {
+        match self {
+            Self::All => {}
+            Self::Values(values) => values.extend(iter),
+        }
+    }
+}
 
-    pub(crate) fn update_macros<T>(&mut self, other: T)
-    where
-        T: IntoIterator<Item = String>,
-    {
-        self.macros.extend(other.into_iter());
+impl SkipNameContext {
+    pub(crate) fn update(&mut self, other: Self) {
+        match (self, other) {
+            // If we're already skipping everything, nothing more can be added
+            (Self::All, _) => {}
+            // If we want to skip all, set it
+            (this, Self::All) => {
+                *this = Self::All;
+            }
+            // If we have some new values to skip, add them
+            (Self::Values(existing_values), Self::Values(new_values)) => {
+                existing_values.extend(new_values)
+            }
+        }
     }
 
-    pub(crate) fn skip_macro(&self, name: &str) -> bool {
-        self.all_macros || self.macros.iter().any(|n| n == name)
+    pub(crate) fn skip(&self, name: &str) -> bool {
+        match self {
+            Self::All => true,
+            Self::Values(values) => values.contains(name),
+        }
     }
 
-    pub(crate) fn skip_attribute(&self, name: &str) -> bool {
-        self.attributes.iter().any(|n| n == name)
+    pub(crate) fn skip_all(&mut self) {
+        *self = Self::All;
     }
 }
 

src/visitor.rs

@@ -775,10 +775,10 @@ impl<'b, 'a: 'b> FmtVisitor<'a> {
         for macro_selector in config.skip_macro_invocations().0 {
             match macro_selector {
                 MacroSelector::Name(name) => macro_names.push(name.to_string()),
-                MacroSelector::All => skip_context.all_macros = true,
+                MacroSelector::All => skip_context.macros.skip_all(),
             }
         }
-        skip_context.update_macros(macro_names);
+        skip_context.macros.extend(macro_names);
         FmtVisitor {
             parent_context: None,
             parse_sess: parse_session,