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

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-llvm-support

Author: Jakub Kuderski (kuhar)

<details>
<summary>Changes</summary>

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

---
Full diff: https://github.com/llvm/llvm-project/pull/136318.diff


2 Files Affected:

- (modified) llvm/include/llvm/Support/InterleavedRange.h (+2-2) 
- (modified) llvm/unittests/Support/InterleavedRangeTest.cpp (+22) 


``````````diff
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

``````````

</details>


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