[dslx:fmt] Better formatting of a long struct field expression.

xls/dslx/fmt/ast_fmt.cc

@@ -52,6 +52,9 @@
 namespace xls::dslx {
 namespace {
 
+DocRef FmtBlockedExprLeader(const Expr& e, const Comments& comments,
+                            DocArena& arena);
+
 // Note: if a comment doc is emitted (i.e. return value has_value()) it does not
 // have a trailing hard-line. This is for consistency with other emission
 // routines which generally don't emit any whitespace afterwards, just their
@@ -1234,9 +1237,34 @@ static DocRef FmtStructMembersBreak(
           return arena.MakeText(name);
         }
 
-        return ConcatNGroup(arena,
-                            {arena.MakeText(name), arena.colon(), arena.space(),
-                             arena.MakeGroup(Fmt(*expr, comments, arena))});
+        DocRef field_expr = Fmt(*expr, comments, arena);
+
+        // This is the document we want to emit both when we:
+        // - Know it fits in flat mode
+        // - Know the start of the document (i.e. the leader on the RHS
+        //   expression) can be emitted in flat mode
+        //
+        // That's why it has a `break1` in it (instead of a space) and a
+        // reassessment of whether to enter break mode for the field
+        // expression.
+        DocRef on_flat =
+            ConcatN(arena, {arena.MakeText(name), arena.colon(), arena.break1(),
+                            arena.MakeGroup(field_expr)});
+        DocRef nest_field_expr =
+            ConcatN(arena, {arena.MakeText(name), arena.colon(),
+                            arena.hard_line(), arena.MakeNest(field_expr)});
+
+        DocRef on_other;
+        if (expr->IsBlockedExprWithLeader()) {
+          DocRef leader = arena.MakeConcat(
+              arena.space(), FmtBlockedExprLeader(*expr, comments, arena));
+          on_other = arena.MakeModeSelect(leader, /*on_flat=*/on_flat,
+                                          /*on_break=*/nest_field_expr);
+        } else {
+          on_other = arena.MakeFlatChoice(on_flat, nest_field_expr);
+        }
+
+        return arena.MakeNestIfFlatFits(on_flat, on_other);
       },
       comments, arena);
 }

xls/dslx/fmt/ast_fmt_test.cc

@@ -2228,5 +2228,22 @@ pub fn f(a: A) -> A { if a.B[0].value == u32:0 { zero!<A>() } else { a } }
 )");
 }
 
+TEST_F(ModuleFmtTest, LongStructInstanceFieldExpr) {
+  Run(R"(struct X { xxxxxxxxxxxxxxxxxxx: bool, yyyyyyyyyyyyyyyyyyy: u32 }
+
+const aaaaaaaaaaaaaaaaaaaaaaaaa = u32:0;
+const bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb = u32:0;
+const ccccccccccccccccccccccccc = bool:false;
+
+fn f() -> X {
+    X {
+        xxxxxxxxxxxxxxxxxxx:
+            aaaaaaaaaaaaaaaaaaaaaaaaa == bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb,
+        yyyyyyyyyyyyyyyyyyy: u32:0,
+    }
+}
+)");
+}
+
 }  // namespace
 }  // namespace xls::dslx