Merge pull request #1933 from topecongiro/attributes-on-stmt

Format and preserve attributes on `ast::Stmt`

src/expr.rs

@@ -130,11 +130,10 @@ pub fn format_expr(
                 ExprType::Statement => {
                     if is_unsafe_block(block) {
                         block.rewrite(context, shape)
-                    } else {
+                    } else if let rw @ Some(_) = rewrite_empty_block(context, block, shape) {
                         // Rewrite block without trying to put it in a single line.
-                        if let rw @ Some(_) = rewrite_empty_block(context, block, shape) {
-                            return rw;
-                        }
+                        rw
+                    } else {
                         let prefix = try_opt!(block_prefix(context, block, shape));
                         rewrite_block_with_visitor(context, &prefix, block, shape)
                     }
@@ -181,17 +180,11 @@ pub fn format_expr(
                 )
             }
         }
-        ast::ExprKind::Yield(ref opt_expr) => {
-            if let Some(ref expr) = *opt_expr {
-                rewrite_unary_prefix(context, "yield ", &**expr, shape)
-            } else {
-                wrap_str(
-                    "yield".to_string(),
-                    context.config.max_width(),
-                    shape,
-                )
-            }
-        }
+        ast::ExprKind::Yield(ref opt_expr) => if let Some(ref expr) = *opt_expr {
+            rewrite_unary_prefix(context, "yield ", &**expr, shape)
+        } else {
+            wrap_str("yield".to_string(), context.config.max_width(), shape)
+        },
         ast::ExprKind::Closure(capture, ref fn_decl, ref body, _) => {
             rewrite_closure(capture, fn_decl, body, expr.span, context, shape)
         }
@@ -293,17 +286,17 @@ pub fn format_expr(
             shape,
         ),
         ast::ExprKind::Catch(ref block) => {
-            if let rewrite @ Some(_) = rewrite_single_line_block(context, "do catch ", block, shape)
-            {
-                return rewrite;
+            if let rw @ Some(_) = rewrite_single_line_block(context, "do catch ", block, shape) {
+                rw
+            } else {
+                // 9 = `do catch `
+                let budget = shape.width.checked_sub(9).unwrap_or(0);
+                Some(format!(
+                    "{}{}",
+                    "do catch ",
+                    try_opt!(block.rewrite(&context, Shape::legacy(budget, shape.indent)))
+                ))
             }
-            // 9 = `do catch `
-            let budget = shape.width.checked_sub(9).unwrap_or(0);
-            Some(format!(
-                "{}{}",
-                "do catch ",
-                try_opt!(block.rewrite(&context, Shape::legacy(budget, shape.indent)))
-            ))
         }
     };
 
@@ -883,16 +876,13 @@ impl Rewrite for ast::Stmt {
                     ""
                 };
 
-                format_expr(
-                    ex,
-                    match self.node {
-                        ast::StmtKind::Expr(_) => ExprType::SubExpression,
-                        ast::StmtKind::Semi(_) => ExprType::Statement,
-                        _ => unreachable!(),
-                    },
-                    context,
-                    try_opt!(shape.sub_width(suffix.len())),
-                ).map(|s| s + suffix)
+                let expr_type = match self.node {
+                    ast::StmtKind::Expr(_) => ExprType::SubExpression,
+                    ast::StmtKind::Semi(_) => ExprType::Statement,
+                    _ => unreachable!(),
+                };
+                let shape = try_opt!(shape.sub_width(suffix.len()));
+                format_expr(ex, expr_type, context, shape).map(|s| s + suffix)
             }
             ast::StmtKind::Mac(..) | ast::StmtKind::Item(..) => None,
         };

src/imports.rs

