Visit last if_stmt

Author: csmoe

src/expr.rs

@@ -545,13 +545,33 @@ impl Rewrite for ast::Stmt {
                 };
 
                 let shape = shape.sub_width(suffix.len())?;
-                let expr_type =
-                    if stmt_is_expr(&self) && context.snippet(self.span).starts_with("if ") {
-                        ExprType::SubExpression
-                    } else {
-                        ExprType::Statement
-                    };
-                format_expr(ex, expr_type, context, shape).map(|s| s + suffix)
+                format_expr(ex, ExprType::Statement, context, shape).map(|s| s + suffix)
+            }
+            ast::StmtKind::Mac(..) | ast::StmtKind::Item(..) => None,
+        };
+        result.and_then(|res| recover_comment_removed(res, self.span(), context))
+    }
+}
+
+pub(crate) trait RewriteLastStmt {
+    fn rewrite_last_stmt(&self, context: &RewriteContext, shape: Shape) -> Option<String>;
+}
+
+impl RewriteLastStmt for ast::Stmt {
+    fn rewrite_last_stmt(&self, context: &RewriteContext, shape: Shape) -> Option<String> {
+        skip_out_of_file_lines_range!(context, self.span());
+
+        let result = match self.node {
+            ast::StmtKind::Local(ref local) => local.rewrite(context, shape),
+            ast::StmtKind::Expr(ref ex) | ast::StmtKind::Semi(ref ex) => {
+                let suffix = if semicolon_for_stmt(context, self) {
+                    ";"
+                } else {
+                    ""
+                };
+
+                let shape = shape.sub_width(suffix.len())?;
+                format_expr(ex, ExprType::SubExpression, context, shape).map(|s| s + suffix)
             }
             ast::StmtKind::Mac(..) | ast::StmtKind::Item(..) => None,
         };
@@ -1136,26 +1156,16 @@ pub fn stmt_is_expr(stmt: &ast::Stmt) -> bool {
     }
 }
 
-fn stmt_is_if(stmt: &ast::Stmt) -> bool {
+pub(crate) fn stmt_is_if(stmt: &ast::Stmt) -> bool {
     match stmt.node {
-        ast::StmtKind::Semi(ref e) | ast::StmtKind::Expr(ref e) => {
-            match e.node {
-                ast::ExprKind::If(..) => true,
-                _ => false,
-            }
+        ast::StmtKind::Semi(ref e) | ast::StmtKind::Expr(ref e) => match e.node {
+            ast::ExprKind::If(..) => true,
+            _ => false,
         },
         _ => false,
     }
 }
 
-fn block_last_stmt_is_if(block: &ast::Block) -> bool {
-    if let Some(ref stmt) = block.stmts.last() {
-        stmt_is_if(stmt)
-    } else {
-        false
-    }
-}
-
 pub fn is_unsafe_block(block: &ast::Block) -> bool {
     if let ast::BlockCheckMode::Unsafe(..) = block.rules {
         true

src/visitor.rs

@@ -17,6 +17,7 @@ use attr::*;
 use codemap::{LineRangeUtils, SpanUtils};
 use comment::{CodeCharKind, CommentCodeSlices, FindUncommented};
 use config::{BraceStyle, Config};
+use expr::{stmt_is_if, RewriteLastStmt};
 use items::{
     format_impl, format_trait, format_trait_alias, is_mod_decl, is_use_item,
     rewrite_associated_impl_type, rewrite_associated_type, rewrite_extern_crate,
@@ -79,7 +80,7 @@ impl<'b, 'a: 'b> FmtVisitor<'a> {
         Shape::indented(self.block_indent, self.config)
     }
 
-    fn visit_stmt(&mut self, stmt: &ast::Stmt) {
+    fn visit_stmt(&mut self, stmt: &ast::Stmt, if_is_last_stmt: bool) {
         debug!(
             "visit_stmt: {:?} {:?}",
             self.codemap.lookup_char_pos(stmt.span.lo()),
@@ -94,7 +95,11 @@ impl<'b, 'a: 'b> FmtVisitor<'a> {
                 if contains_skip(get_attrs_from_stmt(stmt)) {
                     self.push_skipped_with_span(stmt.span());
                 } else {
-                    let rewrite = stmt.rewrite(&self.get_context(), self.shape());
+                    let rewrite = if if_is_last_stmt {
+                        stmt.rewrite_last_stmt(&self.get_context(), self.shape())
+                    } else {
+                        stmt.rewrite(&self.get_context(), self.shape())
+                    };
                     self.push_rewrite(stmt.span(), rewrite)
                 }
             }
@@ -662,7 +667,8 @@ impl<'b, 'a: 'b> FmtVisitor<'a> {
             .collect();
 
         if items.is_empty() {
-            self.visit_stmt(&stmts[0]);
+            let if_is_last_stmt = stmts[0] == stmts[stmts.len() - 1] && stmt_is_if(&stmts[0]);
+            self.visit_stmt(&stmts[0], if_is_last_stmt);
             self.walk_stmts(&stmts[1..]);
         } else {
             self.visit_items_with_reordering(&items);