1 files changed, 65 insertions, 14 deletions

diff --git a/src/tools/rustfmt/src/skip.rs b/src/tools/rustfmt/src/skip.rs
index 032922d42..68f85b2ad 100644
--- a/src/tools/rustfmt/src/skip.rs
+++ b/src/tools/rustfmt/src/skip.rs
@@ -2,33 +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 {
-    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),
+        }
+    }
+}
+
+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.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;
     }
 }