Improve StringMap iterator support.

[zturner]

Summary

Make the two StringMap iterators support iterator_facade_base. This fixes some usage scenarios when passing begin and end to various STL algorithms. After doing that, it's easy to make a StringMapKeysIterator that only iterates keys, so I've done that as well.

[dblaikie]

Seems weird to have 'keys' on a StringSet, perhaps? (& to test StringSet in the StringMap unit test - as much as I realize the two types are very closely related)

Also, I'm wondering if keys() is necessary even on StringMap - I'd have thought a general-purpose "iterator of the first element of a pair" would suffice, but then remembered that StringMaps are special and their values are not pair<key, value> but have special accessors for the key, etc. So perhaps a lambda-driven iterator adapter... at which point, yeah, I can understand the utility of having that wrapped up in a reusable form.

llvm/include/llvm/ADT/StringExtras.h

255 ↗

(On Diff #92433)

most of the other "iterator pair or range" algorithms use the same name for both (eg: all the llvm::any/all/none_of clones that take a range use the same name) - maybe use the same name here too ('join' rather than 'join_range') unless that causes some ambiguities?

Also - this seems like an unrelated change for this patch (& arguably could use test coverage if it is going to be committed). Guess maybe it slipped into this change.

llvm/include/llvm/ADT/StringMap.h

474

Generally it's advised that any op overload that can be a non-member should be, to allow equal implicit conversions on both sides. (it can still be a friend, though:

friend bool operator==(LHS, RHS) { return LHS.Ptr == RHS.Ptr; }

539

op* that returns by value - does that cause problems for op->? I forget...

llvm/unittests/Support/StringMap.cpp

38 ↗

(On Diff #92433)

In retrospect, maybe this function should be called To(Small)Vector? (Since there's no standard library function it's trying to emulate, etc - and it actually produces a SmallVector, not a std::vector, so the name is a bit vague/misleading (& makes the int non-type template parameter unclear/confusing))

[zturner]

llvm/include/llvm/ADT/StringExtras.h

255 ↗

(On Diff #92433)

Yea I almost called it join, but then saw that we have join and join_items so join_range seemed to fit nicely. That said, join_items was introduced because of an ambiguity with join, but join_range won't have that ambiguity, so we could still call it join.

llvm/include/llvm/ADT/StringMap.h

539

Should be fine, the StringRef should live until the end of the sequence point.

[zturner]

1. Fixed issue with returning by value from operator*.
2. Moved the test from Support to ADT, since there was already an existing StringMapTest class there that I hadn't noticed.
3. Removed the join changes that had snuck in. I'll commit those separately.

[dblaikie]

Some optional bits/pieces. (also - could include test coverage for op->. Or perhaps better/alternatively: put a static_assert in iterator_adapter_base to ensure that the return type of op* is a reference type?)

llvm/unittests/ADT/StringMapTest.cpp

287–300

Not sure this test is necessary (I mean, depending on your perspective) - since it's testing the same API in StringMap that's inherited by StringSet - there's no new/interesting code under test here, compared to the previous test.

(should StringSet's iterators actually be key iterators, rather than a 'keys()' operation on a set, which doesn't quite make sense - it only has keys, no values, so why would there be two different ranges to iterate over?)

[zturner]

llvm/unittests/ADT/StringMapTest.cpp

287–300

Not sure this test is necessary (I mean, depending on your perspective) - since it's testing the same API in StringMap that's inherited by StringSet - there's no new/interesting code under test here, compared to the previous test.

True, but the same could be set for join_range vs join in the other patch. It doesn't matter to me either way though.

should StringSet's iterators actually be key iterators, rather than a 'keys()' operation on a set, which doesn't quite make sense - it only has keys, no values, so why would there be two different ranges to iterate over

I could rename the function in StringSet by bringing keys() into protected scope with a using declaration, then make a function in StringSet called values() which simply returns the result of keys(). Thoughts?

llvm/include/llvm/ADT/STLExtras.h

 auto partition(R &&Range, UnaryPredicate P) -> decltype(std::begin(Range)) {
   return std::partition(std::begin(Range), std::end(Range), P);
 }
 
 /// \brief Given a range of type R, iterate the entire range and return a
 /// 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)};
 }
 
 /// Provide a container algorithm similar to C++ Library Fundamentals v2's
 /// `erase_if` which is equivalent to:
 ///
 ///   C.erase(remove_if(C, pred), C.end());
 ///

llvm/include/llvm/ADT/StringMap.h

 // This file defines the StringMap class.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_ADT_STRINGMAP_H
 #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>
 
 namespace llvm {
 
   template<typename ValueT>
   class StringMapConstIterator;
   template<typename ValueT>
   class StringMapIterator;
+  template <typename ValueT> class StringMapKeyIterator;
   template<typename ValueTy>
   class StringMapEntry;
 
 /// StringMapEntryBase - Shared base class of StringMapEntry instances.
 class StringMapEntryBase {
   unsigned StrLen;
 
 public:
public:
   }
   const_iterator begin() const {
     return const_iterator(TheTable, NumBuckets == 0);
   }
   const_iterator end() const {
     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();
     return iterator(TheTable+Bucket, true);
   }
 
   const_iterator find(StringRef Key) const {
     int Bucket = FindKey(Key);
if (!empty()) {
           static_cast<MapEntryTy*>(Bucket)->Destroy(Allocator);
         }
       }
     }
     free(TheTable);
   }
 };
 
-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)
+  explicit StringMapIterBase(StringMapEntryBase **Bucket,
       : 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:
   void AdvancePastEmptyBuckets() {
     while (*Ptr == nullptr || *Ptr == StringMapImpl::getTombstoneVal())
       ++Ptr;
   }
 };
 
-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);
+  }
+
+  operator StringMapConstIterator<ValueTy>() const {
+    return StringMapConstIterator<ValueTy>(Ptr, false);
   }
-  StringMapEntry<ValueTy> *operator->() const {
-    return static_cast<StringMapEntry<ValueTy>*>(*this->Ptr);
+};
+
+template <typename ValueTy>
+class StringMapKeyIterator
+    : public iterator_adaptor_base<StringMapKeyIterator<ValueTy>,
+                                   StringMapConstIterator<ValueTy>,
+                                   std::forward_iterator_tag, StringRef> {
+  using BaseT = 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

llvm/unittests/ADT/StringMapTest.cpp

 //===- llvm/unittest/ADT/StringMapMap.cpp - StringMap unit tests ----------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ADT/StringMap.h"
+#include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Support/DataTypes.h"
 #include "gtest/gtest.h"
 #include <tuple>
 using namespace llvm;
 
 namespace {
 
TEST_F(StringMapTest, InsertRehashingPairTest) {
 
   StringMap<uint32_t>::iterator It =
     t.insert(std::make_pair("abcdef", 42)).first;
   EXPECT_EQ(16u, t.getNumBuckets());
   EXPECT_EQ("abcdef", It->first());
   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) {}
   StringMapTestStruct() = delete;
   int i;
 };
 
 TEST_F(StringMapTest, NonDefaultConstructable) {