refactor: maintain more AST info when formatting a RHS

src/expr.rs

@@ -196,9 +196,10 @@ pub(crate) fn format_expr(
                 capture, is_async, movability, fn_decl, body, expr.span, context, shape,
             )
         }
-        ast::ExprKind::Try(..) | ast::ExprKind::Field(..) | ast::ExprKind::MethodCall(..) => {
-            rewrite_chain(expr, context, shape)
-        }
+        ast::ExprKind::Try(..)
+        | ast::ExprKind::Field(..)
+        | ast::ExprKind::MethodCall(..)
+        | ast::ExprKind::Await(_) => rewrite_chain(expr, context, shape),
         ast::ExprKind::MacCall(ref mac) => {
             rewrite_macro(mac, None, context, shape, MacroPosition::Expression).or_else(|| {
                 wrap_str(
@@ -377,7 +378,6 @@ pub(crate) fn format_expr(
                 ))
             }
         }
-        ast::ExprKind::Await(_) => rewrite_chain(expr, context, shape),
         ast::ExprKind::Underscore => Some("_".to_owned()),
         ast::ExprKind::Err => None,
     };
@@ -829,6 +829,7 @@ impl<'a> ControlFlow<'a> {
                 &format!("{}{}{}", matcher, pat_string, self.connector),
                 expr,
                 cond_shape,
+                &RhsAssignKind::Expr(&expr.kind, expr.span),
                 RhsTactics::Default,
                 comments_span,
                 true,
@@ -1839,6 +1840,34 @@ fn rewrite_unary_op(
     rewrite_unary_prefix(context, ast::UnOp::to_string(op), expr, shape)
 }
 
+pub(crate) enum RhsAssignKind<'ast> {
+    Expr(&'ast ast::ExprKind, Span),
+    Bounds,
+    Ty,
+}
+
+impl<'ast> RhsAssignKind<'ast> {
+    // TODO(calebcartwright)
+    // Preemptive addition for handling RHS with chains, not yet utilized.
+    // It may make more sense to construct the chain first and then check
+    // whether there are actually chain elements.
+    #[allow(dead_code)]
+    fn is_chain(&self) -> bool {
+        match self {
+            RhsAssignKind::Expr(kind, _) => {
+                matches!(
+                    kind,
+                    ast::ExprKind::Try(..)
+                        | ast::ExprKind::Field(..)
+                        | ast::ExprKind::MethodCall(..)
+                        | ast::ExprKind::Await(_)
+                )
+            }
+            _ => false,
+        }
+    }
+}
+
 fn rewrite_assignment(
     context: &RewriteContext<'_>,
     lhs: &ast::Expr,
@@ -1855,7 +1884,13 @@ fn rewrite_assignment(
     let lhs_shape = shape.sub_width(operator_str.len() + 1)?;
     let lhs_str = format!("{} {}", lhs.rewrite(context, lhs_shape)?, operator_str);
 
-    rewrite_assign_rhs(context, lhs_str, rhs, shape)
+    rewrite_assign_rhs(
+        context,
+        lhs_str,
+        rhs,
+        &RhsAssignKind::Expr(&rhs.kind, rhs.span),
+        shape,
+    )
 }
 
 /// Controls where to put the rhs.
@@ -1876,16 +1911,18 @@ pub(crate) fn rewrite_assign_rhs<S: Into<String>, R: Rewrite>(
     context: &RewriteContext<'_>,
     lhs: S,
     ex: &R,
+    rhs_kind: &RhsAssignKind<'_>,
     shape: Shape,
 ) -> Option<String> {
-    rewrite_assign_rhs_with(context, lhs, ex, shape, RhsTactics::Default)
+    rewrite_assign_rhs_with(context, lhs, ex, shape, rhs_kind, RhsTactics::Default)
 }
 
 pub(crate) fn rewrite_assign_rhs_expr<R: Rewrite>(
     context: &RewriteContext<'_>,
     lhs: &str,
     ex: &R,
     shape: Shape,
+    rhs_kind: &RhsAssignKind<'_>,
     rhs_tactics: RhsTactics,
 ) -> Option<String> {
     let last_line_width = last_line_width(lhs).saturating_sub(if lhs.contains('\n') {
@@ -1910,6 +1947,7 @@ pub(crate) fn rewrite_assign_rhs_expr<R: Rewrite>(
         ex,
         orig_shape,
         ex.rewrite(context, orig_shape),
+        rhs_kind,
         rhs_tactics,
         has_rhs_comment,
     )
@@ -1920,10 +1958,11 @@ pub(crate) fn rewrite_assign_rhs_with<S: Into<String>, R: Rewrite>(
     lhs: S,
     ex: &R,
     shape: Shape,
+    rhs_kind: &RhsAssignKind<'_>,
     rhs_tactics: RhsTactics,
 ) -> Option<String> {
     let lhs = lhs.into();
-    let rhs = rewrite_assign_rhs_expr(context, &lhs, ex, shape, rhs_tactics)?;
+    let rhs = rewrite_assign_rhs_expr(context, &lhs, ex, shape, rhs_kind, rhs_tactics)?;
     Some(lhs + &rhs)
 }
 
@@ -1932,6 +1971,7 @@ pub(crate) fn rewrite_assign_rhs_with_comments<S: Into<String>, R: Rewrite>(
     lhs: S,
     ex: &R,
     shape: Shape,
+    rhs_kind: &RhsAssignKind<'_>,
     rhs_tactics: RhsTactics,
     between_span: Span,
     allow_extend: bool,
@@ -1943,7 +1983,7 @@ pub(crate) fn rewrite_assign_rhs_with_comments<S: Into<String>, R: Rewrite>(
     } else {
         shape
     };
-    let rhs = rewrite_assign_rhs_expr(context, &lhs, ex, shape, rhs_tactics)?;
+    let rhs = rewrite_assign_rhs_expr(context, &lhs, ex, shape, rhs_kind, rhs_tactics)?;
 
     if contains_comment {
         let rhs = rhs.trim_start();
@@ -1958,6 +1998,7 @@ fn choose_rhs<R: Rewrite>(
     expr: &R,
     shape: Shape,
     orig_rhs: Option<String>,
+    _rhs_kind: &RhsAssignKind<'_>,
     rhs_tactics: RhsTactics,
     has_rhs_comment: bool,
 ) -> Option<String> {

src/items.rs

@@ -18,7 +18,7 @@ use crate::config::lists::*;
 use crate::config::{BraceStyle, Config, IndentStyle, Version};
 use crate::expr::{
     is_empty_block, is_simple_block_stmt, rewrite_assign_rhs, rewrite_assign_rhs_with,
-    rewrite_assign_rhs_with_comments, RhsTactics,
+    rewrite_assign_rhs_with_comments, RhsAssignKind, RhsTactics,
 };
 use crate::lists::{definitive_tactic, itemize_list, write_list, ListFormatting, Separator};
 use crate::macros::{rewrite_macro, MacroPosition};
@@ -116,7 +116,13 @@ impl Rewrite for ast::Local {
             // 1 = trailing semicolon;
             let nested_shape = shape.sub_width(1)?;
 
-            result = rewrite_assign_rhs(context, result, init, nested_shape)?;
+            result = rewrite_assign_rhs(
+                context,
+                result,
+                init,
+                &RhsAssignKind::Expr(&init.kind, init.span),
+                nested_shape,
+            )?;
             // todo else
         }
 
@@ -564,11 +570,13 @@ impl<'a> FmtVisitor<'a> {
 
         let variant_body = if let Some(ref expr) = field.disr_expr {
             let lhs = format!("{:1$} =", variant_body, pad_discrim_ident_to);
+            let ex = &*expr.value;
             rewrite_assign_rhs_with(
                 &context,
                 lhs,
-                &*expr.value,
+                ex,
                 shape,
+                &RhsAssignKind::Expr(&ex.kind, ex.span),
                 RhsTactics::AllowOverflow,
             )?
         } else {
@@ -1033,6 +1041,7 @@ pub(crate) fn format_trait(
                 result + ":",
                 bounds,
                 shape,
+                &RhsAssignKind::Bounds,
                 RhsTactics::ForceNextLineWithoutIndent,
             )?;
         }
@@ -1213,7 +1222,14 @@ pub(crate) fn format_trait_alias(
         generic_bounds,
         generics,
     };
-    rewrite_assign_rhs(context, lhs, &trait_alias_bounds, shape.sub_width(1)?).map(|s| s + ";")
+    rewrite_assign_rhs(
+        context,
+        lhs,
+        &trait_alias_bounds,
+        &RhsAssignKind::Bounds,
+        shape.sub_width(1)?,
+    )
+    .map(|s| s + ";")
 }
 
 fn format_unit_struct(
@@ -1630,7 +1646,7 @@ fn rewrite_ty<R: Rewrite>(
 
         // 1 = `;`
         let shape = Shape::indented(indent, context.config).sub_width(1)?;
-        rewrite_assign_rhs(context, lhs, &*ty, shape).map(|s| s + ";")
+        rewrite_assign_rhs(context, lhs, &*ty, &RhsAssignKind::Ty, shape).map(|s| s + ";")
     } else {
         Some(format!("{};", result))
     }
@@ -1720,7 +1736,7 @@ pub(crate) fn rewrite_struct_field(
 
     let is_prefix_empty = prefix.is_empty();
     // We must use multiline. We are going to put attributes and a field on different lines.
-    let field_str = rewrite_assign_rhs(context, prefix, &*field.ty, shape)?;
+    let field_str = rewrite_assign_rhs(context, prefix, &*field.ty, &RhsAssignKind::Ty, shape)?;
     // Remove a leading white-space from `rewrite_assign_rhs()` when rewriting a tuple struct.
     let field_str = if is_prefix_empty {
         field_str.trim_start()
@@ -1850,6 +1866,7 @@ fn rewrite_static(
             &lhs,
             &**expr,
             Shape::legacy(remaining_width, offset.block_only()),
+            &RhsAssignKind::Expr(&expr.kind, expr.span),
             RhsTactics::Default,
             comments_span,
             true,
@@ -3147,7 +3164,14 @@ impl Rewrite for ast::ForeignItem {
                     rewrite_ident(context, self.ident)
                 );
                 // 1 = ;
-                rewrite_assign_rhs(context, prefix, &**ty, shape.sub_width(1)?).map(|s| s + ";")
+                rewrite_assign_rhs(
+                    context,
+                    prefix,
+                    &**ty,
+                    &RhsAssignKind::Ty,
+                    shape.sub_width(1)?,
+                )
+                .map(|s| s + ";")
             }
             ast::ForeignItemKind::TyAlias(ref ty_alias) => {
                 let (kind, span) = (&ItemVisitorKind::ForeignItem(&self), self.span);

src/macros.rs

@@ -27,7 +27,7 @@ use crate::comment::{
     contains_comment, CharClasses, FindUncommented, FullCodeCharKind, LineClasses,
 };
 use crate::config::lists::*;
-use crate::expr::rewrite_array;
+use crate::expr::{rewrite_array, rewrite_assign_rhs, RhsAssignKind};
 use crate::lists::{itemize_list, write_list, ListFormatting};
 use crate::overflow;
 use crate::rewrite::{Rewrite, RewriteContext};
@@ -1468,10 +1468,11 @@ fn format_lazy_static(
             id,
             ty.rewrite(context, nested_shape)?
         ));
-        result.push_str(&crate::expr::rewrite_assign_rhs(
+        result.push_str(&rewrite_assign_rhs(
             context,
             stmt,
             &*expr,
+            &RhsAssignKind::Expr(&expr.kind, expr.span),
             nested_shape.sub_width(1)?,
         )?);
         result.push(';');

src/types.rs

@@ -10,6 +10,7 @@ use crate::config::lists::*;
 use crate::config::{IndentStyle, TypeDensity, Version};
 use crate::expr::{
     format_expr, rewrite_assign_rhs, rewrite_call, rewrite_tuple, rewrite_unary_prefix, ExprType,
+    RhsAssignKind,
 };
 use crate::lists::{
     definitive_tactic, itemize_list, write_list, ListFormatting, ListItem, Separator,
@@ -430,7 +431,7 @@ impl Rewrite for ast::WherePredicate {
                     format!("{}{}", type_str, colon)
                 };
 
-                rewrite_assign_rhs(context, lhs, bounds, shape)?
+                rewrite_assign_rhs(context, lhs, bounds, &RhsAssignKind::Bounds, shape)?
             }
             ast::WherePredicate::RegionPredicate(ast::WhereRegionPredicate {
                 ref lifetime,
@@ -443,7 +444,7 @@ impl Rewrite for ast::WherePredicate {
                 ..
             }) => {
                 let lhs_ty_str = lhs_ty.rewrite(context, shape).map(|lhs| lhs + " =")?;
-                rewrite_assign_rhs(context, lhs_ty_str, &**rhs_ty, shape)?
+                rewrite_assign_rhs(context, lhs_ty_str, &**rhs_ty, &RhsAssignKind::Ty, shape)?
             }
         };