Try to discover mod items in cfg_match! arms (#6522)

Author: CAD97

Diff for: 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,
+    }
+}

Diff for: 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 {

Diff for: src/parse/macros/cfg_match.rs

@@ -0,0 +1,80 @@
+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 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)
+}

Diff for: 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> {

Diff for: src/test/mod.rs

@@ -42,6 +42,8 @@ const FILE_SKIP_LIST: &[&str] = &[
     "issue-3253/foo.rs",
     "issue-3253/bar.rs",
     "issue-3253/paths",
+    // This directory is directly tested by format_files_find_new_files_via_cfg_match
+    "cfg_match",
     // These files and directory are a part of modules defined inside `cfg_attr(..)`.
     "cfg_mod/dir",
     "cfg_mod/bar.rs",
@@ -468,6 +470,45 @@ fn format_files_find_new_files_via_cfg_if() {
     });
 }
 
+#[test]
+fn format_files_find_new_files_via_cfg_match() {
+    init_log();
+    run_test_with(&TestSetting::default(), || {
+        // We load these two files into the same session to test cfg_match!
+        // transparent mod discovery, and to ensure that it does not suffer
+        // from a similar issue as cfg_if! support did with issue-4656.
+        let files = vec![
+            Path::new("tests/source/cfg_match/lib2.rs"),
+            Path::new("tests/source/cfg_match/lib.rs"),
+        ];
+
+        let config = Config::default();
+        let mut session = Session::<io::Stdout>::new(config, None);
+
+        let mut write_result = HashMap::new();
+        for file in files {
+            assert!(file.exists());
+            let result = session.format(Input::File(file.into())).unwrap();
+            assert!(!session.has_formatting_errors());
+            assert!(!result.has_warnings());
+            let mut source_file = SourceFile::new();
+            mem::swap(&mut session.source_file, &mut source_file);
+
+            for (filename, text) in source_file {
+                if let FileName::Real(ref filename) = filename {
+                    write_result.insert(filename.to_owned(), text);
+                }
+            }
+        }
+        assert_eq!(
+            6,
+            write_result.len(),
+            "Should have uncovered an extra file (format_me_please_x.rs) via lib.rs"
+        );
+        assert!(handle_result(write_result, None).is_ok());
+    });
+}
+
 #[test]
 fn stdin_formatting_smoke_test() {
     init_log();

Diff for: tests/source/cfg_match/format_me_please_1.rs

@@ -0,0 +1,2 @@
+
+pub fn     format_me_please_1(   )  {     }

Diff for: tests/source/cfg_match/format_me_please_2.rs

@@ -0,0 +1,2 @@
+
+pub fn     format_me_please_2(   )  {     }

Diff for: tests/source/cfg_match/format_me_please_3.rs

@@ -0,0 +1,2 @@
+
+pub fn     format_me_please_3(   )  {     }

Diff for: tests/source/cfg_match/format_me_please_4.rs

@@ -0,0 +1,2 @@
+
+pub fn     format_me_please_4(   )  {     }

Diff for: tests/source/cfg_match/lib.rs

@@ -0,0 +1,16 @@
+#![feature(cfg_match)]
+
+std::cfg_match! {
+    test => {
+        mod format_me_please_1;
+    }
+    target_family = "unix" => {
+        mod format_me_please_2;
+    }
+    cfg(target_pointer_width = "32") => {
+        mod format_me_please_3;
+    }
+    _ => {
+        mod format_me_please_4;
+    }
+}

Diff for: tests/source/cfg_match/lib2.rs

@@ -0,0 +1,3 @@
+its_a_macro! {
+    // Contents
+}

Diff for: tests/target/cfg_match/format_me_please_1.rs

@@ -0,0 +1 @@
+pub fn format_me_please_1() {}

Diff for: tests/target/cfg_match/format_me_please_2.rs

@@ -0,0 +1 @@
+pub fn format_me_please_2() {}

Diff for: tests/target/cfg_match/format_me_please_3.rs

@@ -0,0 +1 @@
+pub fn format_me_please_3() {}

Diff for: tests/target/cfg_match/format_me_please_4.rs

@@ -0,0 +1 @@
+pub fn format_me_please_4() {}

Diff for: tests/target/cfg_match/lib.rs

@@ -0,0 +1,16 @@
+#![feature(cfg_match)]
+
+std::cfg_match! {
+    test => {
+        mod format_me_please_1;
+    }
+    target_family = "unix" => {
+        mod format_me_please_2;
+    }
+    cfg(target_pointer_width = "32") => {
+        mod format_me_please_3;
+    }
+    _ => {
+        mod format_me_please_4;
+    }
+}

Diff for: tests/target/cfg_match/lib2.rs

@@ -0,0 +1,3 @@
+its_a_macro! {
+    // Contents
+}