Extend cfg_if! support to cfg_match!

Author: CAD97

src/modules.rs

@@ -167,6 +167,25 @@ impl<'ast, 'psess, 'c> ModResolver<'ast, 'psess> {
         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);
+        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(())
+    }
+
     /// Visit modules defined inside macro calls.
     fn visit_mod_outside_ast(
         &mut self,
@@ -178,6 +197,11 @@ impl<'ast, 'psess, 'c> ModResolver<'ast, 'psess> {
                 continue;
             }
 
+            if is_cfg_match(&item) {
+                self.visit_cfg_match(Cow::Owned(item.into_inner()))?;
+                continue;
+            }
+
             if let ast::ItemKind::Mod(_, ref sub_mod_kind) = item.kind {
                 let span = item.span;
                 self.visit_sub_mod(
@@ -204,6 +228,10 @@ impl<'ast, 'psess, 'c> ModResolver<'ast, 'psess> {
                 self.visit_cfg_if(Cow::Borrowed(item))?;
             }
 
+            if is_cfg_match(item) {
+                self.visit_cfg_match(Cow::Borrowed(item))?;
+            }
+
             if let ast::ItemKind::Mod(_, ref sub_mod_kind) = item.kind {
                 let span = item.span;
                 self.visit_sub_mod(
@@ -575,3 +603,17 @@ fn is_cfg_if(item: &ast::Item) -> bool {
         _ => false,
     }
 }
+
+fn is_cfg_match(item: &ast::Item) -> bool {
+    match item.kind {
+        ast::ItemKind::MacCall(ref mac) => {
+            if let Some(last_segment) = mac.path.segments.last() {
+                if last_segment.ident.name == Symbol::intern("cfg_match") {
+                    return true;
+                }
+            }
+            false
+        }
+        _ => false,
+    }
+}

src/modules/visitor.rs

@@ -5,6 +5,7 @@ use tracing::debug;
 
 use crate::attr::MetaVisitor;
 use crate::parse::macros::cfg_if::parse_cfg_if;
+use crate::parse::macros::cfg_match::parse_cfg_match;
 use crate::parse::session::ParseSess;
 
 pub(crate) struct ModItem {
@@ -71,6 +72,65 @@ impl<'a, 'ast: 'a> CfgIfVisitor<'a> {
     }
 }
 
+/// 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> {
+        // Support both:
+        // ```
+        // std::cfg_match! {..}
+        // core::cfg_match! {..}
+        // ```
+        // And:
+        // ```
+        // use std::cfg_match;
+        // cfg_match! {..}
+        // ```
+        match mac.path.segments.last() {
+            Some(last_segment) => {
+                if last_segment.ident.name != Symbol::intern("cfg_match") {
+                    return Err("Expected cfg_match");
+                }
+            }
+            None => {
+                return Err("Expected cfg_match");
+            }
+        };
+
+        let items = parse_cfg_match(self.psess, mac)?;
+        self.mods
+            .append(&mut items.into_iter().map(|item| ModItem { item }).collect());
+
+        Ok(())
+    }
+}
+
 /// Extracts `path = "foo.rs"` from attributes.
 #[derive(Default)]
 pub(crate) struct PathVisitor {

src/parse/macros/cfg_match.rs

@@ -0,0 +1,81 @@
+use std::panic::{AssertUnwindSafe, catch_unwind};
+
+use rustc_ast::ast;
+use rustc_ast::token::{Delimiter, TokenKind};
+use rustc_parse::exp;
+use rustc_parse::parser::ForceCollect;
+use rustc_span::symbol::kw;
+
+use crate::parse::macros::build_stream_parser;
+use crate::parse::session::ParseSess;
+
+pub(crate) fn parse_cfg_match<'a>(
+    psess: &'a ParseSess,
+    mac: &'a ast::MacCall,
+) -> Result<Vec<ast::Item>, &'static str> {
+    match catch_unwind(AssertUnwindSafe(|| parse_cfg_match_inner(psess, mac))) {
+        Ok(Ok(items)) => Ok(items),
+        Ok(err @ Err(_)) => err,
+        Err(..) => Err("failed to parse cfg_match!"),
+    }
+}
+
+fn parse_cfg_match_inner<'a>(
+    psess: &'a ParseSess,
+    mac: &'a ast::MacCall,
+) -> Result<Vec<ast::Item>, &'static str> {
+    let ts = mac.args.tokens.clone();
+    let mut parser = build_stream_parser(psess.inner(), ts);
+
+    if parser.token == TokenKind::OpenDelim(Delimiter::Brace) {
+        return Err("Expression position cfg_match! not yet supported");
+    }
+
+    let mut items = vec![];
+
+    while parser.token.kind != TokenKind::Eof {
+        if !parser.eat_keyword(exp!(Underscore)) {
+            parser.parse_attr_item(ForceCollect::No).map_err(|e| {
+                e.cancel();
+                "Failed to parse attr item"
+            })?;
+        }
+
+        if !parser.eat(exp!(FatArrow)) {
+            return Err("Expected a fat arrow");
+        }
+
+        if !parser.eat(exp!(OpenBrace)) {
+            return Err("Expected an opening brace");
+        }
+
+        while parser.token != TokenKind::CloseDelim(Delimiter::Brace)
+            && parser.token.kind != TokenKind::Eof
+        {
+            let item = match parser.parse_item(ForceCollect::No) {
+                Ok(Some(item_ptr)) => item_ptr.into_inner(),
+                Ok(None) => continue,
+                Err(err) => {
+                    err.cancel();
+                    parser.psess.dcx().reset_err_count();
+                    return Err(
+                        "Expected item inside cfg_match block, but failed to parse it as an item",
+                    );
+                }
+            };
+            if let ast::ItemKind::Mod(..) = item.kind {
+                items.push(item);
+            }
+        }
+
+        if !parser.eat(exp!(CloseBrace)) {
+            return Err("Expected a closing brace");
+        }
+
+        if parser.eat(exp!(Eof)) {
+            break;
+        }
+    }
+
+    Ok(items)
+}

src/parse/macros/mod.rs

@@ -11,6 +11,7 @@ use crate::rewrite::RewriteContext;
 
 pub(crate) mod asm;
 pub(crate) mod cfg_if;
+pub(crate) mod cfg_match;
 pub(crate) mod lazy_static;
 
 fn build_stream_parser<'a>(psess: &'a ParseSess, tokens: TokenStream) -> Parser<'a> {