2 files changed, 36 insertions, 33 deletions

diff --git a/clang/include/clang/Sema/SemaInternal.h b/clang/include/clang/Sema/SemaInternal.h
index 005d882c341..14e119c843f 100644
--- a/clang/include/clang/Sema/SemaInternal.h
+++ b/clang/include/clang/Sema/SemaInternal.h
@@ -221,18 +221,13 @@ private:
     std::string CurNameSpecifier;
     SmallVector<const IdentifierInfo*, 4> CurContextIdentifiers;
     SmallVector<const IdentifierInfo*, 4> CurNameSpecifierIdentifiers;
-    bool isSorted;
 
-    SpecifierInfoList Specifiers;
-    llvm::SmallSetVector<unsigned, 4> Distances;
-    llvm::DenseMap<unsigned, SpecifierInfoList> DistanceMap;
+    std::map<unsigned, SpecifierInfoList> DistanceMap;
 
     /// \brief Helper for building the list of DeclContexts between the current
     /// context and the top of the translation unit
     static DeclContextList buildContextChain(DeclContext *Start);
 
-    void sortNamespaces();
-
     unsigned buildNestedNameSpecifier(DeclContextList &DeclChain,
                                       NestedNameSpecifier *&NNS);
 
@@ -244,12 +239,40 @@ private:
     /// the corresponding NestedNameSpecifier and its distance in the process.
     void addNameSpecifier(DeclContext *Ctx);
 
-    typedef SpecifierInfoList::iterator iterator;
-    iterator begin() {
-      if (!isSorted) sortNamespaces();
-      return Specifiers.begin();
-    }
-    iterator end() { return Specifiers.end(); }
+    /// \brief Provides flat iteration over specifiers, sorted by distance.
+    class iterator
+        : public llvm::iterator_facade_base<iterator, std::forward_iterator_tag,
+                                            SpecifierInfo> {
+      /// Always points to the last element in the distance map.
+      const std::map<unsigned, SpecifierInfoList>::iterator OuterBack;
+      /// Iterator on the distance map.
+      std::map<unsigned, SpecifierInfoList>::iterator Outer;
+      /// Iterator on an element in the distance map.
+      SpecifierInfoList::iterator Inner;
+
+    public:
+      iterator(NamespaceSpecifierSet &Set, bool IsAtEnd)
+          : OuterBack(std::prev(Set.DistanceMap.end())),
+            Outer(Set.DistanceMap.begin()),
+            Inner(!IsAtEnd ? Outer->second.begin() : OuterBack->second.end()) {
+        assert(!Set.DistanceMap.empty());
+      }
+
+      iterator &operator++() {
+        ++Inner;
+        if (Inner == Outer->second.end() && Outer != OuterBack) {
+          ++Outer;
+          Inner = Outer->second.begin();
+        }
+        return *this;
+      }
+
+      SpecifierInfo &operator*() { return *Inner; }
+      bool operator==(const iterator &RHS) const { return Inner == RHS.Inner; }
+    };
+
+    iterator begin() { return iterator(*this, /*IsAtEnd=*/false); }
+    iterator end() { return iterator(*this, /*IsAtEnd=*/true); }
   };
 
   void addName(StringRef Name, NamedDecl *ND,
diff --git a/clang/lib/Sema/SemaLookup.cpp b/clang/lib/Sema/SemaLookup.cpp
index d9ac445a6a1..e15e532be19 100644
--- a/clang/lib/Sema/SemaLookup.cpp
+++ b/clang/lib/Sema/SemaLookup.cpp
@@ -3685,8 +3685,7 @@ void TypoCorrectionConsumer::performQualifiedLookups() {
 
 TypoCorrectionConsumer::NamespaceSpecifierSet::NamespaceSpecifierSet(
     ASTContext &Context, DeclContext *CurContext, CXXScopeSpec *CurScopeSpec)
-    : Context(Context), CurContextChain(buildContextChain(CurContext)),
-      isSorted(false) {
+    : Context(Context), CurContextChain(buildContextChain(CurContext)) {
   if (NestedNameSpecifier *NNS =
           CurScopeSpec ? CurScopeSpec->getScopeRep() : nullptr) {
     llvm::raw_string_ostream SpecifierOStream(CurNameSpecifier);
@@ -3705,7 +3704,6 @@ TypoCorrectionConsumer::NamespaceSpecifierSet::NamespaceSpecifierSet(
   }
 
   // Add the global context as a NestedNameSpecifier
-  Distances.insert(1);
   SpecifierInfo SI = {cast<DeclContext>(Context.getTranslationUnitDecl()),
                       NestedNameSpecifier::GlobalSpecifier(Context), 1};
   DistanceMap[1].push_back(SI);
@@ -3725,22 +3723,6 @@ auto TypoCorrectionConsumer::NamespaceSpecifierSet::buildContextChain(
   return Chain;
 }
 
-void TypoCorrectionConsumer::NamespaceSpecifierSet::sortNamespaces() {
-  SmallVector<unsigned, 4> sortedDistances;
-  sortedDistances.append(Distances.begin(), Distances.end());
-
-  if (sortedDistances.size() > 1)
-    std::sort(sortedDistances.begin(), sortedDistances.end());
-
-  Specifiers.clear();
-  for (auto D : sortedDistances) {
-    SpecifierInfoList &SpecList = DistanceMap[D];
-    Specifiers.append(SpecList.begin(), SpecList.end());
-  }
-
-  isSorted = true;
-}
-
 unsigned
 TypoCorrectionConsumer::NamespaceSpecifierSet::buildNestedNameSpecifier(
     DeclContextList &DeclChain, NestedNameSpecifier *&NNS) {
@@ -3821,8 +3803,6 @@ void TypoCorrectionConsumer::NamespaceSpecifierSet::addNameSpecifier(
         llvm::makeArrayRef(NewNameSpecifierIdentifiers));
   }
 
-  isSorted = false;
-  Distances.insert(NumSpecifiers);
   SpecifierInfo SI = {Ctx, NNS, NumSpecifiers};
   DistanceMap[NumSpecifiers].push_back(SI);
 }