[ADT] Refactor StringMap iterators (NFC)

clang-tools-extra/clang-tidy/fuchsia/MultipleInheritanceCheck.cpp

@@ -39,7 +39,7 @@ bool MultipleInheritanceCheck::getInterfaceStatus(const CXXRecordDecl *Node,
                                                   bool &IsInterface) const {
   assert(Node->getIdentifier());
   StringRef Name = Node->getIdentifier()->getName();
-  llvm::StringMapConstIterator<bool> Pair = InterfaceMap.find(Name);
+  auto Pair = InterfaceMap.find(Name);
   if (Pair == InterfaceMap.end())
     return false;
   IsInterface = Pair->second;

clang/lib/Driver/OffloadBundler.cpp

@@ -1936,8 +1936,7 @@ Error OffloadBundler::UnbundleArchive() {
   /// Write out an archive for each target
   for (auto &Target : BundlerConfig.TargetNames) {
     StringRef FileName = TargetOutputFileNameMap[Target];
-    StringMapIterator<std::vector<llvm::NewArchiveMember>> CurArchiveMembers =
-        OutputArchivesMap.find(Target);
+    auto CurArchiveMembers = OutputArchivesMap.find(Target);
     if (CurArchiveMembers != OutputArchivesMap.end()) {
       if (Error WriteErr = writeArchive(FileName, CurArchiveMembers->getValue(),
                                         SymtabWritingMode::NormalSymtab,

llvm/include/llvm/ADT/StringMap.h

@@ -21,11 +21,11 @@
 #include "llvm/Support/PointerLikeTypeTraits.h"
 #include <initializer_list>
 #include <iterator>
+#include <type_traits>
 
 namespace llvm {
 
-template <typename ValueTy> class StringMapConstIterator;
-template <typename ValueTy> class StringMapIterator;
+template <typename ValueTy, bool IsConst> class StringMapIterBase;
 template <typename ValueTy> class StringMapKeyIterator;
 
 /// StringMapImpl - This is the base class of StringMap that is shared among
@@ -217,8 +217,8 @@ class LLVM_ALLOCATORHOLDER_EMPTYBASE StringMap
   using value_type = StringMapEntry<ValueTy>;
   using size_type = size_t;
 
-  using const_iterator = StringMapConstIterator<ValueTy>;
-  using iterator = StringMapIterator<ValueTy>;
+  using const_iterator = StringMapIterBase<ValueTy, true>;
+  using iterator = StringMapIterBase<ValueTy, false>;
 
   iterator begin() { return iterator(TheTable, NumBuckets == 0); }
   iterator end() { return iterator(TheTable + NumBuckets, true); }
@@ -431,14 +431,17 @@ class LLVM_ALLOCATORHOLDER_EMPTYBASE StringMap
   }
 };
 
-template <typename DerivedTy, typename ValueTy>
-class StringMapIterBase
-    : public iterator_facade_base<DerivedTy, std::forward_iterator_tag,
-                                  ValueTy> {
-protected:
+template <typename ValueTy, bool IsConst> class StringMapIterBase {
   StringMapEntryBase **Ptr = nullptr;
 
 public:
+  using iterator_category = std::forward_iterator_tag;
+  using value_type = StringMapEntry<ValueTy>;
+  using difference_type = std::ptrdiff_t;
+  using pointer = std::conditional_t<IsConst, const value_type *, value_type *>;
+  using reference =
+      std::conditional_t<IsConst, const value_type &, value_type &>;
+
   StringMapIterBase() = default;
 
   explicit StringMapIterBase(StringMapEntryBase **Bucket,
@@ -448,85 +451,56 @@ class StringMapIterBase
       AdvancePastEmptyBuckets();
   }
 
-  DerivedTy &operator=(const DerivedTy &Other) {
-    Ptr = Other.Ptr;
-    return static_cast<DerivedTy &>(*this);
-  }
-
-  friend bool operator==(const DerivedTy &LHS, const DerivedTy &RHS) {
-    return LHS.Ptr == RHS.Ptr;
-  }
+  reference operator*() const { return *static_cast<value_type *>(*Ptr); }
+  pointer operator->() const { return static_cast<value_type *>(*Ptr); }
 
-  DerivedTy &operator++() { // Preincrement
+  StringMapIterBase &operator++() { // Preincrement
     ++Ptr;
     AdvancePastEmptyBuckets();
-    return static_cast<DerivedTy &>(*this);
+    return *this;
   }
 
-  DerivedTy operator++(int) { // Post-increment
-    DerivedTy Tmp(Ptr);
+  StringMapIterBase operator++(int) { // Post-increment
+    StringMapIterBase Tmp(*this);
     ++*this;
     return Tmp;
   }
 
-private:
-  void AdvancePastEmptyBuckets() {
-    while (*Ptr == nullptr || *Ptr == StringMapImpl::getTombstoneVal())
-      ++Ptr;
+  template <bool ToConst,
+            typename = typename std::enable_if<!IsConst && ToConst>::type>
+  operator StringMapIterBase<ValueTy, ToConst>() const {
+    return StringMapIterBase<ValueTy, ToConst>(Ptr, true);
   }
-};
-
-template <typename ValueTy>
-class StringMapConstIterator
-    : public StringMapIterBase<StringMapConstIterator<ValueTy>,
-                               const StringMapEntry<ValueTy>> {
-  using base = StringMapIterBase<StringMapConstIterator<ValueTy>,
-                                 const StringMapEntry<ValueTy>>;
-
-public:
-  StringMapConstIterator() = default;
-  explicit StringMapConstIterator(StringMapEntryBase **Bucket,
-                                  bool NoAdvance = false)
-      : base(Bucket, NoAdvance) {}
 
-  const StringMapEntry<ValueTy> &operator*() const {
-    return *static_cast<const StringMapEntry<ValueTy> *>(*this->Ptr);
+  friend bool operator==(const StringMapIterBase &LHS,
+                         const StringMapIterBase &RHS) {
+    return LHS.Ptr == RHS.Ptr;
   }
-};
-
-template <typename ValueTy>
-class StringMapIterator : public StringMapIterBase<StringMapIterator<ValueTy>,
-                                                   StringMapEntry<ValueTy>> {
-  using base =
-      StringMapIterBase<StringMapIterator<ValueTy>, StringMapEntry<ValueTy>>;
 
-public:
-  StringMapIterator() = default;
-  explicit StringMapIterator(StringMapEntryBase **Bucket,
-                             bool NoAdvance = false)
-      : base(Bucket, NoAdvance) {}
-
-  StringMapEntry<ValueTy> &operator*() const {
-    return *static_cast<StringMapEntry<ValueTy> *>(*this->Ptr);
+  friend bool operator!=(const StringMapIterBase &LHS,
+                         const StringMapIterBase &RHS) {
+    return !(LHS == RHS);
   }
 
-  operator StringMapConstIterator<ValueTy>() const {
-    return StringMapConstIterator<ValueTy>(this->Ptr, true);
+private:
+  void AdvancePastEmptyBuckets() {
+    while (*Ptr == nullptr || *Ptr == StringMapImpl::getTombstoneVal())
+      ++Ptr;
   }
 };
 
 template <typename ValueTy>
 class StringMapKeyIterator
     : public iterator_adaptor_base<StringMapKeyIterator<ValueTy>,
-                                   StringMapConstIterator<ValueTy>,
+                                   StringMapIterBase<ValueTy, true>,
                                    std::forward_iterator_tag, StringRef> {
   using base = iterator_adaptor_base<StringMapKeyIterator<ValueTy>,
-                                     StringMapConstIterator<ValueTy>,
+                                     StringMapIterBase<ValueTy, true>,
                                      std::forward_iterator_tag, StringRef>;
 
 public:
   StringMapKeyIterator() = default;
-  explicit StringMapKeyIterator(StringMapConstIterator<ValueTy> Iter)
+  explicit StringMapKeyIterator(StringMapIterBase<ValueTy, true> Iter)
       : base(std::move(Iter)) {}
 
   StringRef operator*() const { return this->wrapped()->getKey(); }

llvm/unittests/ADT/StringMapTest.cpp

@@ -693,4 +693,58 @@ TEST(StringMapCustomTest, StringMapEntrySize) {
   EXPECT_EQ(LargeValue, Key.size());
 }
 
+TEST(StringMapCustomTest, NonConstIterator) {
+  StringMap<int> Map;
+  Map["key"] = 1;
+
+  // Check that Map.begin() returns a non-const iterator.
+  static_assert(
+      std::is_same_v<decltype(Map.begin()), StringMap<int>::iterator>);
+
+  // Check that the value_type of a non-const iterator is not a const type.
+  static_assert(
+      !std::is_const_v<StringMap<int>::iterator::value_type>,
+      "The value_type of a non-const iterator should not be a const type.");
+
+  // Check that pointer and reference types are not const.
+  static_assert(std::is_same_v<StringMap<int>::iterator::pointer,
+                               StringMap<int>::iterator::value_type *>);
+  static_assert(std::is_same_v<StringMap<int>::iterator::reference,
+                               StringMap<int>::iterator::value_type &>);
+
+  // Check that we can construct a const_iterator from an iterator.
+  static_assert(std::is_constructible_v<StringMap<int>::const_iterator,
+                                        StringMap<int>::iterator>);
+
+  // Double check that we can actually construct a const_iterator.
+  StringMap<int>::const_iterator const_it = Map.begin();
+  (void)const_it;
+}
+
+TEST(StringMapCustomTest, ConstIterator) {
+  StringMap<int> Map;
+  const StringMap<int> &ConstMap = Map;
+
+  // Check that ConstMap.begin() returns a const_iterator.
+  static_assert(std::is_same_v<decltype(ConstMap.begin()),
+                               StringMap<int>::const_iterator>);
+
+  // Check that the value_type of a const iterator is not a const type.
+  static_assert(
+      !std::is_const_v<StringMap<int>::const_iterator::value_type>,
+      "The value_type of a const iterator should not be a const type.");
+
+  // Check that pointer and reference types are const.
+  static_assert(
+      std::is_same_v<StringMap<int>::const_iterator::pointer,
+                     const StringMap<int>::const_iterator::value_type *>);
+  static_assert(
+      std::is_same_v<StringMap<int>::const_iterator::reference,
+                     const StringMap<int>::const_iterator::value_type &>);
+
+  // Check that we cannot construct an iterator from a const_iterator.
+  static_assert(!std::is_constructible_v<StringMap<int>::iterator,
+                                         StringMap<int>::const_iterator>);
+}
+
 } // end anonymous namespace