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

Author: kuhar

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);

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