diff --git a/llvm/include/llvm/ADT/StringMap.h b/llvm/include/llvm/ADT/StringMap.h --- a/llvm/include/llvm/ADT/StringMap.h +++ b/llvm/include/llvm/ADT/StringMap.h @@ -14,12 +14,16 @@ #ifndef LLVM_ADT_STRINGMAP_H #define LLVM_ADT_STRINGMAP_H +#include "llvm/ADT/STLExtras.h" #include "llvm/ADT/StringMapEntry.h" #include "llvm/ADT/iterator.h" #include "llvm/Support/AllocatorBase.h" +#include "llvm/Support/FormatVariadic.h" #include "llvm/Support/PointerLikeTypeTraits.h" #include #include +#include +#include namespace llvm { @@ -483,6 +487,27 @@ StringRef operator*() const { return this->wrapped()->getKey(); } }; +template +std::ostream &operator<<(std::ostream &OS, const StringMap &M) { + if (M.empty()) { + return OS << "{ }"; + } + + std::vector Lines; + for (const auto &E : M) { + std::ostringstream SS; + SS << E << ","; + Lines.push_back(SS.str()); + } + llvm::sort(Lines.begin(), Lines.end()); + Lines.insert(Lines.begin(), "{"); + Lines.insert(Lines.end(), "}"); + + return OS << llvm::formatv("{0:$[\n]}", + make_range(Lines.begin(), Lines.end())) + .str(); +} + } // end namespace llvm #endif // LLVM_ADT_STRINGMAP_H diff --git a/llvm/include/llvm/ADT/StringMapEntry.h b/llvm/include/llvm/ADT/StringMapEntry.h --- a/llvm/include/llvm/ADT/StringMapEntry.h +++ b/llvm/include/llvm/ADT/StringMapEntry.h @@ -19,6 +19,7 @@ #include "llvm/ADT/None.h" #include "llvm/ADT/STLFunctionalExtras.h" #include "llvm/ADT/StringRef.h" +#include namespace llvm { @@ -145,6 +146,11 @@ } }; +template +std::ostream &operator<<(std::ostream &OS, const StringMapEntry &S) { + return OS << "{\"" << S.getKey().data() << "\": " << S.getValue() << "}"; +} + } // end namespace llvm #endif // LLVM_ADT_STRINGMAPENTRY_H diff --git a/llvm/unittests/ADT/StringMapTest.cpp b/llvm/unittests/ADT/StringMapTest.cpp --- a/llvm/unittests/ADT/StringMapTest.cpp +++ b/llvm/unittests/ADT/StringMapTest.cpp @@ -12,6 +12,8 @@ #include "llvm/Support/DataTypes.h" #include "gtest/gtest.h" #include +#include +#include #include using namespace llvm; @@ -635,4 +637,40 @@ EXPECT_EQ(LargeValue, Key.size()); } +TEST(StringMapCustomTest, StringMapEntryStream) { + std::ostringstream OS; + StringMap Map; + Map["A"] = 42; + Map["Z"] = 35; + Map["B"] = 7; + + OS << Map; + EXPECT_EQ(OS.str(), R"({ +{"A": 42}, +{"B": 7}, +{"Z": 35}, +})"); +} + +TEST(StringMapCustomTest, NestedStringMapEntryStream) { + std::ostringstream OS; + StringMap> Map; + Map["Z"]; + Map["A"]["Apple"] = 5; + Map["B"]["Bee"] = 3; + Map["A"]["Axe"] = 3; + + OS << Map; + EXPECT_EQ(OS.str(), R"({ +{"A": { +{"Apple": 5}, +{"Axe": 3}, +}}, +{"B": { +{"Bee": 3}, +}}, +{"Z": { }}, +})"); +} + } // end anonymous namespace