From af1660fa85a5bfa393bcbbf5d9741c594a9a3499 Mon Sep 17 00:00:00 2001
From: "Christopher D. Leary" <cdleary@openai.com>
Date: Sat, 5 Oct 2024 17:23:35 -0700
Subject: [PATCH 1/2] [dslx:fmt] Better formatting of a long struct field
 expression.

Fixes google/xls#1633
---
 xls/dslx/fmt/ast_fmt.cc      | 26 +++++++++++++++++++++++---
 xls/dslx/fmt/ast_fmt_test.cc | 17 +++++++++++++++++
 2 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/xls/dslx/fmt/ast_fmt.cc b/xls/dslx/fmt/ast_fmt.cc
index e68e1f0c40..f928f245cc 100644
--- a/xls/dslx/fmt/ast_fmt.cc
+++ b/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,26 @@ 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);
+
+        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);
 }
diff --git a/xls/dslx/fmt/ast_fmt_test.cc b/xls/dslx/fmt/ast_fmt_test.cc
index 952852d49a..6308e4c960 100644
--- a/xls/dslx/fmt/ast_fmt_test.cc
+++ b/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

From 1b2eabe6033f3db20c19ad921fc6e39f710c9a37 Mon Sep 17 00:00:00 2001
From: "Christopher D. Leary" <cdleary@openai.com>
Date: Sat, 5 Oct 2024 17:35:08 -0700
Subject: [PATCH 2/2] Add comment.

---
 xls/dslx/fmt/ast_fmt.cc | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/xls/dslx/fmt/ast_fmt.cc b/xls/dslx/fmt/ast_fmt.cc
index f928f245cc..0434ba27ce 100644
--- a/xls/dslx/fmt/ast_fmt.cc
+++ b/xls/dslx/fmt/ast_fmt.cc
@@ -1239,6 +1239,14 @@ static DocRef FmtStructMembersBreak(
 
         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)});