From d9a09925d7ecfb46fb6a218dbf5479b51acc67c6 Mon Sep 17 00:00:00 2001
From: Cosmic Horror <CosmicHorrorDev@pm.me>
Date: Sat, 24 Jun 2023 12:35:20 -0600
Subject: [PATCH] Switch `ast::Stmt` rewriting use to
 `stmt::Stmt::from_simple_block`

---
 src/expr.rs  | 23 ++++++++++++-----------
 src/stmt.rs  | 27 ++++++++++++++++++---------
 src/utils.rs |  4 ++--
 3 files changed, 32 insertions(+), 22 deletions(-)

diff --git a/src/expr.rs b/src/expr.rs
index 41cc3651992..9a605d6427e 100644
--- a/src/expr.rs
+++ b/src/expr.rs
@@ -26,6 +26,7 @@
 use crate::shape::{Indent, Shape};
 use crate::source_map::{LineRangeUtils, SpanUtils};
 use crate::spanned::Spanned;
+use crate::stmt;
 use crate::string::{rewrite_string, StringFormat};
 use crate::types::{rewrite_path, PathContext};
 use crate::utils::{
@@ -515,9 +516,9 @@ fn rewrite_single_line_block(
     label: Option<ast::Label>,
     shape: Shape,
 ) -> Option<String> {
-    if is_simple_block(context, block, attrs) {
+    if let Some(block_expr) = stmt::Stmt::from_simple_block(context, block, attrs) {
         let expr_shape = shape.offset_left(last_line_width(prefix))?;
-        let expr_str = block.stmts[0].rewrite(context, expr_shape)?;
+        let expr_str = block_expr.rewrite(context, expr_shape)?;
         let label_str = rewrite_label(label);
         let result = format!("{}{}{{ {} }}", prefix, label_str, expr_str);
         if result.len() <= shape.width && !result.contains('\n') {
@@ -799,19 +800,19 @@ fn rewrite_single_line(
         let fixed_cost = self.keyword.len() + "  {  } else {  }".len();
 
         if let ast::ExprKind::Block(ref else_node, _) = else_block.kind {
-            if !is_simple_block(context, self.block, None)
-                || !is_simple_block(context, else_node, None)
-                || pat_expr_str.contains('\n')
-            {
-                return None;
-            }
+            let (if_expr, else_expr) = match (
+                stmt::Stmt::from_simple_block(context, self.block, None),
+                stmt::Stmt::from_simple_block(context, else_node, None),
+                pat_expr_str.contains('\n'),
+            ) {
+                (Some(if_expr), Some(else_expr), false) => (if_expr, else_expr),
+                _ => return None,
+            };
 
             let new_width = width.checked_sub(pat_expr_str.len() + fixed_cost)?;
-            let expr = &self.block.stmts[0];
-            let if_str = expr.rewrite(context, Shape::legacy(new_width, Indent::empty()))?;
+            let if_str = if_expr.rewrite(context, Shape::legacy(new_width, Indent::empty()))?;
 
             let new_width = new_width.checked_sub(if_str.len())?;
-            let else_expr = &else_node.stmts[0];
             let else_str = else_expr.rewrite(context, Shape::legacy(new_width, Indent::empty()))?;
 
             if if_str.contains('\n') || else_str.contains('\n') {
diff --git a/src/stmt.rs b/src/stmt.rs
index 0148de8190b..e3fe4ebca11 100644
--- a/src/stmt.rs
+++ b/src/stmt.rs
@@ -3,7 +3,7 @@
 
 use crate::comment::recover_comment_removed;
 use crate::config::Version;
-use crate::expr::{format_expr, ExprType};
+use crate::expr::{format_expr, is_simple_block, ExprType};
 use crate::rewrite::{Rewrite, RewriteContext};
 use crate::shape::Shape;
 use crate::source_map::LineRangeUtils;
@@ -33,6 +33,21 @@ pub(crate) fn to_item(&self) -> Option<&ast::Item> {
         }
     }
 
+    pub(crate) fn from_simple_block(
+        context: &RewriteContext<'_>,
+        block: &'a ast::Block,
+        attrs: Option<&[ast::Attribute]>,
+    ) -> Option<Self> {
+        if is_simple_block(context, block, attrs) {
+            let inner = &block.stmts[0];
+            // Simple blocks only contain one expr and no stmts
+            let is_last = true;
+            Some(Stmt { inner, is_last })
+        } else {
+            None
+        }
+    }
+
     pub(crate) fn from_ast_node(inner: &'a ast::Stmt, is_last: bool) -> Self {
         Stmt { inner, is_last }
     }
@@ -85,23 +100,17 @@ fn rewrite(&self, context: &RewriteContext<'_>, shape: Shape) -> Option<String>
             shape,
             self.as_ast_node(),
             expr_type,
-            Some(self.is_last_expr()),
+            self.is_last_expr(),
         )
     }
 }
 
-impl Rewrite for ast::Stmt {
-    fn rewrite(&self, context: &RewriteContext<'_>, shape: Shape) -> Option<String> {
-        format_stmt(context, shape, self, ExprType::Statement, None)
-    }
-}
-
 fn format_stmt(
     context: &RewriteContext<'_>,
     shape: Shape,
     stmt: &ast::Stmt,
     expr_type: ExprType,
-    is_last_expr: Option<bool>,
+    is_last_expr: bool,
 ) -> Option<String> {
     skip_out_of_file_lines_range!(context, stmt.span());
 
diff --git a/src/utils.rs b/src/utils.rs
index e768ae0a1bc..b8a44d4bade 100644
--- a/src/utils.rs
+++ b/src/utils.rs
@@ -295,7 +295,7 @@ pub(crate) fn semicolon_for_expr(context: &RewriteContext<'_>, expr: &ast::Expr)
 pub(crate) fn semicolon_for_stmt(
     context: &RewriteContext<'_>,
     stmt: &ast::Stmt,
-    is_last_expr: Option<bool>,
+    is_last_expr: bool,
 ) -> bool {
     match stmt.kind {
         ast::StmtKind::Semi(ref expr) => match expr.kind {
@@ -305,7 +305,7 @@ pub(crate) fn semicolon_for_stmt(
             ast::ExprKind::Break(..) | ast::ExprKind::Continue(..) | ast::ExprKind::Ret(..) => {
                 // The only time we can skip the semi-colon is if the config option is set to false
                 // **and** this is the last expr (even though any following exprs are unreachable)
-                context.config.trailing_semicolon() || !is_last_expr.unwrap_or(false)
+                context.config.trailing_semicolon() || !is_last_expr
             }
             _ => true,
         },