fix(fmt): preserve indexed callee when it fits (#12270)

* fix(fmt): preserve indexed callee when it fits

* style: typos

Author: 0xrusowsky

crates/fmt/src/state/mod.rs

@@ -126,6 +126,8 @@ pub(super) struct State<'sess, 'ast> {
     return_bin_expr: bool,
     // Whether inside a call with call options and at least one argument.
     call_with_opts_and_args: bool,
+    // Whether to skip the index soft breaks because the callee fits inline.
+    skip_index_break: bool,
     // Whether inside an `emit` or `revert` call with a qualified path, or not.
     emit_or_revert: bool,
     // Whether inside a variable initialization expression, or not.
@@ -219,6 +221,7 @@ impl<'sess> State<'sess, '_> {
             contract: None,
             single_line_stmt: None,
             call_with_opts_and_args: false,
+            skip_index_break: false,
             binary_expr: None,
             return_bin_expr: false,
             emit_or_revert: false,

crates/fmt/src/state/sol.rs

@@ -1274,11 +1274,10 @@ impl<'ast> State<'_, 'ast> {
                     MemberOrCallArgs::Member(self.estimate_size(ident.span)),
                     |s| {
                         s.print_trailing_comment(member_expr.span.hi(), Some(ident.span.lo()));
-                        if !matches!(
-                            member_expr.kind,
-                            ast::ExprKind::Ident(_) | ast::ExprKind::Type(_)
-                        ) {
-                            s.zerobreak();
+                        match member_expr.kind {
+                            ast::ExprKind::Ident(_) | ast::ExprKind::Type(_) => (),
+                            ast::ExprKind::Index(..) if s.skip_index_break => (),
+                            _ => s.zerobreak(),
                         }
                         s.word(".");
                         s.print_ident(ident);
@@ -1442,10 +1441,16 @@ impl<'ast> State<'_, 'ast> {
         self.s.cbox(self.ind);
 
         let mut skip_break = false;
-
+        let mut zerobreak = |this: &mut Self| {
+            if this.skip_index_break {
+                skip_break = true;
+            } else {
+                this.zerobreak();
+            }
+        };
         match kind {
             ast::IndexKind::Index(Some(inner_expr)) => {
-                self.zerobreak();
+                zerobreak(self);
                 self.print_expr(inner_expr);
             }
             ast::IndexKind::Index(None) => {}
@@ -1455,19 +1460,19 @@ impl<'ast> State<'_, 'ast> {
                         .print_comments(start_expr.span.lo(), CommentConfig::skip_ws())
                         .is_none_or(|s| s.is_mixed())
                     {
-                        self.zerobreak();
+                        zerobreak(self);
                     }
                     self.print_expr(start_expr);
                 } else {
-                    self.zerobreak();
+                    zerobreak(self);
                 }
 
                 self.word(":");
 
                 if let Some(end_expr) = end {
                     self.s.ibox(self.ind);
                     if start.is_some() {
-                        self.zerobreak();
+                        zerobreak(self);
                     }
                     self.print_comments(
                         end_expr.span.lo(),
@@ -1596,7 +1601,7 @@ impl<'ast> State<'_, 'ast> {
             }
         }
 
-        let mut extra_box = false;
+        let (mut extra_box, skip_cache) = (false, self.skip_index_break);
         let parent_is_chain = self.call_stack.last().copied().is_some_and(|call| call.is_chained());
         if !parent_is_chain {
             // Estimate sizes of callee and optional member
@@ -1626,6 +1631,7 @@ impl<'ast> State<'_, 'ast> {
                     // calls with cmnts between the args always break
                     || (total_fits_line && !member_or_args.has_comments()))
             {
+                self.skip_index_break = true;
                 self.cbox(0);
             } else {
                 self.s.ibox(self.ind);
@@ -1650,6 +1656,11 @@ impl<'ast> State<'_, 'ast> {
             }
             self.end();
         }
+
+        // Restore cache
+        if self.skip_index_break {
+            self.skip_index_break = skip_cache;
+        }
     }
 
     fn print_call_args(
@@ -2859,6 +2870,16 @@ pub(super) fn get_callee_head_size(callee: &ast::Expr<'_>) -> usize {
         ast::ExprKind::Type(ast::Type { kind: ast::TypeKind::Elementary(ty), .. }) => {
             ty.to_abi_str().len()
         }
+        ast::ExprKind::Index(base, idx) => {
+            let idx_len = match idx {
+                ast::IndexKind::Index(expr) => expr.as_ref().map_or(0, |e| get_callee_head_size(e)),
+                ast::IndexKind::Range(e1, e2) => {
+                    1 + e1.as_ref().map_or(0, |e| get_callee_head_size(e))
+                        + e2.as_ref().map_or(0, |e| get_callee_head_size(e))
+                }
+            };
+            get_callee_head_size(base) + 2 + idx_len
+        }
         ast::ExprKind::Member(base, member_ident) => {
             match &base.kind {
                 ast::ExprKind::Ident(..) | ast::ExprKind::Type(..) => {

crates/fmt/testdata/VariableAssignment/bracket-spacing.fmt.sol

@@ -38,4 +38,21 @@ contract TestContract {
             "0x0000000000000000000000000000000000000000000000000000000000000000,"
             "0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF";
     }
+
+    function test_longIndexedCall() {
+        // https://github.com/foundry-rs/foundry/issues/12254
+        bytes memory message = mailboxes[destinationDomain].buildMessage(
+            originDomain,
+            bytes32(0),
+            address(inbox).toBytes32(),
+            abi.encode(orderId, bytes32(0), address(0))
+        );
+        // should have identicall behavior when call of the same size without indexing
+        bytes memory message = mailboxes_destinationDomains.buildMessage(
+            originDomain,
+            bytes32(0),
+            address(inbox).toBytes32(),
+            abi.encode(orderId, bytes32(0), address(0))
+        );
+    }
 }

crates/fmt/testdata/VariableAssignment/fmt.sol

@@ -37,4 +37,21 @@ contract TestContract {
             "0x0000000000000000000000000000000000000000000000000000000000000000,"
             "0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF";
     }
+
+    function test_longIndexedCall() {
+        // https://github.com/foundry-rs/foundry/issues/12254
+        bytes memory message = mailboxes[destinationDomain].buildMessage(
+            originDomain,
+            bytes32(0),
+            address(inbox).toBytes32(),
+            abi.encode(orderId, bytes32(0), address(0))
+        );
+        // should have identicall behavior when call of the same size without indexing
+        bytes memory message = mailboxes_destinationDomains.buildMessage(
+            originDomain,
+            bytes32(0),
+            address(inbox).toBytes32(),
+            abi.encode(orderId, bytes32(0), address(0))
+        );
+    }
 }

crates/fmt/testdata/VariableAssignment/original.sol

@@ -36,4 +36,11 @@ contract TestContract {
             "0x0000000000000000000000000000000000000000000000000000000000000000,"
             "0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFF";
     }
+
+    function test_longIndexedCall() {
+        // https://github.com/foundry-rs/foundry/issues/12254
+        bytes memory message = mailboxes[destinationDomain].buildMessage(originDomain, bytes32(0), address(inbox).toBytes32(), abi.encode(orderId, bytes32(0), address(0)));
+        // should have identicall behavior when call of the same size without indexing
+        bytes memory message = mailboxes_destinationDomains.buildMessage(originDomain, bytes32(0), address(inbox).toBytes32(), abi.encode(orderId, bytes32(0), address(0)));
+    }
 }