fix: return + bin ops + calls

Author: 0xrusowsky

crates/fmt/src/lib.rs

@@ -140,9 +140,9 @@ fn format_inner(
     let first_result = format_once(config.clone(), compiler, mk_file);
 
     // If first pass was not successful, return the result
-    // if first_result.is_err() {
-    return first_result;
-    // }
+    if first_result.is_err() {
+        return first_result;
+    }
     let Some(first_formatted) = first_result.ok_ref() else { return first_result };
 
     // Second pass formatting

crates/fmt/src/state/mod.rs

@@ -133,6 +133,7 @@ pub(super) struct State<'sess, 'ast> {
     single_line_stmt: Option<bool>,
     named_call_expr: bool,
     binary_expr: Option<BinOpGroup>,
+    return_bin_expr: bool,
     var_init: bool,
     fn_body: bool,
     call_stack: CallStack,
@@ -220,6 +221,7 @@ impl<'sess> State<'sess, '_> {
             single_line_stmt: None,
             named_call_expr: false,
             binary_expr: None,
+            return_bin_expr: false,
             var_init: false,
             fn_body: false,
             call_stack: CallStack::default(),

crates/fmt/src/state/sol.rs

@@ -1273,9 +1273,11 @@ impl<'ast> State<'_, 'ast> {
                         && call_args.len() == 1
                         && !with_single_call_chain_child
                         && callee_fits;
-                    let without_ind = !all_fits
+                    let without_ind = (!all_fits
                         && (callee_fits || s.call_stack.first().is_some_and(|c| c.is_chained()))
-                        && with_single_call_chain_child;
+                        && with_single_call_chain_child)
+                        // exception for when returning a binary expr that has a call in one of its members
+                        || s.return_bin_expr;
 
                     s.print_call_args(
                         call_args,
@@ -1824,6 +1826,7 @@ impl<'ast> State<'_, 'ast> {
                 let fits_alone = space_left >= expr_size + 1;
 
                 if let Some(expr) = expr {
+                    self.return_bin_expr = matches!(&expr.kind, ast::ExprKind::Binary(..));
                     let is_lit_or_ident =
                         matches!(&expr.kind, ast::ExprKind::Lit(..) | ast::ExprKind::Ident(..));
                     if is_lit_or_ident || (overflows && fits_alone) {
@@ -1848,6 +1851,7 @@ impl<'ast> State<'_, 'ast> {
                     }
                     self.print_expr(expr);
                     self.end();
+                    self.return_bin_expr = false;
                 } else {
                     self.print_word("return");
                 }
@@ -2616,6 +2620,17 @@ fn is_call_chain(expr_kind: &ast::ExprKind<'_>, must_have_child: bool) -> bool {
     }
 }
 
+fn is_call_chain_traverse_bin_ops(expr_kind: &ast::ExprKind<'_>, must_have_child: bool) -> bool {
+    match expr_kind {
+        ast::ExprKind::Binary(lhs, _, rhs) => {
+            is_call_chain_traverse_bin_ops(&lhs.kind, false)
+                || is_call_chain_traverse_bin_ops(&rhs.kind, false)
+        }
+        ast::ExprKind::Member(child, ..) => is_call_chain_traverse_bin_ops(&child.kind, false),
+        _ => !must_have_child && is_call(expr_kind),
+    }
+}
+
 #[derive(Debug)]
 struct Decision {
     outcome: bool,

crates/fmt/testdata/OperatorExpressions/120.fmt.sol

@@ -58,3 +58,11 @@ function test_nested() {
 
     return spender == ownerOf(tokenId) || getApproved[tokenId] == spender || isApprovedForAll[ownerOf(tokenId)][spender];
 }
+
+function new_y(uint256 x, uint256 dx, uint256 x_basis, uint256 y, uint256 y_basis) external pure returns (uint256) {
+    return _get_y(
+        x * _VELODROME_TOKEN_BASIS / x_basis,
+        dx * _VELODROME_TOKEN_BASIS / x_basis,
+        y * _VELODROME_TOKEN_BASIS / y_basis
+    ) * y_basis / _VELODROME_TOKEN_BASIS * aReallyLongIdentifierThatMakesTheOperatorExpressionBreak;
+}

crates/fmt/testdata/OperatorExpressions/fmt.sol

@@ -61,3 +61,18 @@ function test_nested() {
     return spender == ownerOf(tokenId) || getApproved[tokenId] == spender
         || isApprovedForAll[ownerOf(tokenId)][spender];
 }
+
+function new_y(
+    uint256 x,
+    uint256 dx,
+    uint256 x_basis,
+    uint256 y,
+    uint256 y_basis
+) external pure returns (uint256) {
+    return _get_y(
+        x * _VELODROME_TOKEN_BASIS / x_basis,
+        dx * _VELODROME_TOKEN_BASIS / x_basis,
+        y * _VELODROME_TOKEN_BASIS / y_basis
+    ) * y_basis / _VELODROME_TOKEN_BASIS
+        * aReallyLongIdentifierThatMakesTheOperatorExpressionBreak;
+}

crates/fmt/testdata/OperatorExpressions/original.sol

@@ -46,3 +46,15 @@ function test_nested() {
     return spender == ownerOf(tokenId) || getApproved[tokenId] == spender
         || isApprovedForAll[ownerOf(tokenId)][spender];
 }
+
+function new_y(uint256 x, uint256 dx, uint256 x_basis, uint256 y, uint256 y_basis)
+    external
+    pure
+    returns (uint256)
+{
+    return _get_y(
+        x * _VELODROME_TOKEN_BASIS / x_basis,
+        dx * _VELODROME_TOKEN_BASIS / x_basis,
+        y * _VELODROME_TOKEN_BASIS / y_basis
+    ) * y_basis / _VELODROME_TOKEN_BASIS * aReallyLongIdentifierThatMakesTheOperatorExpressionBreak;
+}

crates/fmt/testdata/OperatorExpressions/pow-no-space.fmt.sol

@@ -62,3 +62,18 @@ function test_nested() {
     return spender == ownerOf(tokenId) || getApproved[tokenId] == spender
         || isApprovedForAll[ownerOf(tokenId)][spender];
 }
+
+function new_y(
+    uint256 x,
+    uint256 dx,
+    uint256 x_basis,
+    uint256 y,
+    uint256 y_basis
+) external pure returns (uint256) {
+    return _get_y(
+        x * _VELODROME_TOKEN_BASIS / x_basis,
+        dx * _VELODROME_TOKEN_BASIS / x_basis,
+        y * _VELODROME_TOKEN_BASIS / y_basis
+    ) * y_basis / _VELODROME_TOKEN_BASIS
+        * aReallyLongIdentifierThatMakesTheOperatorExpressionBreak;
+}

crates/fmt/testdata/ReprosCalls/original.sol

@@ -78,15 +78,3 @@ function returnLongBinaryOp() returns (bytes32) {
     return
         bytes32(uint256(Feature.unwrap(feature)) << 128 | uint256(block.chainid) << 64 | uint256(Nonce.unwrap(nonce)));
 }
-
-function new_y(uint256 x, uint256 dx, uint256 x_basis, uint256 y, uint256 y_basis)
-    external
-    pure
-    returns (uint256)
-{
-    return _get_y(
-        x * _VELODROME_TOKEN_BASIS / x_basis,
-        dx * _VELODROME_TOKEN_BASIS / x_basis,
-        y * _VELODROME_TOKEN_BASIS / y_basis
-    );
-}