Format the last expression-statement as expression

src/expr.rs

@@ -33,7 +33,7 @@ use crate::types::{rewrite_path, PathContext};
 use crate::utils::{
     colon_spaces, contains_skip, count_newlines, first_line_ends_with, inner_attributes,
     last_line_extendable, last_line_width, mk_sp, outer_attributes, ptr_vec_to_ref_vec,
-    semicolon_for_expr, semicolon_for_stmt, unicode_str_width, wrap_str,
+    semicolon_for_expr, unicode_str_width, wrap_str,
 };
 use crate::vertical::rewrite_with_alignment;
 use crate::visitor::FmtVisitor;
@@ -568,28 +568,6 @@ fn rewrite_block(
     result
 }
 
-impl Rewrite for ast::Stmt {
-    fn rewrite(&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::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))
-    }
-}
-
 // Rewrite condition if the given expression has one.
 pub(crate) fn rewrite_cond(
     context: &RewriteContext<'_>,
@@ -1189,16 +1167,6 @@ pub(crate) fn stmt_is_expr(stmt: &ast::Stmt) -> bool {
     }
 }
 
-pub(crate) fn stmt_is_if(stmt: &ast::Stmt) -> bool {
-    match stmt.node {
-        ast::StmtKind::Expr(ref e) => match e.node {
-            ast::ExprKind::If(..) => true,
-            _ => false,
-        },
-        _ => false,
-    }
-}
-
 pub(crate) fn is_unsafe_block(block: &ast::Block) -> bool {
     if let ast::BlockCheckMode::Unsafe(..) = block.rules {
         true

src/items.rs

@@ -18,8 +18,7 @@ use crate::comment::{
 use crate::config::lists::*;
 use crate::config::{BraceStyle, Config, IndentStyle, Version};
 use crate::expr::{
-    format_expr, is_empty_block, is_simple_block_stmt, rewrite_assign_rhs, rewrite_assign_rhs_with,
-    ExprType, RhsTactics,
+    is_empty_block, is_simple_block_stmt, rewrite_assign_rhs, rewrite_assign_rhs_with, RhsTactics,
 };
 use crate::lists::{definitive_tactic, itemize_list, write_list, ListFormatting, Separator};
 use crate::macros::{rewrite_macro, MacroPosition};
@@ -28,6 +27,7 @@ use crate::rewrite::{Rewrite, RewriteContext};
 use crate::shape::{Indent, Shape};
 use crate::source_map::{LineRangeUtils, SpanUtils};
 use crate::spanned::Spanned;
+use crate::stmt::Stmt;
 use crate::utils::*;
 use crate::vertical::rewrite_with_alignment;
 use crate::visitor::FmtVisitor;
@@ -394,20 +394,8 @@ impl<'a> FmtVisitor<'a> {
             return None;
         }
 
-        let stmt = block.stmts.first()?;
-        let res = match stmt_expr(stmt) {
-            Some(e) => {
-                let suffix = if semicolon_for_expr(&self.get_context(), e) {
-                    ";"
-                } else {
-                    ""
-                };
-
-                format_expr(e, ExprType::Statement, &self.get_context(), self.shape())
-                    .map(|s| s + suffix)?
-            }
-            None => stmt.rewrite(&self.get_context(), self.shape())?,
-        };
+        let res = Stmt::from_ast_node(block.stmts.first()?, true)
+            .rewrite(&self.get_context(), self.shape())?;
 
         let width = self.block_indent.width() + fn_str.len() + res.len() + 5;
         if !res.contains('\n') && width <= self.config.max_width() {

src/lib.rs

@@ -72,6 +72,7 @@ mod shape;
 pub(crate) mod source_file;
 pub(crate) mod source_map;
 mod spanned;
+mod stmt;
 mod string;
 #[cfg(test)]
 mod test;

src/stmt.rs

@@ -0,0 +1,112 @@
+use syntax::ast;
+use syntax_pos::Span;
+
+use crate::comment::recover_comment_removed;
+use crate::config::Version;
+use crate::expr::{format_expr, ExprType};
+use crate::rewrite::{Rewrite, RewriteContext};
+use crate::shape::Shape;
+use crate::source_map::LineRangeUtils;
+use crate::spanned::Spanned;
+use crate::utils::semicolon_for_stmt;
+
+pub(crate) struct Stmt<'a> {
+    inner: &'a ast::Stmt,
+    is_last: bool,
+}
+
+impl<'a> Spanned for Stmt<'a> {
+    fn span(&self) -> Span {
+        self.inner.span()
+    }
+}
+
+impl<'a> Stmt<'a> {
+    pub(crate) fn as_ast_node(&self) -> &ast::Stmt {
+        self.inner
+    }
+
+    pub(crate) fn to_item(&self) -> Option<&ast::Item> {
+        match self.inner.node {
+            ast::StmtKind::Item(ref item) => Some(&**item),
+            _ => None,
+        }
+    }
+
+    pub(crate) fn from_ast_node(inner: &'a ast::Stmt, is_last: bool) -> Self {
+        Stmt { inner, is_last }
+    }
+
+    pub(crate) fn from_ast_nodes<I>(iter: I) -> Vec<Self>
+    where
+        I: Iterator<Item = &'a ast::Stmt>,
+    {
+        let mut result = vec![];
+        let mut iter = iter.peekable();
+        while iter.peek().is_some() {
+            result.push(Stmt {
+                inner: iter.next().unwrap(),
+                is_last: iter.peek().is_none(),
+            })
+        }
+        result
+    }
+
+    fn is_last_expr(&self) -> bool {
+        if !self.is_last {
+            return false;
+        }
+
+        match self.as_ast_node().node {
+            ast::StmtKind::Expr(ref expr) => match expr.node {
+                ast::ExprKind::Ret(..) | ast::ExprKind::Continue(..) | ast::ExprKind::Break(..) => {
+                    false
+                }
+                _ => true,
+            },
+            _ => false,
+        }
+    }
+}
+
+impl<'a> Rewrite for Stmt<'a> {
+    fn rewrite(&self, context: &RewriteContext<'_>, shape: Shape) -> Option<String> {
+        let expr_type = if context.config.version() == Version::Two && self.is_last_expr() {
+            ExprType::SubExpression
+        } else {
+            ExprType::Statement
+        };
+        format_stmt(context, shape, self.as_ast_node(), expr_type)
+    }
+}
+
+impl Rewrite for ast::Stmt {
+    fn rewrite(&self, context: &RewriteContext<'_>, shape: Shape) -> Option<String> {
+        format_stmt(context, shape, self, ExprType::Statement)
+    }
+}
+
+fn format_stmt(
+    context: &RewriteContext<'_>,
+    shape: Shape,
+    stmt: &ast::Stmt,
+    expr_type: ExprType,
+) -> Option<String> {
+    skip_out_of_file_lines_range!(context, stmt.span());
+
+    let result = match stmt.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, stmt) {
+                ";"
+            } else {
+                ""
+            };
+
+            let shape = shape.sub_width(suffix.len())?;
+            format_expr(ex, expr_type, context, shape).map(|s| s + suffix)
+        }
+        ast::StmtKind::Mac(..) | ast::StmtKind::Item(..) => None,
+    };
+    result.and_then(|res| recover_comment_removed(res, stmt.span(), context))
+}

