[llvm] [Support] Allow `llvm::interleaved` with custom ostream types (PR #136318)

[Jakub Kuderski via llvm-commits]

https://github.com/kuhar created https://github.com/llvm/llvm-project/pull/136318

This makes `llvm::interleaved` useable with ostream types that define custom stream operators that print in a different format from `raw_ostream`. For example, MLIR's OpAsmPrinter prints values as operands: https://github.com/llvm/llvm-project/blob/6c5f50f18694a4d91d7ce53a14188c54ee7c6f3b/mlir/include/mlir/IR/OpImplementation.h#L534-L552

>From e46161b65655b3314ddf449ecc7634bf5d11dc81 Mon Sep 17 00:00:00 2001
From: Jakub Kuderski <jakub at nod-labs.com>
Date: Fri, 18 Apr 2025 11:55:53 -0400
Subject: [PATCH] [Support] Allow `llvm::interleaved` with custom ostream types

This makes `llvm::interleaved` useable with ostream types that define
custom stream operators that print in a different format from
`raw_ostream`. For example, MLIR's OpAsmPrinter prints values as
operands: https://github.com/llvm/llvm-project/blob/6c5f50f18694a4d91d7ce53a14188c54ee7c6f3b/mlir/include/mlir/IR/OpImplementation.h#L534-L552
---
 llvm/include/llvm/Support/InterleavedRange.h  |  4 ++--
 .../Support/InterleavedRangeTest.cpp          | 22 +++++++++++++++++++
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/llvm/include/llvm/Support/InterleavedRange.h b/llvm/include/llvm/Support/InterleavedRange.h
index 4e70028504806..ce7cac98bcdd4 100644
--- a/llvm/include/llvm/Support/InterleavedRange.h
+++ b/llvm/include/llvm/Support/InterleavedRange.h
@@ -42,8 +42,8 @@ template <typename Range> class InterleavedRange {
                    StringRef Suffix)
       : TheRange(R), Separator(Separator), Prefix(Prefix), Suffix(Suffix) {}
 
-  friend raw_ostream &operator<<(raw_ostream &OS,
-                                 const InterleavedRange &Interleaved) {
+  template <typename OStream>
+  friend OStream &operator<<(OStream &OS, const InterleavedRange &Interleaved) {
     if (!Interleaved.Prefix.empty())
       OS << Interleaved.Prefix;
     llvm::interleave(Interleaved.TheRange, OS, Interleaved.Separator);
diff --git a/llvm/unittests/Support/InterleavedRangeTest.cpp b/llvm/unittests/Support/InterleavedRangeTest.cpp
index 8640b81fe8ad8..f258eac04e9c6 100644
--- a/llvm/unittests/Support/InterleavedRangeTest.cpp
+++ b/llvm/unittests/Support/InterleavedRangeTest.cpp
@@ -67,4 +67,26 @@ TEST(InterleavedRangeTest, CustomPrint) {
   EXPECT_EQ("[$$3##, $$4##, $$5##]", interleaved_array(V).str());
 }
 
+struct CustomDoublingOStream : raw_string_ostream {
+  unsigned NumCalled = 0;
+  using raw_string_ostream::raw_string_ostream;
+
+  friend CustomDoublingOStream &operator<<(CustomDoublingOStream &OS, int V) {
+    ++OS.NumCalled;
+    static_cast<raw_string_ostream &>(OS) << (2 * V);
+    return OS;
+  }
+};
+
+TEST(InterleavedRangeTest, CustomOStream) {
+  // Make sure that interleaved calls the stream operator on the derived class,
+  // and that it returns a reference to the same stream type.
+  int V[] = {3, 4, 5};
+  std::string Buf;
+  CustomDoublingOStream OS(Buf);
+  OS << interleaved(V) << 22;
+  EXPECT_EQ("6, 8, 1044", Buf);
+  EXPECT_EQ(OS.NumCalled, 4u);
+}
+
 } // namespace