diff --git a/llvm/include/llvm/ADT/DenseMap.h b/llvm/include/llvm/ADT/DenseMap.h index bb99a41646b08..3175b3ece467c 100644 --- a/llvm/include/llvm/ADT/DenseMap.h +++ b/llvm/include/llvm/ADT/DenseMap.h @@ -17,6 +17,7 @@ #include "llvm/ADT/ADL.h" #include "llvm/ADT/DenseMapInfo.h" #include "llvm/ADT/EpochTracker.h" +#include "llvm/ADT/STLExtras.h" #include "llvm/Support/AlignOf.h" #include "llvm/Support/Compiler.h" #include "llvm/Support/MathExtras.h" @@ -96,6 +97,24 @@ class DenseMapBase : public DebugEpochBase { return makeConstIterator(getBucketsEnd(), getBucketsEnd(), *this, true); } + // Return an iterator to iterate over keys in the map. + inline auto keys() { + return map_range(*this, [](const BucketT &P) { return P.getFirst(); }); + } + + // Return an iterator to iterate over values in the map. + inline auto values() { + return map_range(*this, [](const BucketT &P) { return P.getSecond(); }); + } + + inline auto keys() const { + return map_range(*this, [](const BucketT &P) { return P.getFirst(); }); + } + + inline auto values() const { + return map_range(*this, [](const BucketT &P) { return P.getSecond(); }); + } + [[nodiscard]] bool empty() const { return getNumEntries() == 0; } unsigned size() const { return getNumEntries(); } diff --git a/llvm/unittests/ADT/DenseMapTest.cpp b/llvm/unittests/ADT/DenseMapTest.cpp index a4c045585fc28..b9d519a23c9be 100644 --- a/llvm/unittests/ADT/DenseMapTest.cpp +++ b/llvm/unittests/ADT/DenseMapTest.cpp @@ -10,6 +10,7 @@ #include "CountCopyAndMove.h" #include "llvm/ADT/DenseMapInfo.h" #include "llvm/ADT/DenseMapInfoVariant.h" +#include "llvm/ADT/SmallSet.h" #include "llvm/ADT/StringRef.h" #include "gmock/gmock.h" #include "gtest/gtest.h" @@ -359,6 +360,51 @@ TYPED_TEST(DenseMapTest, ConstIteratorTest) { EXPECT_TRUE(cit == cit2); } +TYPED_TEST(DenseMapTest, KeysValuesIterator) { + SmallSet Keys; + SmallSet Values; + for (int I = 0; I < 10; ++I) { + auto K = this->getKey(I); + auto V = this->getValue(I); + Keys.insert(K); + Values.insert(V); + this->Map[K] = V; + } + + SmallSet ActualKeys; + SmallSet ActualValues; + for (auto K : this->Map.keys()) + ActualKeys.insert(K); + for (auto V : this->Map.values()) + ActualValues.insert(V); + + EXPECT_EQ(Keys, ActualKeys); + EXPECT_EQ(Values, ActualValues); +} + +TYPED_TEST(DenseMapTest, ConstKeysValuesIterator) { + SmallSet Keys; + SmallSet Values; + for (int I = 0; I < 10; ++I) { + auto K = this->getKey(I); + auto V = this->getValue(I); + Keys.insert(K); + Values.insert(V); + this->Map[K] = V; + } + + const TypeParam &ConstMap = this->Map; + SmallSet ActualKeys; + SmallSet ActualValues; + for (auto K : ConstMap.keys()) + ActualKeys.insert(K); + for (auto V : ConstMap.values()) + ActualValues.insert(V); + + EXPECT_EQ(Keys, ActualKeys); + EXPECT_EQ(Values, ActualValues); +} + // Test initializer list construction. TEST(DenseMapCustomTest, InitializerList) { DenseMap M({{0, 0}, {0, 1}, {1, 2}});