Added comment reflow functionality

xls/dslx/fmt/ast_fmt.cc

@@ -88,38 +88,49 @@ std::optional<DocRef> Formatter::FormatCommentsBetween(
 
   std::vector<const CommentData*> items = comments_.GetComments(span);
   VLOG(3) << "Found " << items.size() << " comment data items";
+
   std::optional<Span> previous_comment_span;
+  std::string running_comment = "";
   for (size_t i = 0; i < items.size(); ++i) {
     const CommentData* comment_data = items[i];
     comments_.PlaceComment(comment_data);
 
-    // If the previous comment line and this comment line are abutted (i.e.
-    // contiguous lines with comments), we don't put a newline between them.
+    // If the previous comment line and this comment line are not abutted
+    // (i.e. not contiguous lines with comments), we flush running_comment.
     if (previous_comment_span.has_value() &&
         previous_comment_span->start().lineno() + 1 !=
             comment_data->span.start().lineno()) {
-      VLOG(3) << "previous comment span: "
+      VLOG(3) << "previous block, last comment span: "
               << previous_comment_span.value().ToString(file_table)
-              << " this comment span: "
+              << " new block, first comment span: "
               << comment_data->span.ToString(file_table)
-              << " -- inserting hard line";
+              << " -- flushing previous block, inserting hard line";
+      pieces.push_back(arena_.MakePrefixedReflow(
+          "//",
+          std::string{absl::StripTrailingAsciiWhitespace(running_comment)}));
       pieces.push_back(arena_.hard_line());
-    }
-
-    pieces.push_back(arena_.MakePrefixedReflow(
-        "//",
-        std::string{absl::StripTrailingAsciiWhitespace(comment_data->text)}));
-
-    if (i + 1 != items.size()) {
       pieces.push_back(arena_.hard_line());
+
+      running_comment = "";
     }
 
+    running_comment += comment_data->text;
+
     previous_comment_span = comment_data->span;
     if (last_comment_span != nullptr) {
       *last_comment_span = comment_data->span;
     }
   }
 
+  // If running_comment is not empty, flush it. We don't add a hardline
+  // to maintain the property that emitted comment docs do not have
+  // trailing hardlines.
+  if (running_comment != "") {
+    pieces.push_back(arena_.MakePrefixedReflow(
+        "//",
+        std::string{absl::StripTrailingAsciiWhitespace(running_comment)}));
+  }
+
   if (pieces.empty()) {
     return std::nullopt;
   }

xls/dslx/fmt/ast_fmt_test.cc

@@ -1282,6 +1282,70 @@ TEST_F(FunctionFmtTest, CommentParagraphThenStatement) {
   EXPECT_EQ(got, original);
 }
 
+// A blank line separating two comment paragraphs is preserved.
+TEST_F(FunctionFmtTest, TwoCommentParagraphsKeepBlankLine) {
+  const std::string_view original =
+      R"(fn f() {
+    // paragraph one
+    // still paragraph one
+
+    // paragraph two
+    let x = u32:42;
+})";
+  XLS_ASSERT_OK_AND_ASSIGN(std::string got, DoFmt(original));
+  EXPECT_EQ(got, original);
+}
+
+// A block of abutted comment lines that each fit stays one line
+// per source line.
+TEST_F(FunctionFmtTest, AbuttedCommentBlockPreserved) {
+  const std::string_view original =
+      R"(fn f() {
+    // line one
+    // line two
+    // line three
+    let x = u32:42;
+})";
+  XLS_ASSERT_OK_AND_ASSIGN(std::string got, DoFmt(original));
+  EXPECT_EQ(got, original);
+}
+
+// A single comment line that exceeds the column budget is reflowed across
+// multiple lines, each carrying the "//" prefix.
+TEST_F(FunctionFmtTest, OverlongCommentReflowsToMultipleLines) {
+  const std::string_view original =
+      R"(fn f() {
+    // this is a really long explanatory comment that is going to exceed the hundred column budget for sure yes it will
+    let x = u32:42;
+})";
+  const std::string_view want =
+      R"(fn f() {
+    // this is a really long explanatory comment that is going to exceed the hundred column budget
+    // for sure yes it will
+    let x = u32:42;
+})";
+  XLS_ASSERT_OK_AND_ASSIGN(std::string got, DoFmt(original));
+  EXPECT_EQ(got, want);
+}
+
+// Comments reflow smoothly onto next line when there is space.
+TEST_F(FunctionFmtTest, SecondCommentLineReflowsUpIntoAvailableSpace) {
+  const std::string_view original =
+      R"(fn f() {
+    // one two three four five six seven eight nine ten eleven twelve thirteen fourteen fifteen sixteen seventeen eighteen
+    // nineteen twenty
+    let x = u32:42;
+})";
+  const std::string_view want =
+      R"(fn f() {
+    // one two three four five six seven eight nine ten eleven twelve thirteen fourteen fifteen
+    // sixteen seventeen eighteen nineteen twenty
+    let x = u32:42;
+})";
+  XLS_ASSERT_OK_AND_ASSIGN(std::string got, DoFmt(original));
+  EXPECT_EQ(got, want);
+}
+
 TEST_F(FunctionFmtTest, LetRhsIsOverLongFor) {
   const std::string_view original =
       R"(fn f() {

xls/dslx/fmt/pretty_print.cc

@@ -203,6 +203,9 @@ void PrettyPrintInternal(const DocArena& arena, const Doc& doc,
                   absl::StrSplit(prefixed.text, '\n');
               const std::string& prefix = prefixed.prefix;
 
+              // Determines whether to emit prefix.
+              bool need_prefix = true;
+
               for (size_t i = 0; i < lines.size(); ++i) {
                 std::string_view line = lines[i];
 
@@ -215,10 +218,19 @@ void PrettyPrintInternal(const DocArena& arena, const Doc& doc,
                 if (prefix.size() + line.size() < remaining_cols) {
                   // If it all fits in available cols, place it there in its
                   // entirety.
-                  emit(absl::StrCat(prefix, line));
+                  if (need_prefix) {
+                    emit(absl::StrCat(prefix, line));
+                  } else {
+                    emit(line);
+                  }
+
                   if (i + 1 != lines.size()) {
                     emit_cr(entry.indent());
                   }
+
+                  // Since we are on a new line due to emit_cr, we will need
+                  // a prefix.
+                  need_prefix = true;
                 } else {
                   // Otherwise, place tokens until we encounter EOL and then
                   // wrap. We make sure we put at least one token on each line
@@ -240,16 +252,22 @@ void PrettyPrintInternal(const DocArena& arena, const Doc& doc,
                       absl::MakeConstSpan(toks);
 
                   while (!remaining_toks.empty()) {
-                    emit(prefix);
-                    emit(std::string(leading_whitespace_size, ' '));
+                    if (need_prefix) {
+                      emit(prefix);
+                      emit(std::string(leading_whitespace_size, ' '));
+                    }
 
-                    // After we emit the prefix we make sure we emit at least
-                    // one token.
                     while (!remaining_toks.empty()) {
-                      std::string_view tok = remaining_toks.front();
-                      remaining_toks.remove_prefix(1);
+                      // Emit one guaranteed token only on a fresh line.
+                      if (need_prefix) {
+                        std::string_view tok = remaining_toks.front();
+                        remaining_toks.remove_prefix(1);
 
-                      emit(tok);
+                        emit(tok);
+                      }
+
+                      // We don't need a prefix in the middle of a line.
+                      need_prefix = false;
 
                       if (!remaining_toks.empty()) {
                         // If the next token isn't going to fit we make a
@@ -267,12 +285,19 @@ void PrettyPrintInternal(const DocArena& arena, const Doc& doc,
                                   << " tok width: " << next_tok.size()
                                   << " text width: " << entry.text_width();
                           emit_cr(entry.indent());
+                          need_prefix = true;  // Prefix needed on newline.
                           break;
-                        }
+                        } else {
+                          // If the next token is going to fit we just put a
+                          // space char.
+                          emit(" ");
 
-                        // If the next token is going to fit we just put a
-                        // space char.
-                        emit(" ");
+                          // Emit token.
+                          std::string_view tok = remaining_toks.front();
+                          remaining_toks.remove_prefix(1);
+
+                          emit(tok);
+                        }
                       }
                     }
                   }

xls/dslx/fmt/pretty_print_test.cc

@@ -213,6 +213,36 @@ TEST(PrettyPrintTest, PrefixedReflowCustomSpacingBeforeToken) {
   EXPECT_EQ(PrettyPrint(arena, ref, 40), R"(//     I like this many spaces)");
 }
 
+// A block of abutted comment lines that each fit is preserved one-per-line,
+// with the prefix reintroduced on every line.
+TEST(PrettyPrintTest, PrefixedReflowMultiLineAllFit) {
+  FileTable file_table;
+  DocArena arena(file_table);
+  DocRef ref = arena.MakePrefixedReflow("//", " alpha\n beta\n gamma");
+  EXPECT_EQ(PrettyPrint(arena, ref, 20), R"(// alpha
+// beta
+// gamma)");
+}
+
+// When a line overflows and the next line fits in the remaining columns,
+// the next line reflows onto the wrapped line WITHOUT a second prefix.
+TEST(PrettyPrintTest, PrefixedReflowMultiLineOverflowThenFitContinues) {
+  FileTable file_table;
+  DocArena arena(file_table);
+  DocRef ref = arena.MakePrefixedReflow("//", " aaaa bbbb cccc\n dddd");
+  EXPECT_EQ(PrettyPrint(arena, ref, 16), R"(// aaaa bbbb
+// cccc dddd)");
+}
+
+// A lower line does not reflow upward into a line that did not overflow.
+TEST(PrettyPrintTest, PrefixedReflowMultiLineBoundaryPreserved) {
+  FileTable file_table;
+  DocArena arena(file_table);
+  DocRef ref = arena.MakePrefixedReflow("//", " one\n two");
+  EXPECT_EQ(PrettyPrint(arena, ref, 40), R"(// one
+// two)");
+}
+
 // Scenario where we use NestIfFlatFits and the "on_other_ref" DOES NOT fits
 // inline into the current line so we emit the "on_nested_flat_ref" into the
 // subsequent line (it does fit there).