[Support] Allow `llvm::interleaved` with custom ostream types (#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

1 files changed, 22 insertions, 0 deletions

diff --git a/llvm/unittests/Support/InterleavedRangeTest.cpp b/llvm/unittests/Support/InterleavedRangeTest.cpp
index 8640b81..f258eac 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