[llvm] Implement a custom stream for LDBG macro to handle newlines (PR #150750)

[Mehdi Amini via llvm-commits]

https://github.com/joker-eph updated https://github.com/llvm/llvm-project/pull/150750

>From 54c4f13db4237f50eeb21793489889fd74384338 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Sat, 26 Jul 2025 03:51:17 -0700
Subject: [PATCH] Implement a custom stream for LDBG macro to handle newlines

This prints the prefix on every new line, allowing for an output that looks like:

[dead-code-analysis] DeadCodeAnalysis.cpp:284 Visiting program point: 0x55ea5be2d4c8 <after operation>:func.func private @private_1() -> (i32, i32) {...}
[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 }
---
 llvm/include/llvm/Support/DebugLog.h | 80 +++++++++++++++++++++-------
 1 file changed, 62 insertions(+), 18 deletions(-)

diff --git a/llvm/include/llvm/Support/DebugLog.h b/llvm/include/llvm/Support/DebugLog.h
index b1b17e3ede950..f1fca33932f71 100644
--- a/llvm/include/llvm/Support/DebugLog.h
+++ b/llvm/include/llvm/Support/DebugLog.h
@@ -30,30 +30,62 @@ namespace llvm {
 #define DEBUGLOG_WITH_STREAM_AND_TYPE(STREAM, TYPE)                            \
   for (bool _c = (::llvm::DebugFlag && ::llvm::isCurrentDebugType(TYPE)); _c;  \
        _c = false)                                                             \
-  ::llvm::impl::LogWithNewline(TYPE, __SHORT_FILE__, __LINE__, (STREAM))
+  ::llvm::impl::LogWithNewline(TYPE, __SHORT_FILE__, __LINE__, (STREAM))       \
+      .stream()
 #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))
+  ::llvm::impl::LogWithNewline(TYPE, __FILE__, __LINE__, (STREAM)).stream()
 #endif
 
 namespace impl {
-class LogWithNewline {
+
+/// A raw_ostream that tracks `\n` and print the prefix.
+class LLVM_ABI raw_ldbg_ostream : public raw_ostream {
+  std::string Prefix;
+  raw_ostream &Os;
+
+  /// 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);
+    auto Eol = Str.find('\n');
+    while (Eol != StringRef::npos) {
+      StringRef Line = Str.take_front(Eol + 1);
+      Os.write(Line.data(), Line.size());
+      Os.write(Prefix.c_str(), Prefix.size());
+
+      Str = Str.drop_front(Eol + 1);
+      Eol = Str.find('\n');
+    }
+    Os.write(Str.data(), Str.size());
+  }
+
 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 << " ";
+  explicit raw_ldbg_ostream(std::string Prefix, raw_ostream &Os)
+      : Prefix(std::move(Prefix)), Os(Os) {
+    SetUnbuffered();
   }
-  ~LogWithNewline() { os << '\n'; }
-  template <typename T> raw_ostream &operator<<(const T &t) && {
-    return os << t;
+
+  void emitPrefix() { Os.write(Prefix.c_str(), Prefix.size()); }
+
+  raw_ostream &getOs() { return Os; }
+
+  /// Forward the current_pos method to the underlying stream.
+  uint64_t current_pos() const final { return Os.tell(); }
+};
+
+class LogWithNewline {
+public:
+  LogWithNewline(const char *DebugType, const char *File, int Line,
+                 raw_ostream &Os)
+      : Os(computePrefix(DebugType, File, Line), Os) {}
+  ~LogWithNewline() { Os.getOs() << '\n'; }
+
+  raw_ostream &stream() {
+    Os.emitPrefix();
+    return Os;
   }
 
   // Prevent copying, as this class manages newline responsibility and is
@@ -61,6 +93,8 @@ class LogWithNewline {
   LogWithNewline(const LogWithNewline &) = delete;
   LogWithNewline &operator=(const LogWithNewline &) = delete;
   LogWithNewline &operator=(LogWithNewline &&) = delete;
+
+private:
   static constexpr const char *getShortFileName(const char *path) {
     // Remove the path prefix from the file name.
     const char *filename = path;
@@ -71,9 +105,19 @@ class LogWithNewline {
     }
     return filename;
   }
-
-private:
-  raw_ostream &os;
+  static std::string computePrefix(const char *DebugType, const char *File,
+                                   int Line) {
+    std::string Prefix;
+    raw_string_ostream OsPrefix(Prefix);
+#if !defined(__SHORT_FILE__)
+    file = ::llvm::impl::LogWithNewline::getShortFileName(file);
+#endif
+    if (DebugType)
+      OsPrefix << "[" << DebugType << "] ";
+    OsPrefix << File << ":" << Line << " ";
+    return OsPrefix.str();
+  }
+  raw_ldbg_ostream Os;
 };
 } // end namespace impl
 #else