Index: llvm/trunk/include/llvm/ADT/STLExtras.h =================================================================== --- llvm/trunk/include/llvm/ADT/STLExtras.h +++ llvm/trunk/include/llvm/ADT/STLExtras.h @@ -893,7 +893,8 @@ /// SmallVector with elements of the vector. This is useful, for example, /// when you want to iterate a range and then sort the results. template -SmallVector, Size> to_vector(R &&Range) { +SmallVector>::type, Size> +to_vector(R &&Range) { return {std::begin(Range), std::end(Range)}; } Index: llvm/trunk/include/llvm/ADT/StringMap.h =================================================================== --- llvm/trunk/include/llvm/ADT/StringMap.h +++ llvm/trunk/include/llvm/ADT/StringMap.h @@ -15,13 +15,13 @@ #define LLVM_ADT_STRINGMAP_H #include "llvm/ADT/StringRef.h" +#include "llvm/ADT/iterator.h" #include "llvm/Support/Allocator.h" #include "llvm/Support/PointerLikeTypeTraits.h" #include #include #include #include -#include #include #include #include @@ -32,6 +32,7 @@ class StringMapConstIterator; template class StringMapIterator; + template class StringMapKeyIterator; template class StringMapEntry; @@ -312,6 +313,11 @@ return const_iterator(TheTable+NumBuckets, true); } + llvm::iterator_range> keys() const { + return make_range(StringMapKeyIterator(begin()), + StringMapKeyIterator(end())); + } + iterator find(StringRef Key) { int Bucket = FindKey(Key); if (Bucket == -1) return end(); @@ -444,42 +450,39 @@ } }; -template class StringMapConstIterator { +template +class StringMapIterBase + : public iterator_facade_base { protected: StringMapEntryBase **Ptr = nullptr; public: - typedef StringMapEntry value_type; + StringMapIterBase() = default; - StringMapConstIterator() = default; - - explicit StringMapConstIterator(StringMapEntryBase **Bucket, - bool NoAdvance = false) - : Ptr(Bucket) { + explicit StringMapIterBase(StringMapEntryBase **Bucket, + bool NoAdvance = false) + : Ptr(Bucket) { if (!NoAdvance) AdvancePastEmptyBuckets(); } - const value_type &operator*() const { - return *static_cast*>(*Ptr); - } - const value_type *operator->() const { - return static_cast*>(*Ptr); + DerivedTy &operator=(const DerivedTy &Other) { + Ptr = Other.Ptr; + return static_cast(*this); } - bool operator==(const StringMapConstIterator &RHS) const { - return Ptr == RHS.Ptr; - } - bool operator!=(const StringMapConstIterator &RHS) const { - return Ptr != RHS.Ptr; - } + bool operator==(const DerivedTy &RHS) const { return Ptr == RHS.Ptr; } - inline StringMapConstIterator& operator++() { // Preincrement + DerivedTy &operator++() { // Preincrement ++Ptr; AdvancePastEmptyBuckets(); - return *this; + return static_cast(*this); } - StringMapConstIterator operator++(int) { // Postincrement - StringMapConstIterator tmp = *this; ++*this; return tmp; + + DerivedTy operator++(int) { // Post-increment + DerivedTy Tmp(Ptr); + ++*this; + return Tmp; } private: @@ -489,24 +492,59 @@ } }; -template -class StringMapIterator : public StringMapConstIterator { +template +class StringMapConstIterator + : public StringMapIterBase, + const StringMapEntry> { public: - StringMapIterator() = default; + StringMapConstIterator() = default; + explicit StringMapConstIterator(StringMapEntryBase **Bucket, + bool NoAdvance = false) + : StringMapIterBase(Bucket, NoAdvance) {} + + const StringMapEntry &operator*() const { + return *static_cast *>(*Ptr); + } +}; +template +class StringMapIterator : public StringMapIterBase, + StringMapEntry> { +public: + StringMapIterator() = default; explicit StringMapIterator(StringMapEntryBase **Bucket, bool NoAdvance = false) - : StringMapConstIterator(Bucket, NoAdvance) { - } + : StringMapIterBase(Bucket, NoAdvance) {} StringMapEntry &operator*() const { - return *static_cast*>(*this->Ptr); + return *static_cast *>(*Ptr); } - StringMapEntry *operator->() const { - return static_cast*>(*this->Ptr); + + operator StringMapConstIterator() const { + return StringMapConstIterator(Ptr, false); } }; +template +class StringMapKeyIterator + : public iterator_adaptor_base, + StringMapConstIterator, + std::forward_iterator_tag, StringRef> { +public: + StringMapKeyIterator() = default; + + explicit StringMapKeyIterator(StringMapConstIterator Iter) + : iterator_adaptor_base(std::move(Iter)) {} + + StringRef &operator*() { + Key = wrapped()->getKey(); + return Key; + } + +private: + StringRef Key; +}; + } // end namespace llvm #endif // LLVM_ADT_STRINGMAP_H Index: llvm/trunk/unittests/ADT/StringMapTest.cpp =================================================================== --- llvm/trunk/unittests/ADT/StringMapTest.cpp +++ llvm/trunk/unittests/ADT/StringMapTest.cpp @@ -8,6 +8,7 @@ //===----------------------------------------------------------------------===// #include "llvm/ADT/StringMap.h" +#include "llvm/ADT/StringSet.h" #include "llvm/ADT/Twine.h" #include "llvm/Support/DataTypes.h" #include "gtest/gtest.h" @@ -269,6 +270,34 @@ EXPECT_EQ(42u, It->second); } +TEST_F(StringMapTest, IterMapKeys) { + StringMap Map; + Map["A"] = 1; + Map["B"] = 2; + Map["C"] = 3; + Map["D"] = 3; + + auto Keys = to_vector<4>(Map.keys()); + std::sort(Keys.begin(), Keys.end()); + + SmallVector Expected = {"A", "B", "C", "D"}; + EXPECT_EQ(Expected, Keys); +} + +TEST_F(StringMapTest, IterSetKeys) { + StringSet<> Set; + Set.insert("A"); + Set.insert("B"); + Set.insert("C"); + Set.insert("D"); + + auto Keys = to_vector<4>(Set.keys()); + std::sort(Keys.begin(), Keys.end()); + + SmallVector Expected = {"A", "B", "C", "D"}; + EXPECT_EQ(Expected, Keys); +} + // Create a non-default constructable value struct StringMapTestStruct { StringMapTestStruct(int i) : i(i) {}