fix: inner attribute formatting

calebcartwright · calebcartwright · commit 62dc7c5c6b78 · 2020-09-04T19:01:18.000-05:00
diff --git a/src/formatting.rs b/src/formatting.rs
@@ -11,6 +11,7 @@ use self::newline_style::apply_newline_style;
 use crate::comment::{CharClasses, FullCodeCharKind};
 use crate::config::{Config, FileName, Verbosity};
 use crate::issues::BadIssueSeeker;
+use crate::modules::Module;
 use crate::syntux::parser::{DirectoryOwnership, Parser, ParserError};
 use crate::syntux::session::ParseSess;
 use crate::utils::count_newlines;
@@ -102,8 +103,7 @@ fn format_project<T: FormatHandler>(
             continue;
         }
         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)?;
     }
     timer = timer.done_formatting();
 
@@ -134,13 +134,8 @@ impl<'a, T: FormatHandler + 'a> FormatContext<'a, T> {
     }
 
     // Formats a single file/module.
-    fn format_file(
-        &mut self,
-        path: FileName,
-        module: &ast::Mod,
-        is_root: bool,
-    ) -> Result<(), ErrorKind> {
-        let snippet_provider = self.parse_session.snippet_provider(module.inner);
+    fn format_file(&mut self, path: FileName, module: &Module<'_>) -> Result<(), ErrorKind> {
+        let snippet_provider = self.parse_session.snippet_provider(module.as_ref().inner);
         let mut visitor = FmtVisitor::from_parse_sess(
             &self.parse_session,
             &self.config,
@@ -149,19 +144,9 @@ impl<'a, T: FormatHandler + 'a> FormatContext<'a, T> {
         );
         visitor.skip_context.update_with_attrs(&self.krate.attrs);
 
-        // Format inner attributes if available.
-        if !self.krate.attrs.is_empty() && is_root {
-            visitor.skip_empty_lines(snippet_provider.end_pos());
-            if visitor.visit_attrs(&self.krate.attrs, ast::AttrStyle::Inner) {
-                visitor.push_rewrite(module.inner, None);
-            } else {
-                visitor.format_separate_mod(module, snippet_provider.end_pos());
-            }
-        } else {
-            visitor.last_pos = snippet_provider.start_pos();
-            visitor.skip_empty_lines(snippet_provider.end_pos());
-            visitor.format_separate_mod(module, snippet_provider.end_pos());
-        };
+        visitor.last_pos = snippet_provider.start_pos();
+        visitor.skip_empty_lines(snippet_provider.end_pos());
+        visitor.format_separate_mod(module, snippet_provider.end_pos());
 
         debug_assert_eq!(
             visitor.line_number,
diff --git a/src/modules.rs b/src/modules.rs
@@ -3,6 +3,7 @@ use std::collections::BTreeMap;
 use std::path::{Path, PathBuf};
 
 use rustc_ast::ast;
+use rustc_ast::attr::HasAttrs;
 use rustc_ast::visit::Visitor;
 use rustc_span::symbol::{self, sym, Symbol};
 use thiserror::Error;
@@ -18,12 +19,48 @@ use crate::utils::contains_skip;
 
 mod visitor;
 
-type FileModMap<'ast> = BTreeMap<FileName, Cow<'ast, ast::Mod>>;
+type FileModMap<'ast> = BTreeMap<FileName, Module<'ast>>;
 
 lazy_static! {
     static ref CFG_IF: Symbol = Symbol::intern("cfg_if");
 }
 
+/// Represents module with its inner attributes.
+#[derive(Debug, Clone)]
+pub(crate) struct Module<'a> {
+    ast_mod: Cow<'a, ast::Mod>,
+    inner_attr: Vec<ast::Attribute>,
+}
+
+impl<'a> Module<'a> {
+    pub(crate) fn new(ast_mod: Cow<'a, ast::Mod>, attrs: &[ast::Attribute]) -> Self {
+        let inner_attr = attrs
+            .iter()
+            .filter(|attr| attr.style == ast::AttrStyle::Inner)
+            .cloned()
+            .collect();
+        Module {
+            ast_mod,
+            inner_attr,
+        }
+    }
+}
+
+impl<'a> HasAttrs for Module<'a> {
+    fn attrs(&self) -> &[ast::Attribute] {
+        &self.inner_attr
+    }
+    fn visit_attrs(&mut self, f: impl FnOnce(&mut Vec<ast::Attribute>)) {
+        f(&mut self.inner_attr)
+    }
+}
+
+impl<'a> AsRef<ast::Mod> for Module<'a> {
+    fn as_ref(&self) -> &ast::Mod {
+        &self.ast_mod
+    }
+}
+
 /// Maps each module to the corresponding file.
 pub(crate) struct ModResolver<'ast, 'sess> {
     parse_sess: &'sess ParseSess,
@@ -53,9 +90,9 @@ pub(crate) enum ModuleResolutionErrorKind {
 #[derive(Clone)]
 enum SubModKind<'a, 'ast> {
     /// `mod foo;`
-    External(PathBuf, DirectoryOwnership, Cow<'ast, ast::Mod>),
+    External(PathBuf, DirectoryOwnership, Module<'ast>),
     /// `mod foo;` with multiple sources.
-    MultiExternal(Vec<(PathBuf, DirectoryOwnership, Cow<'ast, ast::Mod>)>),
+    MultiExternal(Vec<(PathBuf, DirectoryOwnership, Module<'ast>)>),
     /// `mod foo {}`
     Internal(&'a ast::Item),
 }
@@ -94,8 +131,10 @@ impl<'ast, 'sess, 'c> ModResolver<'ast, 'sess> {
             self.visit_mod_from_ast(&krate.module)?;
         }
 
-        self.file_map
-            .insert(root_filename, Cow::Borrowed(&krate.module));
+        self.file_map.insert(
+            root_filename,
+            Module::new(Cow::Borrowed(&krate.module), &krate.attrs),
+        );
         Ok(self.file_map)
     }
 
@@ -105,7 +144,10 @@ impl<'ast, 'sess, 'c> ModResolver<'ast, 'sess> {
         visitor.visit_item(&item);
         for module_item in visitor.mods() {
             if let ast::ItemKind::Mod(ref sub_mod) = module_item.item.kind {
-                self.visit_sub_mod(&module_item.item, Cow::Owned(sub_mod.clone()))?;
+                self.visit_sub_mod(
+                    &module_item.item,
+                    Module::new(Cow::Owned(sub_mod.clone()), &module_item.item.attrs),
+                )?;
             }
         }
         Ok(())
@@ -120,7 +162,7 @@ impl<'ast, 'sess, 'c> ModResolver<'ast, 'sess> {
             }
 
             if let ast::ItemKind::Mod(ref sub_mod) = item.kind {
-                self.visit_sub_mod(&item, Cow::Owned(sub_mod.clone()))?;
+                self.visit_sub_mod(&item, Module::new(Cow::Owned(sub_mod.clone()), &item.attrs))?;
             }
         }
         Ok(())
@@ -134,7 +176,7 @@ impl<'ast, 'sess, 'c> ModResolver<'ast, 'sess> {
             }
 
             if let ast::ItemKind::Mod(ref sub_mod) = item.kind {
-                self.visit_sub_mod(item, Cow::Borrowed(sub_mod))?;
+                self.visit_sub_mod(item, Module::new(Cow::Borrowed(sub_mod), &item.attrs))?;
             }
         }
         Ok(())
@@ -143,12 +185,12 @@ impl<'ast, 'sess, 'c> ModResolver<'ast, 'sess> {
     fn visit_sub_mod(
         &mut self,
         item: &'c ast::Item,
-        sub_mod: Cow<'ast, ast::Mod>,
+        sub_mod: Module<'ast>,
     ) -> Result<(), ModuleResolutionError> {
         let old_directory = self.directory.clone();
         let sub_mod_kind = self.peek_sub_mod(item, &sub_mod)?;
         if let Some(sub_mod_kind) = sub_mod_kind {
-            self.insert_sub_mod(sub_mod_kind.clone(), sub_mod.clone())?;
+            self.insert_sub_mod(sub_mod_kind.clone())?;
             self.visit_sub_mod_inner(sub_mod, sub_mod_kind)?;
         }
         self.directory = old_directory;
@@ -159,7 +201,7 @@ impl<'ast, 'sess, 'c> ModResolver<'ast, 'sess> {
     fn peek_sub_mod(
         &self,
         item: &'c ast::Item,
-        sub_mod: &Cow<'ast, ast::Mod>,
+        sub_mod: &Module<'ast>,
     ) -> Result<Option<SubModKind<'c, 'ast>>, ModuleResolutionError> {
         if contains_skip(&item.attrs) {
             return Ok(None);
@@ -178,7 +220,6 @@ impl<'ast, 'sess, 'c> ModResolver<'ast, 'sess> {
     fn insert_sub_mod(
         &mut self,
         sub_mod_kind: SubModKind<'c, 'ast>,
-        _sub_mod: Cow<'ast, ast::Mod>,
     ) -> Result<(), ModuleResolutionError> {
         match sub_mod_kind {
             SubModKind::External(mod_path, _, sub_mod) => {
@@ -200,7 +241,7 @@ impl<'ast, 'sess, 'c> ModResolver<'ast, 'sess> {
 
     fn visit_sub_mod_inner(
         &mut self,
-        sub_mod: Cow<'ast, ast::Mod>,
+        sub_mod: Module<'ast>,
         sub_mod_kind: SubModKind<'c, 'ast>,
     ) -> Result<(), ModuleResolutionError> {
         match sub_mod_kind {
@@ -230,13 +271,13 @@ impl<'ast, 'sess, 'c> ModResolver<'ast, 'sess> {
 
     fn visit_sub_mod_after_directory_update(
         &mut self,
-        sub_mod: Cow<'ast, ast::Mod>,
+        sub_mod: Module<'ast>,
         directory: Option<Directory>,
     ) -> Result<(), ModuleResolutionError> {
         if let Some(directory) = directory {
             self.directory = directory;
         }
-        match sub_mod {
+        match sub_mod.ast_mod {
             Cow::Borrowed(sub_mod) => self.visit_mod_from_ast(sub_mod),
             Cow::Owned(sub_mod) => self.visit_mod_outside_ast(sub_mod),
         }
@@ -247,7 +288,7 @@ impl<'ast, 'sess, 'c> ModResolver<'ast, 'sess> {
         &self,
         mod_name: symbol::Ident,
         attrs: &[ast::Attribute],
-        sub_mod: &Cow<'ast, ast::Mod>,
+        sub_mod: &Module<'ast>,
     ) -> Result<Option<SubModKind<'c, 'ast>>, ModuleResolutionError> {
         let relative = match self.directory.ownership {
             DirectoryOwnership::Owned { relative } => relative,
@@ -257,12 +298,13 @@ impl<'ast, 'sess, 'c> ModResolver<'ast, 'sess> {
             if self.parse_sess.is_file_parsed(&path) {
                 return Ok(None);
             }
-            return match Parser::parse_file_as_module(self.parse_sess, &path, sub_mod.inner) {
+            return match Parser::parse_file_as_module(self.parse_sess, &path, sub_mod.ast_mod.inner)
+            {
                 Ok((_, ref attrs)) if contains_skip(attrs) => Ok(None),
-                Ok((m, _)) => Ok(Some(SubModKind::External(
+                Ok(m) => Ok(Some(SubModKind::External(
                     path,
                     DirectoryOwnership::Owned { relative: None },
-                    Cow::Owned(m),
+                    Module::new(Cow::Owned(m.0), &m.1),
                 ))),
                 Err(ParserError::ParseError) => Err(ModuleResolutionError {
                     module: mod_name.to_string(),
@@ -300,13 +342,19 @@ impl<'ast, 'sess, 'c> ModResolver<'ast, 'sess> {
                         return Ok(Some(SubModKind::MultiExternal(mods_outside_ast)));
                     }
                 }
-                match Parser::parse_file_as_module(self.parse_sess, &path, sub_mod.inner) {
+                match Parser::parse_file_as_module(self.parse_sess, &path, sub_mod.ast_mod.inner) {
                     Ok((_, ref attrs)) if contains_skip(attrs) => Ok(None),
-                    Ok((m, _)) if outside_mods_empty => {
-                        Ok(Some(SubModKind::External(path, ownership, Cow::Owned(m))))
-                    }
-                    Ok((m, _)) => {
-                        mods_outside_ast.push((path.clone(), ownership, Cow::Owned(m)));
+                    Ok(m) if outside_mods_empty => Ok(Some(SubModKind::External(
+                        path,
+                        ownership,
+                        Module::new(Cow::Owned(m.0), &m.1),
+                    ))),
+                    Ok(m) => {
+                        mods_outside_ast.push((
+                            path.clone(),
+                            ownership,
+                            Module::new(Cow::Owned(m.0), &m.1),
+                        ));
                         if should_insert {
                             mods_outside_ast.push((path, ownership, sub_mod.clone()));
                         }
@@ -368,8 +416,8 @@ impl<'ast, 'sess, 'c> ModResolver<'ast, 'sess> {
     fn find_mods_outside_of_ast(
         &self,
         attrs: &[ast::Attribute],
-        sub_mod: &Cow<'ast, ast::Mod>,
-    ) -> Vec<(PathBuf, DirectoryOwnership, Cow<'ast, ast::Mod>)> {
+        sub_mod: &Module<'ast>,
+    ) -> Vec<(PathBuf, DirectoryOwnership, Module<'ast>)> {
         // Filter nested path, like `#[cfg_attr(feature = "foo", path = "bar.rs")]`.
         let mut path_visitor = visitor::PathVisitor::default();
         for attr in attrs.iter() {
@@ -393,17 +441,20 @@ impl<'ast, 'sess, 'c> ModResolver<'ast, 'sess> {
                 ));
                 continue;
             }
-            let m = match Parser::parse_file_as_module(self.parse_sess, &actual_path, sub_mod.inner)
-            {
+            let m = match Parser::parse_file_as_module(
+                self.parse_sess,
+                &actual_path,
+                sub_mod.ast_mod.inner,
+            ) {
                 Ok((_, ref attrs)) if contains_skip(attrs) => continue,
-                Ok((m, _)) => m,
+                Ok(m) => m,
                 Err(..) => continue,
             };
 
             result.push((
                 actual_path,
                 DirectoryOwnership::Owned { relative: None },
-                Cow::Owned(m),
+                Module::new(Cow::Owned(m.0), &m.1),
             ))
         }
         result
diff --git a/src/visitor.rs b/src/visitor.rs
@@ -1,7 +1,7 @@
 use std::cell::{Cell, RefCell};
 use std::rc::Rc;
 
-use rustc_ast::{ast, token::DelimToken, visit};
+use rustc_ast::{ast, attr::HasAttrs, token::DelimToken, visit};
 use rustc_span::{symbol, BytePos, Pos, Span};
 
 use crate::attr::*;
@@ -16,6 +16,7 @@ use crate::items::{
     StaticParts, StructParts,
 };
 use crate::macros::{macro_style, rewrite_macro, rewrite_macro_def, MacroPosition};
+use crate::modules::Module;
 use crate::rewrite::{Rewrite, RewriteContext};
 use crate::shape::{Indent, Shape};
 use crate::skip::{is_skip_attr, SkipContext};
@@ -938,10 +939,14 @@ impl<'b, 'a: 'b> FmtVisitor<'a> {
         }
     }
 
-    pub(crate) fn format_separate_mod(&mut self, m: &ast::Mod, end_pos: BytePos) {
+    pub(crate) fn format_separate_mod(&mut self, m: &Module<'_>, end_pos: BytePos) {
         self.block_indent = Indent::empty();
-        self.walk_mod_items(m);
-        self.format_missing_with_indent(end_pos);
+        if self.visit_attrs(m.attrs(), ast::AttrStyle::Inner) {
+            self.push_skipped_with_span(m.attrs(), m.as_ref().inner, m.as_ref().inner);
+        } else {
+            self.walk_mod_items(m.as_ref());
+            self.format_missing_with_indent(end_pos);
+        }
     }
 
     pub(crate) fn skip_empty_lines(&mut self, end_pos: BytePos) {
diff --git a/tests/target/inner-module-path/lib.rs b/tests/target/inner-module-path/lib.rs
@@ -1,5 +1,3 @@
-// rustfmt-recursive: true
-
 #[path = "."]
 mod a {
     mod b;

Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,3 @@`
`1`		`-// rustfmt-recursive: true`
`2`		`-`
`3`	`1`	`#[path = "."]`
`4`	`2`	`mod a {`
`5`	`3`	`mod b;`