Format modules defined inside cfg_if macro calls

src/formatting.rs

@@ -99,8 +99,8 @@ fn format_project<T: FormatHandler>(
 
     let mut context = FormatContext::new(&krate, report, parse_session, config, handler);
     let files = modules::ModResolver::new(
-        context.parse_session.source_map(),
-        directory_ownership.unwrap_or(parse::DirectoryOwnership::UnownedViaMod(false)),
+        &context.parse_session,
+        directory_ownership.unwrap_or(parse::DirectoryOwnership::UnownedViaMod(true)),
         !(input_is_stdin || config.skip_children()),
     )
     .visit_crate(&krate)
@@ -112,7 +112,7 @@ fn format_project<T: FormatHandler>(
         }
         should_emit_verbose(input_is_stdin, config, || println!("Formatting {}", path));
         let is_root = path == main_file;
-        context.format_file(path, module, is_root)?;
+        context.format_file(path, &module, is_root)?;
     }
     timer = timer.done_formatting();
 

src/modules.rs

@@ -1,23 +1,27 @@
+use std::borrow::Cow;
 use std::collections::BTreeMap;
 use std::path::{Path, PathBuf};
 
 use syntax::ast;
-use syntax::parse::{parser, DirectoryOwnership};
+use syntax::parse::{parser, DirectoryOwnership, ParseSess};
 use syntax::source_map;
 use syntax::symbol::sym;
+use syntax::visit::Visitor;
 use syntax_pos::symbol::Symbol;
 
 use crate::config::FileName;
 use crate::items::is_mod_decl;
 use crate::utils::contains_skip;
 
-type FileModMap<'a> = BTreeMap<FileName, (&'a ast::Mod, String)>;
+mod visitor;
+
+type FileModMap<'ast> = BTreeMap<FileName, (Cow<'ast, ast::Mod>, String)>;
 
 /// Maps each module to the corresponding file.
