Implement a custom stream for LDBG macro to handle newlines (#150750)

This prints the prefix on every new line, allowing for an output that
looks like:
```
[dead-code-analysis] DeadCodeAnalysis.cpp:288 Visiting operation: func.func private @private_1() -> (i32, i32) {
[dead-code-analysis] DeadCodeAnalysis.cpp:288   %c0_i32 = arith.constant 0 : i32
[dead-code-analysis] DeadCodeAnalysis.cpp:288   %0 = arith.addi %c0_i32, %c0_i32 {tag = "one"} : i32
[dead-code-analysis] DeadCodeAnalysis.cpp:288   return %c0_i32, %0 : i32, i32
[dead-code-analysis] DeadCodeAnalysis.cpp:288 }
[dead-code-analysis] DeadCodeAnalysis.cpp:313 Visiting callable operation: func.func private @private_1() -> (i32, i32) {
[dead-code-analysis] DeadCodeAnalysis.cpp:313   %c0_i32 = arith.constant 0 : i32
[dead-code-analysis] DeadCodeAnalysis.cpp:313   %0 = arith.addi %c0_i32, %c0_i32 {tag = "one"} : i32
[dead-code-analysis] DeadCodeAnalysis.cpp:313   return %c0_i32, %0 : i32, i32
[dead-code-analysis] DeadCodeAnalysis.cpp:313 }
```

2 files changed, 109 insertions, 35 deletions

diff --git a/llvm/include/llvm/Support/DebugLog.h b/llvm/include/llvm/Support/DebugLog.h
index b1b17e3..19d3098 100644
--- a/llvm/include/llvm/Support/DebugLog.h
+++ b/llvm/include/llvm/Support/DebugLog.h
@@ -26,55 +26,107 @@ namespace llvm {
 //              << "] " << "Bitset contains: " << Bitset << "\n");
 #define LDBG() DEBUGLOG_WITH_STREAM_AND_TYPE(llvm::dbgs(), DEBUG_TYPE)
 
-#if defined(__SHORT_FILE__)
-#define DEBUGLOG_WITH_STREAM_AND_TYPE(STREAM, TYPE)                            \
+#define DEBUGLOG_WITH_STREAM_TYPE_AND_FILE(STREAM, TYPE, FILE)                 \
   for (bool _c = (::llvm::DebugFlag && ::llvm::isCurrentDebugType(TYPE)); _c;  \
        _c = false)                                                             \
-  ::llvm::impl::LogWithNewline(TYPE, __SHORT_FILE__, __LINE__, (STREAM))
+  ::llvm::impl::raw_ldbg_ostream{                                              \
+      ::llvm::impl::computePrefix(TYPE, FILE, __LINE__), (STREAM)}             \
+      .asLvalue()
+// When __SHORT_FILE__ is not defined, the File is the full path,
+// otherwise __SHORT_FILE__ is defined in CMake to provide the file name
+// without the path prefix.
+#if defined(__SHORT_FILE__)
+#define DEBUGLOG_WITH_STREAM_AND_TYPE(STREAM, TYPE)                            \
+  DEBUGLOG_WITH_STREAM_TYPE_AND_FILE(STREAM, TYPE, __SHORT_FILE__)
 #else
 #define DEBUGLOG_WITH_STREAM_AND_TYPE(STREAM, TYPE)                            \
-  for (bool _c = (::llvm::DebugFlag && ::llvm::isCurrentDebugType(TYPE)); _c;  \
-       _c = false)                                                             \
-  ::llvm::impl::LogWithNewline(TYPE, __FILE__, __LINE__, (STREAM))
+  DEBUGLOG_WITH_STREAM_TYPE_AND_FILE(                                          \
+      STREAM, TYPE, ::llvm::impl::LogWithNewline::getShortFileName(__FILE__))
 #endif
 
 namespace impl {
-class LogWithNewline {
-public:
-  LogWithNewline(const char *debug_type, const char *file, int line,
-                 raw_ostream &os)
-      : os(os) {
-#if !defined(__SHORT_FILE__)
-    file = ::llvm::impl::LogWithNewline::getShortFileName(file);
-#endif
-    if (debug_type)
-      os << "[" << debug_type << "] ";
-    os << file << ":" << line << " ";
+
+/// A raw_ostream that tracks `\n` and print the prefix.
+class LLVM_ABI raw_ldbg_ostream final : public raw_ostream {
+  std::string Prefix;
+  raw_ostream &Os;
+  bool HasPendingNewline = true;
+
+  /// Split the line on newlines and insert the prefix before each newline.
+  /// Forward everything to the underlying stream.
+  void write_impl(const char *Ptr, size_t Size) final {
+    auto Str = StringRef(Ptr, Size);
+    // Handle the initial prefix.
+    if (!Str.empty())
+      writeWithPrefix(StringRef());
+
+    auto Eol = Str.find('\n');
+    while (Eol != StringRef::npos) {
+      StringRef Line = Str.take_front(Eol + 1);
+      if (!Line.empty())
+        writeWithPrefix(Line);
+      HasPendingNewline = true;
+      Str = Str.drop_front(Eol + 1);
+      Eol = Str.find('\n');
+    }
+    if (!Str.empty())
+      writeWithPrefix(Str);
   }
-  ~LogWithNewline() { os << '\n'; }
-  template <typename T> raw_ostream &operator<<(const T &t) && {
-    return os << t;
+  void emitPrefix() { Os.write(Prefix.c_str(), Prefix.size()); }
+  void writeWithPrefix(StringRef Str) {
+    if (HasPendingNewline) {
+      emitPrefix();
+      HasPendingNewline = false;
+    }
+    Os.write(Str.data(), Str.size());
   }
 
-  // Prevent copying, as this class manages newline responsibility and is
-  // intended for use as a temporary.
-  LogWithNewline(const LogWithNewline &) = delete;
-  LogWithNewline &operator=(const LogWithNewline &) = delete;
-  LogWithNewline &operator=(LogWithNewline &&) = delete;
-  static constexpr const char *getShortFileName(const char *path) {
-    // Remove the path prefix from the file name.
-    const char *filename = path;
-    for (const char *p = path; *p != '\0'; ++p) {
-      if (*p == '/' || *p == '\\') {
-        filename = p + 1;
-      }
+public:
+  explicit raw_ldbg_ostream(std::string Prefix, raw_ostream &Os)
+      : Prefix(std::move(Prefix)), Os(Os) {
+    SetUnbuffered();
+  }
+  ~raw_ldbg_ostream() final {
+    flushEol();
+    Os << '\n';
+  }
+  void flushEol() {
+    if (HasPendingNewline) {
+      emitPrefix();
+      HasPendingNewline = false;
     }
-    return filename;
   }
 
-private:
-  raw_ostream &os;
+  /// Forward the current_pos method to the underlying stream.
+  uint64_t current_pos() const final { return Os.tell(); }
+
+  /// Some of the `<<` operators expect an lvalue, so we trick the type system.
+  raw_ldbg_ostream &asLvalue() { return *this; }
 };
+
+/// Remove the path prefix from the file name.
+static LLVM_ATTRIBUTE_UNUSED constexpr const char *
+getShortFileName(const char *path) {
+  const char *filename = path;
+  for (const char *p = path; *p != '\0'; ++p) {
+    if (*p == '/' || *p == '\\')
+      filename = p + 1;
+  }
+  return filename;
+}
+
+/// Compute the prefix for the debug log in the form of:
+/// "[DebugType] File:Line "
+/// Where the File is the file name without the path prefix.
+static LLVM_ATTRIBUTE_UNUSED std::string
+computePrefix(const char *DebugType, const char *File, int Line) {
+  std::string Prefix;
+  raw_string_ostream OsPrefix(Prefix);
+  if (DebugType)
+    OsPrefix << "[" << DebugType << "] ";
+  OsPrefix << File << ":" << Line << " ";
+  return OsPrefix.str();
+}
 } // end namespace impl
 #else
 // As others in Debug, When compiling without assertions, the -debug-* options
diff --git a/llvm/unittests/Support/DebugLogTest.cpp b/llvm/unittests/Support/DebugLogTest.cpp
index b26fa40..c34d888 100644
--- a/llvm/unittests/Support/DebugLogTest.cpp
+++ b/llvm/unittests/Support/DebugLogTest.cpp
@@ -62,6 +62,28 @@ TEST(DebugLogTest, Basic) {
     EXPECT_THAT(count, Eq(1));
   }
 }
+
+TEST(DebugLogTest, StreamPrefix) {
+  llvm::DebugFlag = true;
+  static const char *DT[] = {"A", "B"};
+  setCurrentDebugTypes(DT, 2);
+
+  std::string str;
+  raw_string_ostream os(str);
+  std::string expected = "PrefixA 1\nPrefixA 2\nPrefixA \nPrefixB "
+                         "3\nPrefixB 4\nPrefixA 5";
+  {
+    llvm::impl::raw_ldbg_ostream ldbg_osB("PrefixB ", os);
+    llvm::impl::raw_ldbg_ostream ldbg_osA("PrefixA ", os);
+    ldbg_osA << "1\n2";
+    ldbg_osA << "\n\n";
+    ldbg_osB << "3\n4\n";
+    ldbg_osA << "5";
+    EXPECT_EQ(os.str(), expected);
+  }
+  // After destructors, there was a pending newline for stream B.
+  EXPECT_EQ(os.str(), expected + "\nPrefixB \n");
+}
 #else
 TEST(DebugLogTest, Basic) {
   // LDBG should be compiled out in NDEBUG, so just check it compiles and has