From 7ca5985061caabf5e3334516db0d7f4942a3e056 Mon Sep 17 00:00:00 2001
From: Chris Leary <leary@google.com>
Date: Wed, 11 Oct 2023 16:45:18 -0700
Subject: [PATCH] [DSLX:fmt] Support for SplatStructInstance.

PiperOrigin-RevId: 572724279
---
 xls/dslx/fmt/ast_fmt.cc      | 63 +++++++++++++++++++++++++++---------
 xls/dslx/fmt/ast_fmt_test.cc | 36 +++++++++++++++++++++
 2 files changed, 83 insertions(+), 16 deletions(-)

diff --git a/xls/dslx/fmt/ast_fmt.cc b/xls/dslx/fmt/ast_fmt.cc
index 5a460b85c4..3ca2ac6a0c 100644
--- a/xls/dslx/fmt/ast_fmt.cc
+++ b/xls/dslx/fmt/ast_fmt.cc
@@ -564,11 +564,6 @@ DocRef Fmt(const XlsTuple& n, const Comments& comments, DocArena& arena) {
   return ConcatNGroup(arena, top);
 }
 
-DocRef Fmt(const SplatStructInstance& n, const Comments& comments,
-           DocArena& arena) {
-  XLS_LOG(FATAL) << "handle splat struct instance: " << n.ToString();
-}
-
 static DocRef Fmt(const StructRef& n, const Comments& comments,
                   DocArena& arena) {
   return absl::visit(
@@ -579,18 +574,23 @@ static DocRef Fmt(const StructRef& n, const Comments& comments,
       n);
 }
 
-DocRef Fmt(const StructInstance& n, const Comments& comments, DocArena& arena) {
-  DocRef leader = ConcatNGroup(arena, {
-                                          Fmt(n.struct_def(), comments, arena),
-                                          arena.break1(),
-                                          arena.ocurl(),
-                                      });
-  if (n.GetUnorderedMembers().empty()) {  // empty struct instance
-    return arena.MakeConcat(leader, arena.ccurl());
-  }
+// Note: this does not put any spacing characters after the '{' so we can
+// appropriately handle the case of an empty struct having no spacing in its
+// `S {}` style construct.
+static DocRef FmtStructLeader(const StructRef& struct_ref,
+                              const Comments& comments, DocArena& arena) {
+  return ConcatNGroup(arena, {
+                                 Fmt(struct_ref, comments, arena),
+                                 arena.break1(),
+                                 arena.ocurl(),
+                             });
+}
 
-  DocRef body_pieces = FmtJoin<std::pair<std::string, Expr*>>(
-      n.GetUnorderedMembers(), Joiner::kCommaBreak1,
+static DocRef FmtStructMembers(
+    absl::Span<const std::pair<std::string, Expr*>> members,
+    const Comments& comments, DocArena& arena) {
+  return FmtJoin<std::pair<std::string, Expr*>>(
+      members, Joiner::kCommaBreak1,
       [](const auto& member, const Comments& comments, DocArena& arena) {
         const auto& [name, expr] = member;
         return ConcatNGroup(
@@ -598,12 +598,43 @@ DocRef Fmt(const StructInstance& n, const Comments& comments, DocArena& arena) {
                     Fmt(*expr, comments, arena)});
       },
       comments, arena);
+}
+
+DocRef Fmt(const StructInstance& n, const Comments& comments, DocArena& arena) {
+  DocRef leader = FmtStructLeader(n.struct_def(), comments, arena);
+
+  if (n.GetUnorderedMembers().empty()) {  // empty struct instance
+    return arena.MakeConcat(leader, arena.ccurl());
+  }
+
+  DocRef body_pieces =
+      FmtStructMembers(n.GetUnorderedMembers(), comments, arena);
 
   return ConcatNGroup(arena,
                       {leader, arena.break1(), arena.MakeNest(body_pieces),
                        arena.break1(), arena.ccurl()});
 }
 
+DocRef Fmt(const SplatStructInstance& n, const Comments& comments,
+           DocArena& arena) {
+  DocRef leader = FmtStructLeader(n.struct_ref(), comments, arena);
+  if (n.members().empty()) {
+    return ConcatNGroup(arena, {leader, arena.break1(), arena.dot_dot(),
+                                Fmt(*n.splatted(), comments, arena),
+                                arena.break1(), arena.ccurl()});
+  }
+
+  DocRef body_pieces = FmtStructMembers(n.members(), comments, arena);
+
+  return ConcatNGroup(
+      arena,
+      {leader, arena.break1(), arena.MakeNest(body_pieces), arena.comma(),
+       arena.break1(), arena.dot_dot(), Fmt(*n.splatted(), comments, arena),
+       arena.break1(), arena.ccurl()});
+
+  XLS_LOG(FATAL) << "handle splat struct instance: " << n.ToString();
+}
+
 DocRef Fmt(const String& n, const Comments& comments, DocArena& arena) {
   return arena.MakeText(n.ToString());
 }
diff --git a/xls/dslx/fmt/ast_fmt_test.cc b/xls/dslx/fmt/ast_fmt_test.cc
index 0d05aa580f..d50b5fd783 100644
--- a/xls/dslx/fmt/ast_fmt_test.cc
+++ b/xls/dslx/fmt/ast_fmt_test.cc
@@ -660,5 +660,41 @@ fn get_x(p: Point) -> u32 { p.x }
   }
 }
 
+TEST(ModuleFmtTest, ConstStructInstanceWithSplatVariantOneUpdate) {
+  const std::string_view kProgram =
+      R"(struct Point { x: u32, y: u32 }
+
+const P = Point { x: u32:42, y: u32:64 };
+
+const Q = Point { x: u32:32, ..P };
+)";
+  std::vector<CommentData> comments;
+  XLS_ASSERT_OK_AND_ASSIGN(std::unique_ptr<Module> m,
+                           ParseModule(kProgram, "fake.x", "fake", &comments));
+
+  {
+    std::string got = AutoFmt(*m, Comments::Create(comments));
+    EXPECT_EQ(got, kProgram);
+  }
+}
+
+TEST(ModuleFmtTest, ConstStructInstanceWithSplatVariantNoUpdate) {
+  const std::string_view kProgram =
+      R"(struct Point { x: u32, y: u32 }
+
+const P = Point { x: u32:42, y: u32:64 };
+
+const Q = Point { ..P };
+)";
+  std::vector<CommentData> comments;
+  XLS_ASSERT_OK_AND_ASSIGN(std::unique_ptr<Module> m,
+                           ParseModule(kProgram, "fake.x", "fake", &comments));
+
+  {
+    std::string got = AutoFmt(*m, Comments::Create(comments));
+    EXPECT_EQ(got, kProgram);
+  }
+}
+
 }  // namespace
 }  // namespace xls::dslx