src/visitor.rs

@@ -7,8 +7,7 @@ use syntax::{ast, visit};
 use crate::attr::*;
 use crate::comment::{CodeCharKind, CommentCodeSlices};
 use crate::config::file_lines::FileName;
-use crate::config::{BraceStyle, Config, Version};
-use crate::expr::{format_expr, ExprType};
+use crate::config::{BraceStyle, Config};
 use crate::items::{
     format_impl, format_trait, format_trait_alias, is_mod_decl, is_use_item,
     rewrite_associated_impl_type, rewrite_associated_type, rewrite_existential_impl_type,
@@ -20,6 +19,7 @@ use crate::rewrite::{Rewrite, RewriteContext};
 use crate::shape::{Indent, Shape};
 use crate::source_map::{LineRangeUtils, SpanUtils};
 use crate::spanned::Spanned;
+use crate::stmt::Stmt;
 use crate::utils::{
     self, contains_skip, count_newlines, depr_skip_annotation, get_skip_macro_names,
     inner_attributes, mk_sp, ptr_vec_to_ref_vec, rewrite_ident, stmt_expr,
@@ -89,23 +89,27 @@ 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: &Stmt<'_>) {
         debug!(
             "visit_stmt: {:?} {:?}",
-            self.source_map.lookup_char_pos(stmt.span.lo()),
-            self.source_map.lookup_char_pos(stmt.span.hi())
+            self.source_map.lookup_char_pos(stmt.span().lo()),
+            self.source_map.lookup_char_pos(stmt.span().hi())
         );
 
-        match stmt.node {
+        match stmt.as_ast_node().node {
             ast::StmtKind::Item(ref item) => {
                 self.visit_item(item);
                 // Handle potential `;` after the item.
-                self.format_missing(stmt.span.hi());
+                self.format_missing(stmt.span().hi());
             }
             ast::StmtKind::Local(..) | ast::StmtKind::Expr(..) | ast::StmtKind::Semi(..) => {
-                let attrs = get_attrs_from_stmt(stmt);
+                let attrs = get_attrs_from_stmt(stmt.as_ast_node());
                 if contains_skip(attrs) {
-                    self.push_skipped_with_span(attrs, stmt.span(), get_span_without_attrs(stmt));
+                    self.push_skipped_with_span(
+                        attrs,
+                        stmt.span(),
+                        get_span_without_attrs(stmt.as_ast_node()),
+                    );
                 } else {
                     let shape = self.shape();
                     let rewrite = self.with_context(|ctx| stmt.rewrite(&ctx, shape));
@@ -115,11 +119,15 @@ impl<'b, 'a: 'b> FmtVisitor<'a> {
             ast::StmtKind::Mac(ref mac) => {
                 let (ref mac, _macro_style, ref attrs) = **mac;
                 if self.visit_attrs(attrs, ast::AttrStyle::Outer) {
-                    self.push_skipped_with_span(attrs, stmt.span(), get_span_without_attrs(stmt));
+                    self.push_skipped_with_span(
+                        attrs,
+                        stmt.span(),
+                        get_span_without_attrs(stmt.as_ast_node()),
+                    );
                 } else {
                     self.visit_mac(mac, None, MacroPosition::Statement);
                 }
-                self.format_missing(stmt.span.hi());
+                self.format_missing(stmt.span().hi());
             }
         }
     }
@@ -717,50 +725,29 @@ impl<'b, 'a: 'b> FmtVisitor<'a> {
         self.visit_items_with_reordering(&ptr_vec_to_ref_vec(&m.items));
     }
 
-    fn walk_stmts(&mut self, stmts: &[ast::Stmt]) {
-        fn to_stmt_item(stmt: &ast::Stmt) -> Option<&ast::Item> {
-            match stmt.node {
-                ast::StmtKind::Item(ref item) => Some(&**item),
-                _ => None,
-            }
-        }
-
+    fn walk_stmts(&mut self, stmts: &[Stmt<'_>]) {
         if stmts.is_empty() {
             return;
         }
 
         // Extract leading `use ...;`.
         let items: Vec<_> = stmts
             .iter()
-            .take_while(|stmt| to_stmt_item(stmt).map_or(false, is_use_item))
-            .filter_map(|stmt| to_stmt_item(stmt))
+            .take_while(|stmt| stmt.to_item().map_or(false, is_use_item))
+            .filter_map(|stmt| stmt.to_item())
             .collect();
 
         if items.is_empty() {
-            // The `if` expression at the end of the block should be formatted in a single
-            // line if possible.
-            if self.config.version() == Version::Two
-                && stmts.len() == 1
-                && crate::expr::stmt_is_if(&stmts[0])
-                && !contains_skip(get_attrs_from_stmt(&stmts[0]))
-            {
-                let shape = self.shape();
-                let rewrite = self.with_context(|ctx| {
-                    format_expr(stmt_expr(&stmts[0])?, ExprType::SubExpression, ctx, shape)
-                });
-                self.push_rewrite(stmts[0].span(), rewrite);
-            } else {
-                self.visit_stmt(&stmts[0]);
-                self.walk_stmts(&stmts[1..]);
-            }
+            self.visit_stmt(&stmts[0]);
+            self.walk_stmts(&stmts[1..]);
         } else {
             self.visit_items_with_reordering(&items);
             self.walk_stmts(&stmts[items.len()..]);
         }
     }
 
     fn walk_block_stmts(&mut self, b: &ast::Block) {
-        self.walk_stmts(&b.stmts)
+        self.walk_stmts(&Stmt::from_ast_nodes(b.stmts.iter()))
     }
 
     fn format_mod(

tests/source/fn-single-line.rs → tests/source/fn-single-line/version_one.rs

@@ -1,4 +1,5 @@
 // rustfmt-fn_single_line: true
+// rustfmt-version: One
 // Test single-line functions.
 
 fn foo_expr() {