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

[via llvm-commits]

Author: Jakub Kuderski
Date: 2025-04-18T12:37:16-04:00
New Revision: ad6c23a7b5ecd5582566a6e51c1aa8105fb3f109

URL: https://github.com/llvm/llvm-project/commit/ad6c23a7b5ecd5582566a6e51c1aa8105fb3f109
DIFF: https://github.com/llvm/llvm-project/commit/ad6c23a7b5ecd5582566a6e51c1aa8105fb3f109.diff

LOG: [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

Added: 
    

Modified: 
    llvm/include/llvm/Support/InterleavedRange.h
    llvm/unittests/Support/InterleavedRangeTest.cpp

Removed: 
    


################################################################################
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