feat: nicer skip context for macro/attribute

Author: tommilligan

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,60 @@
 
 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 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.append(get_skip_names("macros", attrs));
+        self.attributes.append(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(Default, Clone)]
+pub(crate) struct SkipNameContext {
+    all: bool,
+    values: HashSet<String>,
+}
+
+impl SkipNameContext {
+    pub(crate) fn append(&mut self, values: Vec<String>) {
+        self.values.extend(values);
     }
 
-    pub(crate) fn update_macros<T>(&mut self, other: T)
-    where
-        T: IntoIterator<Item = String>,
-    {
-        self.macros.extend(other.into_iter());
+    pub(crate) fn update(&mut self, other: Self) {
+        self.all = self.all || other.all;
+        self.values.extend(other.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 {
+        self.all || self.values.contains(name)
     }
 
-    pub(crate) fn skip_attribute(&self, name: &str) -> bool {
-        self.attributes.iter().any(|n| n == name)
+    pub(crate) fn set_all(&mut self, all: bool) {
+        self.all = 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.set_all(true),
             }
         }
-        skip_context.update_macros(macro_names);
+        skip_context.macros.append(macro_names);
         FmtVisitor {
             parent_context: None,
             parse_sess: parse_session,