@@ -260,7 +260,7 @@ impl<'a> FmtVisitor<'a> {
     ) {
         let vis = utils::format_visibility(vis);
         // 4 = `use `, 1 = `;`
-        let rw = Shape::indented(self.block_indent, self.config)
+        let rw = self.shape()
             .offset_left(vis.len() + 4)
             .and_then(|shape| shape.sub_width(1))
             .and_then(|shape| match vp.node {

src/items.rs

@@ -55,7 +55,23 @@ impl Rewrite for ast::Local {
 
         skip_out_of_file_lines_range!(context, self.span);
 
-        let mut result = "let ".to_owned();
+        if contains_skip(&self.attrs) {
+            return None;
+        }
+
+        let attrs_str = try_opt!(self.attrs.rewrite(context, shape));
+        let mut result = if attrs_str.is_empty() {
+            "let ".to_owned()
+        } else {
+            try_opt!(combine_strs_with_missing_comments(
+                context,
+                &attrs_str,
+                "let ",
+                mk_sp(self.attrs.last().map(|a| a.span.hi).unwrap(), self.span.lo),
+                shape,
+                false,
+            ))
+        };
 
         // 4 = "let ".len()
         let pat_shape = try_opt!(shape.offset_left(4));
@@ -187,8 +203,7 @@ impl<'a> FmtVisitor<'a> {
 
 
     fn format_foreign_item(&mut self, item: &ast::ForeignItem) {
-        let shape = Shape::indented(self.block_indent, self.config);
-        let rewrite = item.rewrite(&self.get_context(), shape);
+        let rewrite = item.rewrite(&self.get_context(), self.shape());
         self.push_rewrite(item.span(), rewrite);
         self.last_pos = item.span.hi;
     }
@@ -312,18 +327,11 @@ impl<'a> FmtVisitor<'a> {
                                 ""
                             };
 
-                            format_expr(
-                                &e,
-                                ExprType::Statement,
-                                &self.get_context(),
-                                Shape::indented(self.block_indent, self.config),
-                            ).map(|s| s + suffix)
+                            format_expr(&e, ExprType::Statement, &self.get_context(), self.shape())
+                                .map(|s| s + suffix)
                                 .or_else(|| Some(self.snippet(e.span)))
                         }
-                        None => stmt.rewrite(
-                            &self.get_context(),
-                            Shape::indented(self.block_indent, self.config),
-                        ),
+                        None => stmt.rewrite(&self.get_context(), self.shape()),
                     }
                 } else {
                     None
@@ -421,9 +429,7 @@ impl<'a> FmtVisitor<'a> {
             false,
         );
 
-        let shape = Shape::indented(self.block_indent, self.config)
-            .sub_width(2)
-            .unwrap();
+        let shape = self.shape().sub_width(2).unwrap();
         let fmt = ListFormatting {
             tactic: DefinitiveListTactic::Vertical,
             separator: ",",
@@ -451,7 +457,7 @@ impl<'a> FmtVisitor<'a> {
 
         let context = self.get_context();
         let indent = self.block_indent;
-        let shape = Shape::indented(indent, self.config);
+        let shape = self.shape();
         let attrs_str = try_opt!(field.node.attrs.rewrite(&context, shape));
         let lo = field
             .node

src/lib.rs

@@ -91,32 +91,46 @@ macro_rules! span_with_attrs_lo_hi {
         }
     }
 }
+
 macro_rules! span_with_attrs {
     ($this:ident) => {
         span_with_attrs_lo_hi!($this, $this.span.lo, $this.span.hi)
     }
 }
 
-impl Spanned for ast::Expr {
-    fn span(&self) -> Span {
-        span_with_attrs!(self)
+macro_rules! implement_spanned {
+    ($this:ty) => {
+        impl Spanned for $this {
+            fn span(&self) -> Span {
+                span_with_attrs!(self)
+            }
+        }
     }
 }
 
-impl Spanned for ast::Item {
-    fn span(&self) -> Span {
-        span_with_attrs!(self)
-    }
-}
+// Implement `Spanned` for structs with `attrs` field.
+implement_spanned!(ast::Expr);
+implement_spanned!(ast::Field);
+implement_spanned!(ast::ForeignItem);
+implement_spanned!(ast::Item);
+implement_spanned!(ast::Local);
 
 impl Spanned for ast::Stmt {
     fn span(&self) -> Span {
         match self.node {
-            // Cover attributes
+            ast::StmtKind::Local(ref local) => mk_sp(local.span().lo, self.span.hi),
+            ast::StmtKind::Item(ref item) => mk_sp(item.span().lo, self.span.hi),
             ast::StmtKind::Expr(ref expr) | ast::StmtKind::Semi(ref expr) => {
                 mk_sp(expr.span().lo, self.span.hi)
             }
-            _ => self.span,
+            ast::StmtKind::Mac(ref mac) => {
+                let (_, _, ref attrs) = **mac;
+                if attrs.is_empty() {
+                    self.span
+                } else {
+                    mk_sp(attrs[0].span.lo, self.span.hi)
+                }
+            }
         }
     }
 }
@@ -155,12 +169,6 @@ impl Spanned for ast::StructField {
     }
 }
 
-impl Spanned for ast::Field {
-    fn span(&self) -> Span {
-        span_with_attrs!(self)
-    }
-}
-
 impl Spanned for ast::WherePredicate {
     fn span(&self) -> Span {
         match *self {
@@ -208,12 +216,6 @@ impl Spanned for ast::TyParamBound {
     }
 }
 
-impl Spanned for ast::ForeignItem {
-    fn span(&self) -> Span {
-        span_with_attrs!(self)
-    }
-}
-
 #[derive(Copy, Clone, Debug)]
 pub struct Indent {
     // Width of the block indent, in characters. Must be a multiple of

src/visitor.rs

@@ -30,7 +30,7 @@ use lists::{itemize_list, write_list, DefinitiveListTactic, ListFormatting, Sepa
 use macros::{rewrite_macro, MacroPosition};
 use regex::Regex;
 use rewrite::{Rewrite, RewriteContext};
-use utils::{self, contains_skip, mk_sp};
+use utils::{self, contains_skip, inner_attributes, mk_sp};
 
 fn is_use_item(item: &ast::Item) -> bool {
     match item.node {
@@ -58,6 +58,10 @@ pub struct FmtVisitor<'a> {
 }
 
 impl<'a> FmtVisitor<'a> {
+    pub fn shape(&self) -> Shape {
+        Shape::indented(self.block_indent, self.config)
+    }
+
     fn visit_stmt(&mut self, stmt: &ast::Stmt) {
         debug!(
             "visit_stmt: {:?} {:?}",
@@ -69,47 +73,23 @@ impl<'a> FmtVisitor<'a> {
             ast::StmtKind::Item(ref item) => {
                 self.visit_item(item);
             }
-            ast::StmtKind::Local(ref local) => {
-                let rewrite = if contains_skip(&local.attrs) {
-                    None
-                } else {
-                    stmt.rewrite(
-                        &self.get_context(),
-                        Shape::indented(self.block_indent, self.config),
-                    )
-                };
-                self.push_rewrite(stmt.span, rewrite);
+            ast::StmtKind::Local(..) => {
+                let rewrite = stmt.rewrite(&self.get_context(), self.shape());
+                self.push_rewrite(stmt.span(), rewrite);
             }
             ast::StmtKind::Expr(ref expr) => {
-                let rewrite = format_expr(
-                    expr,
-                    ExprType::Statement,
-                    &self.get_context(),
-                    Shape::indented(self.block_indent, self.config),
-                );
-                let span = if expr.attrs.is_empty() {
-                    stmt.span
-                } else {
-                    mk_sp(expr.span().lo, stmt.span.hi)
-                };
-                self.push_rewrite(span, rewrite)
+                let rewrite =
+                    format_expr(expr, ExprType::Statement, &self.get_context(), self.shape());
+                self.push_rewrite(stmt.span(), rewrite)
             }
-            ast::StmtKind::Semi(ref expr) => {
-                let rewrite = stmt.rewrite(
-                    &self.get_context(),
-                    Shape::indented(self.block_indent, self.config),
-                );
-                let span = if expr.attrs.is_empty() {
-                    stmt.span
-                } else {
-                    mk_sp(expr.span().lo, stmt.span.hi)
-                };
-                self.push_rewrite(span, rewrite)
+            ast::StmtKind::Semi(..) => {
+                let rewrite = stmt.rewrite(&self.get_context(), self.shape());
+                self.push_rewrite(stmt.span(), rewrite)
             }
             ast::StmtKind::Mac(ref mac) => {
                 let (ref mac, _macro_style, ref attrs) = **mac;
-                if contains_skip(attrs) {
-                    self.push_rewrite(mac.span, None);
+                if self.visit_attrs(attrs, ast::AttrStyle::Outer) {
+                    self.push_rewrite(stmt.span(), None);
                 } else {
                     self.visit_mac(mac, None, MacroPosition::Statement);
                 }
@@ -138,9 +118,7 @@ impl<'a> FmtVisitor<'a> {
             if let Some(first_stmt) = b.stmts.first() {
                 let attr_lo = inner_attrs
                     .and_then(|attrs| {
-                        utils::inner_attributes(attrs)
-                            .first()
-                            .map(|attr| attr.span.lo)
+                        inner_attributes(attrs).first().map(|attr| attr.span.lo)
                     })
                     .or_else(|| {
                         // Attributes for an item in a statement position
@@ -218,7 +196,7 @@ impl<'a> FmtVisitor<'a> {
         let mut unindent_comment = self.is_if_else_block && !b.stmts.is_empty();
         if unindent_comment {
             let end_pos = source!(self, b.span).hi - brace_compensation - remove_len;
-            let snippet = self.get_context().snippet(mk_sp(self.last_pos, end_pos));
+            let snippet = self.snippet(mk_sp(self.last_pos, end_pos));
             unindent_comment = snippet.contains("//") || snippet.contains("/*");
         }
         // FIXME: we should compress any newlines here to just one
@@ -336,7 +314,7 @@ impl<'a> FmtVisitor<'a> {
                         self.push_rewrite(item.span, None);
                         return;
                     }
-                } else if utils::contains_skip(&item.attrs) {
+                } else if contains_skip(&item.attrs) {
                     // Module is not inline, but should be skipped.
                     return;
                 } else {
@@ -371,7 +349,7 @@ impl<'a> FmtVisitor<'a> {
             }
             ast::ItemKind::Impl(..) => {
                 self.format_missing_with_indent(source!(self, item.span).lo);
-                let snippet = self.get_context().snippet(item.span);
+                let snippet = self.snippet(item.span);
                 let where_span_end = snippet
                     .find_uncommented("{")
                     .map(|x| (BytePos(x as u32)) + source!(self, item.span).lo);
@@ -635,9 +613,7 @@ impl<'a> FmtVisitor<'a> {
         skip_out_of_file_lines_range_visitor!(self, mac.span);
 
         // 1 = ;
-        let shape = Shape::indented(self.block_indent, self.config)
-            .sub_width(1)
-            .unwrap();
+        let shape = self.shape().sub_width(1).unwrap();
         let rewrite = rewrite_macro(mac, ident, &self.get_context(), shape, pos);
         self.push_rewrite(mac.span, rewrite);
     }
@@ -677,7 +653,7 @@ impl<'a> FmtVisitor<'a> {
 
     // Returns true if we should skip the following item.
     pub fn visit_attrs(&mut self, attrs: &[ast::Attribute], style: ast::AttrStyle) -> bool {
-        if utils::contains_skip(attrs) {
+        if contains_skip(attrs) {
             return true;
         }
 
@@ -686,10 +662,7 @@ impl<'a> FmtVisitor<'a> {
             return false;
         }
 
-        let rewrite = attrs.rewrite(
-            &self.get_context(),
-            Shape::indented(self.block_indent, self.config),
-        );
+        let rewrite = attrs.rewrite(&self.get_context(), self.shape());
         let span = mk_sp(attrs[0].span.lo, attrs[attrs.len() - 1].span.hi);
         self.push_rewrite(span, rewrite);
 

tests/source/attrib.rs

@@ -124,3 +124,25 @@ impl InnerAttributes() {
 mod InnerAttributes {
     #![ this_is_an_inner_attribute ( foo ) ]
 }
+
+fn attributes_on_statements() {
+    // Local
+    # [ attr ( on ( local ) ) ]
+    let x = 3;
+
+    // Item
+    # [ attr ( on ( item ) ) ]
+    use foo;
+
+    // Expr
+    # [ attr ( on ( expr ) ) ]
+    {}
+
+    // Semi
+    # [ attr ( on ( semi ) ) ]
+    foo();
+
+    // Mac
+    # [ attr ( on ( mac ) ) ]
+    foo!();
+}

tests/target/attrib.rs

@@ -124,3 +124,25 @@ impl InnerAttributes() {
 mod InnerAttributes {
     #![this_is_an_inner_attribute(foo)]
 }
+
+fn attributes_on_statements() {
+    // Local
+    #[attr(on(local))]
+    let x = 3;
+
+    // Item
+    #[attr(on(item))]
+    use foo;
+
+    // Expr
+    #[attr(on(expr))]
+    {}
+
+    // Semi
+    #[attr(on(semi))]
+    foo();
+
+    // Mac
+    #[attr(on(mac))]
+    foo!();
+}