-pub(crate) struct ModResolver<'a, 'b> {
-    source_map: &'b source_map::SourceMap,
+pub(crate) struct ModResolver<'ast, 'sess> {
+    parse_sess: &'sess ParseSess,
     directory: Directory,
-    file_map: FileModMap<'a>,
+    file_map: FileModMap<'ast>,
     recursive: bool,
 }
 
@@ -27,10 +31,28 @@ struct Directory {
     ownership: DirectoryOwnership,
 }
 
-impl<'a, 'b> ModResolver<'a, 'b> {
+impl<'a> Directory {
+    fn to_syntax_directory(&'a self) -> syntax::parse::Directory<'a> {
+        syntax::parse::Directory {
+            path: Cow::Borrowed(&self.path),
+            ownership: self.ownership.clone(),
+        }
+    }
+}
+
+enum SubModKind {
+    /// `mod foo;`
+    External(PathBuf, DirectoryOwnership),
+    /// `#[path = "..."] mod foo {}`
+    InternalWithPath(PathBuf),
+    /// `mod foo {}`
+    Internal,
+}
+
+impl<'ast, 'sess, 'c> ModResolver<'ast, 'sess> {
     /// Creates a new `ModResolver`.
     pub(crate) fn new(
-        source_map: &'b source_map::SourceMap,
+        parse_sess: &'sess ParseSess,
         directory_ownership: DirectoryOwnership,
         recursive: bool,
     ) -> Self {
@@ -40,14 +62,17 @@ impl<'a, 'b> ModResolver<'a, 'b> {
                 ownership: directory_ownership,
             },
             file_map: BTreeMap::new(),
-            source_map,
+            parse_sess,
             recursive,
         }
     }
 
     /// Creates a map that maps a file name to the module in AST.
-    pub(crate) fn visit_crate(mut self, krate: &'a ast::Crate) -> Result<FileModMap<'a>, String> {
-        let root_filename = self.source_map.span_to_filename(krate.span);
+    pub(crate) fn visit_crate(
+        mut self,
+        krate: &'ast ast::Crate,
+    ) -> Result<FileModMap<'ast>, String> {
+        let root_filename = self.parse_sess.source_map().span_to_filename(krate.span);
         self.directory.path = match root_filename {
             source_map::FileName::Real(ref path) => path
                 .parent()
@@ -58,54 +83,125 @@ impl<'a, 'b> ModResolver<'a, 'b> {
 
         // Skip visiting sub modules when the input is from stdin.
         if self.recursive {
-            self.visit_mod(&krate.module)?;
+            self.visit_mod_from_ast(&krate.module)?;
         }
 
-        self.file_map
-            .insert(root_filename.into(), (&krate.module, String::new()));
+        self.file_map.insert(
+            root_filename.into(),
+            (Cow::Borrowed(&krate.module), String::new()),
+        );
         Ok(self.file_map)
     }
 
-    fn visit_mod(&mut self, module: &'a ast::Mod) -> Result<(), String> {
+    /// Visit macro calls and look for module declarations. Currently only supports `cfg_if` macro.
+    fn visit_mac(&mut self, item: Cow<'ast, ast::Item>) -> Result<(), String> {
+        let mut visitor =
+            visitor::CfgIfVisitor::new(self.parse_sess, self.directory.to_syntax_directory());
+        visitor.visit_item(&item);
+        for module_item in visitor.mods() {
+            if let ast::ItemKind::Mod(ref sub_mod) = module_item.item.node {
+                self.visit_mod_from_mac_inner(&item, Cow::Owned(sub_mod.clone()))?;
+            }
+        }
+        Ok(())
+    }
+
+    /// Visit modules defined inside macro calls.
+    fn visit_mod_from_macro(&mut self, module: Cow<'ast, ast::Mod>) -> Result<(), String> {
         for item in &module.items {
+            if let ast::ItemKind::Mac(..) = item.node {
+                self.visit_mac(Cow::Owned(item.clone().into_inner()))?;
+            }
+
             if let ast::ItemKind::Mod(ref sub_mod) = item.node {
-                if contains_skip(&item.attrs) {
-                    continue;
-                }
+                self.visit_mod_from_mac_inner(item, Cow::Owned(sub_mod.clone()))?;
+            }
+        }
+        Ok(())
+    }
 
-                let old_direcotry = self.directory.clone();
-                if is_mod_decl(item) {
-                    // mod foo;
-                    // Look for an extern file.
-                    let (mod_path, directory_ownership) =
-                        self.find_external_module(item.ident, &item.attrs)?;
-                    self.file_map.insert(
-                        FileName::Real(mod_path.clone()),
-                        (sub_mod, item.ident.as_str().get().to_owned()),
-                    );
-                    self.directory = Directory {
-                        path: mod_path.parent().unwrap().to_path_buf(),
-                        ownership: directory_ownership,
-                    }
-                } else {
-                    // An internal module (`mod foo { /* ... */ }`);
-                    if let Some(path) = find_path_value(&item.attrs) {
-                        // All `#[path]` files are treated as though they are a `mod.rs` file.
-                        self.directory = Directory {
-                            path: Path::new(&path.as_str()).to_path_buf(),
-                            ownership: DirectoryOwnership::Owned { relative: None },
-                        };
-                    } else {
-                        self.push_inline_mod_directory(item.ident, &item.attrs);
-                    }
-                }
-                self.visit_mod(sub_mod)?;
-                self.directory = old_direcotry;
+    fn visit_mod_from_mac_inner(
+        &mut self,
+        item: &'c ast::Item,
+        sub_mod: Cow<'ast, ast::Mod>,
+    ) -> Result<(), String> {
+        let old_directory = self.directory.clone();
+        self.visit_sub_mod(item, &sub_mod)?;
+        self.visit_mod_from_macro(sub_mod)?;
+        self.directory = old_directory;
+        Ok(())
+    }
+
+    /// Visit modules from AST.
+    fn visit_mod_from_ast(&mut self, module: &'ast ast::Mod) -> Result<(), String> {
+        for item in &module.items {
+            if let ast::ItemKind::Mac(..) = item.node {
+                self.visit_mac(Cow::Borrowed(item))?;
+            }
+
+            if let ast::ItemKind::Mod(ref sub_mod) = item.node {
+                let old_directory = self.directory.clone();
+                self.visit_sub_mod(item, &Cow::Borrowed(sub_mod))?;
+                self.visit_mod_from_ast(sub_mod)?;
+                self.directory = old_directory;
+            }
+        }
+        Ok(())
+    }
+
+    fn visit_sub_mod(
+        &mut self,
+        item: &'c ast::Item,
+        sub_mod: &Cow<'ast, ast::Mod>,
+    ) -> Result<(), String> {
+        match self.peek_sub_mod(item)? {
+            Some(SubModKind::External(mod_path, directory_ownership)) => {
+                self.file_map.insert(
+                    FileName::Real(mod_path.clone()),
+                    (sub_mod.clone(), item.ident.name.as_str().get().to_owned()),
+                );
+                self.directory = Directory {
+                    path: mod_path.parent().unwrap().to_path_buf(),
+                    ownership: directory_ownership,
+                };
+            }
+            Some(SubModKind::InternalWithPath(mod_path)) => {
+                // All `#[path]` files are treated as though they are a `mod.rs` file.
+                self.directory = Directory {
+                    path: mod_path,
+                    ownership: DirectoryOwnership::Owned { relative: None },
+                };
             }
+            Some(SubModKind::Internal) => self.push_inline_mod_directory(item.ident, &item.attrs),
+            None => (), // rustfmt::skip
         }
         Ok(())
     }
 
+    /// Inspect the given sub-module which we are about to visit and returns its kind.
+    fn peek_sub_mod(&self, item: &'c ast::Item) -> Result<Option<SubModKind>, String> {
+        if contains_skip(&item.attrs) {
+            return Ok(None);
+        }
+
+        if is_mod_decl(item) {
+            // mod foo;
+            // Look for an extern file.
+            let (mod_path, directory_ownership) =
+                self.find_external_module(item.ident, &item.attrs)?;
+            Ok(Some(SubModKind::External(mod_path, directory_ownership)))
+        } else {
+            // An internal module (`mod foo { /* ... */ }`);
+            if let Some(path) = find_path_value(&item.attrs) {
+                let path = Path::new(&path.as_str()).to_path_buf();
+                Ok(Some(SubModKind::InternalWithPath(path)))
+            } else {
+                Ok(Some(SubModKind::Internal))
+            }
+        }
+    }
+
+    /// Find a file path in the filesystem which corresponds to the given module.
     fn find_external_module(
         &self,
         mod_name: ast::Ident,
@@ -123,7 +219,7 @@ impl<'a, 'b> ModResolver<'a, 'b> {
             mod_name,
             relative,
             &self.directory.path,
-            self.source_map,
+            self.parse_sess.source_map(),
         )
         .result
         {

src/modules/visitor.rs

@@ -0,0 +1,105 @@
+use syntax::ast;
+use syntax::parse::token::{DelimToken, Token};
+use syntax::parse::{stream_to_parser_with_base_dir, Directory, ParseSess};
+use syntax::symbol::kw;
+use syntax::visit::Visitor;
+use syntax_pos::Symbol;
+
+pub(crate) struct ModItem {
+    pub(crate) item: ast::Item,
+}
+
+/// Traverse `cfg_if!` macro and fetch modules.
+pub(crate) struct CfgIfVisitor<'a> {
+    parse_sess: &'a ParseSess,
+    mods: Vec<ModItem>,
+    base_dir: Directory<'a>,
+}
+
+impl<'a> CfgIfVisitor<'a> {
+    pub(crate) fn new(parse_sess: &'a ParseSess, base_dir: Directory<'a>) -> CfgIfVisitor<'a> {
+        CfgIfVisitor {
+            mods: vec![],
+            parse_sess,
+            base_dir,
+        }
+    }
+
+    pub(crate) fn mods(self) -> Vec<ModItem> {
+        self.mods
+    }
+}
+
+impl<'a, 'ast: 'a> Visitor<'ast> for CfgIfVisitor<'a> {
+    fn visit_mac(&mut self, mac: &'ast ast::Mac) {
+        match self.visit_mac_inner(mac) {
+            Ok(()) => (),
+            Err(e) => debug!("{}", e),
+        }
+    }
+}
+
+impl<'a, 'ast: 'a> CfgIfVisitor<'a> {
+    fn visit_mac_inner(&mut self, mac: &'ast ast::Mac) -> Result<(), &'static str> {
+        if mac.node.path != Symbol::intern("cfg_if") {
+            return Err("Expected cfg_if");
+        }
+
+        let mut parser = stream_to_parser_with_base_dir(
+            self.parse_sess,
+            mac.node.tts.clone(),
+            self.base_dir.clone(),
+        );
+        parser.cfg_mods = false;
+        let mut process_if_cfg = true;
+
+        while parser.token != Token::Eof {
+            if process_if_cfg {
+                if !parser.eat_keyword(kw::If) {
+                    return Err("Expected `if`");
+                }
+                parser
+                    .parse_attribute(false)
+                    .map_err(|_| "Failed to parse attributes")?;
+            }
+
+            if !parser.eat(&Token::OpenDelim(DelimToken::Brace)) {
+                return Err("Expected an opening brace");
+            }
+
+            while parser.token != Token::CloseDelim(DelimToken::Brace) && parser.token != Token::Eof
+            {
+                let item = match parser.parse_item() {
+                    Ok(Some(item_ptr)) => item_ptr.into_inner(),
+                    Ok(None) => continue,
+                    Err(mut err) => {
+                        err.cancel();
+                        parser.sess.span_diagnostic.reset_err_count();
+                        return Err(
+                            "Expected item inside cfg_if block, but failed to parse it as an item",
+                        );
+                    }
+                };
+                if let ast::ItemKind::Mod(..) = item.node {
+                    self.mods.push(ModItem { item });
+                }
+            }
+
+            if !parser.eat(&Token::CloseDelim(DelimToken::Brace)) {
+                return Err("Expected a closing brace");
+            }
+
+            if parser.eat(&Token::Eof) {
+                break;
+            }
+
+            if !parser.eat_keyword(kw::Else) {
+                return Err("Expected `else`");
+            }
+
+            process_if_cfg = parser.token.is_keyword(kw::If);
+        }
+
+        Ok(())
+    }
+}