Clean up macro visiting

Author: CAD97

src/modules.rs

@@ -147,29 +147,10 @@ impl<'ast, 'psess, 'c> ModResolver<'ast, 'psess> {
         Ok(self.file_map)
     }
 
-    /// Visit `cfg_if` macro and look for module declarations.
-    fn visit_cfg_if(&mut self, item: Cow<'ast, ast::Item>) -> Result<(), ModuleResolutionError> {
-        let mut visitor = visitor::CfgIfVisitor::new(self.psess);
-        visitor.visit_item(&item);
-        for module_item in visitor.mods() {
-            if let ast::ItemKind::Mod(_, ref sub_mod_kind) = module_item.item.kind {
-                self.visit_sub_mod(
-                    &module_item.item,
-                    Module::new(
-                        module_item.item.span,
-                        Some(Cow::Owned(sub_mod_kind.clone())),
-                        Cow::Owned(ThinVec::new()),
-                        Cow::Owned(ast::AttrVec::new()),
-                    ),
-                )?;
-            }
-        }
-        Ok(())
-    }
-
-    fn visit_cfg_match(&mut self, item: Cow<'ast, ast::Item>) -> Result<(), ModuleResolutionError> {
-        let mut visitor = visitor::CfgMatchVisitor::new(self.psess);
-        visitor.visit_item(&item);
+    /// Visit calls of well known macros and look for module declarations.
+    fn visit_mac_call(&mut self, mac: &ast::MacCall) -> Result<(), ModuleResolutionError> {
+        let mut visitor = visitor::KnownMacroVisitor::new(self.psess);
+        visitor.visit_mac_call(mac);
         for module_item in visitor.mods() {
             if let ast::ItemKind::Mod(_, ref sub_mod_kind) = module_item.item.kind {
                 self.visit_sub_mod(
@@ -192,14 +173,8 @@ impl<'ast, 'psess, 'c> ModResolver<'ast, 'psess> {
         items: ThinVec<rustc_ast::ptr::P<ast::Item>>,
     ) -> Result<(), ModuleResolutionError> {
         for item in items {
-            if is_cfg_if(&item) {
-                self.visit_cfg_if(Cow::Owned(item.into_inner()))?;
-                continue;
-            }
-
-            if is_cfg_match(&item) {
-                self.visit_cfg_match(Cow::Owned(item.into_inner()))?;
-                continue;
+            if let ast::ItemKind::MacCall(ref mac) = item.kind {
+                self.visit_mac_call(mac)?;
             }
 
             if let ast::ItemKind::Mod(_, ref sub_mod_kind) = item.kind {
@@ -224,12 +199,8 @@ impl<'ast, 'psess, 'c> ModResolver<'ast, 'psess> {
         items: &'ast [rustc_ast::ptr::P<ast::Item>],
     ) -> Result<(), ModuleResolutionError> {
         for item in items {
-            if is_cfg_if(item) {
-                self.visit_cfg_if(Cow::Borrowed(item))?;
-            }
-
-            if is_cfg_match(item) {
-                self.visit_cfg_match(Cow::Borrowed(item))?;
+            if let ast::ItemKind::MacCall(ref mac) = item.kind {
+                self.visit_mac_call(mac)?;
             }
 
             if let ast::ItemKind::Mod(_, ref sub_mod_kind) = item.kind {
@@ -589,24 +560,3 @@ fn path_value(attr: &ast::Attribute) -> Option<Symbol> {
 fn find_path_value(attrs: &[ast::Attribute]) -> Option<Symbol> {
     attrs.iter().flat_map(path_value).next()
 }
-
-fn is_macro_name(mac: &ast::MacCall, name: &str) -> bool {
-    mac.path
-        .segments
-        .last()
-        .map_or(false, |segment| segment.ident.name == Symbol::intern(name))
-}
-
-fn is_cfg_if(item: &ast::Item) -> bool {
-    match item.kind {
-        ast::ItemKind::MacCall(ref mac) => is_macro_name(mac, "cfg_if"),
-        _ => false,
-    }
-}
-
-fn is_cfg_match(item: &ast::Item) -> bool {
-    match item.kind {
-        ast::ItemKind::MacCall(ref mac) => is_macro_name(mac, "cfg_match"),
-        _ => false,
-    }
-}

src/modules/visitor.rs

@@ -1,9 +1,9 @@
 use rustc_ast::ast;
 use rustc_ast::visit::Visitor;
+use rustc_span::Symbol;
 use rustc_span::sym;
 use tracing::debug;
 
-use super::is_macro_name;
 use crate::attr::MetaVisitor;
 use crate::parse::macros::cfg_if::parse_cfg_if;
 use crate::parse::macros::cfg_match::parse_cfg_match;
@@ -13,15 +13,15 @@ pub(crate) struct ModItem {
     pub(crate) item: ast::Item,
 }
 
-/// Traverse `cfg_if!` macro and fetch modules.
-pub(crate) struct CfgIfVisitor<'a> {
+/// Traverse well-known macro names and fetch modules.
+pub(crate) struct KnownMacroVisitor<'a> {
     psess: &'a ParseSess,
     mods: Vec<ModItem>,
 }
 
-impl<'a> CfgIfVisitor<'a> {
-    pub(crate) fn new(psess: &'a ParseSess) -> CfgIfVisitor<'a> {
-        CfgIfVisitor {
+impl<'a> KnownMacroVisitor<'a> {
+    pub(crate) fn new(psess: &'a ParseSess) -> KnownMacroVisitor<'a> {
+        KnownMacroVisitor {
             mods: vec![],
             psess,
         }
@@ -32,7 +32,7 @@ impl<'a> CfgIfVisitor<'a> {
     }
 }
 
-impl<'a, 'ast: 'a> Visitor<'ast> for CfgIfVisitor<'a> {
+impl<'a, 'ast: 'a> Visitor<'ast> for KnownMacroVisitor<'a> {
     fn visit_mac_call(&mut self, mac: &'ast ast::MacCall) {
         match self.visit_mac_inner(mac) {
             Ok(()) => (),
@@ -41,55 +41,24 @@ impl<'a, 'ast: 'a> Visitor<'ast> for CfgIfVisitor<'a> {
     }
 }
 
-impl<'a, 'ast: 'a> CfgIfVisitor<'a> {
+impl<'a, 'ast: 'a> KnownMacroVisitor<'a> {
     fn visit_mac_inner(&mut self, mac: &'ast ast::MacCall) -> Result<(), &'static str> {
-        if !is_macro_name(mac, "cfg_if") {
-            return Err("Expected cfg_if");
+        fn is_macro_name(mac: &ast::MacCall, name: &str) -> bool {
+            mac.path
+                .segments
+                .last()
+                .map_or(false, |segment| segment.ident.name == Symbol::intern(name))
         }
 
-        let items = parse_cfg_if(self.psess, mac)?;
-        self.mods
-            .append(&mut items.into_iter().map(|item| ModItem { item }).collect());
-
-        Ok(())
-    }
-}
-
-/// Traverse `cfg_match!` macro and fetch modules.
-pub(crate) struct CfgMatchVisitor<'a> {
-    psess: &'a ParseSess,
-    mods: Vec<ModItem>,
-}
-
-impl<'a> CfgMatchVisitor<'a> {
-    pub(crate) fn new(psess: &'a ParseSess) -> CfgMatchVisitor<'a> {
-        CfgMatchVisitor {
-            mods: vec![],
-            psess,
-        }
-    }
-
-    pub(crate) fn mods(self) -> Vec<ModItem> {
-        self.mods
-    }
-}
-
-impl<'a, 'ast: 'a> Visitor<'ast> for CfgMatchVisitor<'a> {
-    fn visit_mac_call(&mut self, mac: &'ast ast::MacCall) {
-        match self.visit_mac_inner(mac) {
-            Ok(()) => (),
-            Err(e) => debug!("{}", e),
-        }
-    }
-}
-
-impl<'a, 'ast: 'a> CfgMatchVisitor<'a> {
-    fn visit_mac_inner(&mut self, mac: &'ast ast::MacCall) -> Result<(), &'static str> {
-        if !is_macro_name(mac, "cfg_match") {
-            return Err("Expected cfg_match");
-        }
+        let items = if is_macro_name(mac, "cfg_if") {
+            parse_cfg_if(&self.psess, mac)?
+        } else if is_macro_name(mac, "cfg_match") {
+            parse_cfg_match(&self.psess, mac)?
+        } else {
+            // ignore all other macros
+            vec![]
+        };
 
-        let items = parse_cfg_match(self.psess, mac)?;
         self.mods
             .append(&mut items.into_iter().map(|item| ModItem { item }).collect());