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 <unsigned Size, typename R> -SmallVector<detail::ValueOfRange<R>, Size> to_vector(R &&Range) { +SmallVector<typename std::remove_const<detail::ValueOfRange<R>>::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 <cassert> #include <cstdint> #include <cstdlib> #include <cstring> -#include <utility> #include <initializer_list> #include <new> #include <utility> @@ -32,6 +32,7 @@ class StringMapConstIterator; template<typename ValueT> class StringMapIterator; + template <typename ValueT> class StringMapKeyIterator; template<typename ValueTy> class StringMapEntry; @@ -312,6 +313,11 @@ return const_iterator(TheTable+NumBuckets, true); } + llvm::iterator_range<StringMapKeyIterator<ValueTy>> keys() const { + return make_range(StringMapKeyIterator<ValueTy>(begin()), + StringMapKeyIterator<ValueTy>(end())); + } + iterator find(StringRef Key) { int Bucket = FindKey(Key); if (Bucket == -1) return end(); @@ -444,42 +450,39 @@ } }; -template <typename ValueTy> class StringMapConstIterator { +template <typename DerivedTy, typename ValueTy> +class StringMapIterBase + : public iterator_facade_base<DerivedTy, std::forward_iterator_tag, + ValueTy> { protected: StringMapEntryBase **Ptr = nullptr; public: - typedef StringMapEntry<ValueTy> 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<StringMapEntry<ValueTy>*>(*Ptr); - } - const value_type *operator->() const { - return static_cast<StringMapEntry<ValueTy>*>(*Ptr); + DerivedTy &operator=(const DerivedTy &Other) { + Ptr = Other.Ptr; + return static_cast<DerivedTy &>(*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<DerivedTy &>(*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<typename ValueTy> -class StringMapIterator : public StringMapConstIterator<ValueTy> { +template <typename ValueTy> +class StringMapConstIterator + : public StringMapIterBase<StringMapConstIterator<ValueTy>, + const StringMapEntry<ValueTy>> { public: - StringMapIterator() = default; + StringMapConstIterator() = default; + explicit StringMapConstIterator(StringMapEntryBase **Bucket, + bool NoAdvance = false) + : StringMapIterBase(Bucket, NoAdvance) {} + + const StringMapEntry<ValueTy> &operator*() const { + return *static_cast<const StringMapEntry<ValueTy> *>(*Ptr); + } +}; +template <typename ValueTy> +class StringMapIterator : public StringMapIterBase<StringMapIterator<ValueTy>, + StringMapEntry<ValueTy>> { +public: + StringMapIterator() = default; explicit StringMapIterator(StringMapEntryBase **Bucket, bool NoAdvance = false) - : StringMapConstIterator<ValueTy>(Bucket, NoAdvance) { - } + : StringMapIterBase(Bucket, NoAdvance) {} StringMapEntry<ValueTy> &operator*() const { - return *static_cast<StringMapEntry<ValueTy>*>(*this->Ptr); + return *static_cast<StringMapEntry<ValueTy> *>(*Ptr); } - StringMapEntry<ValueTy> *operator->() const { - return static_cast<StringMapEntry<ValueTy>*>(*this->Ptr); + + operator StringMapConstIterator<ValueTy>() const { + return StringMapConstIterator<ValueTy>(Ptr, false); } }; +template <typename ValueTy> +class StringMapKeyIterator + : public iterator_adaptor_base<StringMapKeyIterator<ValueTy>, + StringMapConstIterator<ValueTy>, + std::forward_iterator_tag, StringRef> { +public: + StringMapKeyIterator() = default; + + explicit StringMapKeyIterator(StringMapConstIterator<ValueTy> 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<int> 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<StringRef, 4> 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<StringRef, 4> Expected = {"A", "B", "C", "D"}; + EXPECT_EQ(Expected, Keys); +} + // Create a non-default constructable value struct StringMapTestStruct { StringMapTestStruct(int i) : i(i) {}