[AST] Use DenseMap for the cache

Author: theblixguy

include/swift/AST/ASTContext.h

@@ -17,9 +17,9 @@
 #ifndef SWIFT_AST_ASTCONTEXT_H
 #define SWIFT_AST_ASTCONTEXT_H
 
-#include "llvm/Support/DataTypes.h"
 #include "swift/AST/ClangModuleLoader.h"
 #include "swift/AST/Evaluator.h"
+#include "swift/AST/GenericSignature.h"
 #include "swift/AST/Identifier.h"
 #include "swift/AST/SearchPathOptions.h"
 #include "swift/AST/Type.h"
@@ -28,13 +28,14 @@
 #include "swift/Basic/Malloc.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/IntrusiveRefCntPtr.h"
+#include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/PointerIntPair.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/TinyPtrVector.h"
 #include "llvm/Support/Allocator.h"
+#include "llvm/Support/DataTypes.h"
 #include <functional>
 #include <memory>
 #include <utility>
@@ -870,25 +871,19 @@ class ASTContext final {
   void getVisibleTopLevelModuleNames(SmallVectorImpl<Identifier> &names) const;
 
   struct OverrideSignatureKey {
-    StringRef baseMethodSigString;
-    StringRef derivedClassSigString;
+    GenericSignature *baseMethodSig;
+    GenericSignature *derivedClassSig;
     Type superclassTy;
 
-    OverrideSignatureKey(StringRef baseMethodSig, StringRef derivedClassSig,
-                         Type derivedClassSuperclassType)
-        : baseMethodSigString(baseMethodSig),
-          derivedClassSigString(derivedClassSig),
-          superclassTy(derivedClassSuperclassType) {}
-
-    bool isEqual(const OverrideSignatureKey other) {
-      return baseMethodSigString == other.baseMethodSigString &&
-             derivedClassSigString == other.derivedClassSigString &&
-             superclassTy.getString() == other.superclassTy.getString();
+    OverrideSignatureKey(GenericSignature *baseMethodSignature,
+                         GenericSignature *derivedClassSignature,
+                         Type superclassType)
+        : baseMethodSig(baseMethodSignature),
+          derivedClassSig(derivedClassSignature), superclassTy(superclassType) {
     }
   };
 
-  std::vector<std::pair<OverrideSignatureKey, GenericSignature *>>
-      overrideSigCache;
+  llvm::DenseMap<OverrideSignatureKey, GenericSignature *> overrideSigCache;
 
 private:
   /// Register the given generic signature builder to be used as the canonical
@@ -919,6 +914,9 @@ class ASTContext final {
   CanGenericSignature getExistentialSignature(CanType existential,
                                               ModuleDecl *mod);
 
+  GenericSignature *getOverrideGenericSignature(ValueDecl *base,
+                                                ValueDecl *derived);
+
   /// Whether our effective Swift version is at least 'major'.
   ///
   /// This is usually the check you want; for example, when introducing
@@ -956,4 +954,39 @@ class ASTContext final {
 
 } // end namespace swift
 
+namespace llvm {
+template <> struct DenseMapInfo<swift::ASTContext::OverrideSignatureKey> {
+  using OverrideSignatureKey = swift::ASTContext::OverrideSignatureKey;
+  using Type = swift::Type;
+  using GenericSignature = swift::GenericSignature;
+
+  static bool isEqual(const OverrideSignatureKey lhs,
+                      const OverrideSignatureKey rhs) {
+    return lhs.baseMethodSig == rhs.baseMethodSig &&
+           lhs.derivedClassSig == rhs.derivedClassSig &&
+           lhs.superclassTy.getString() == rhs.superclassTy.getString();
+  }
+
+  static inline OverrideSignatureKey getEmptyKey() {
+    return OverrideSignatureKey(DenseMapInfo<GenericSignature *>::getEmptyKey(),
+                                DenseMapInfo<GenericSignature *>::getEmptyKey(),
+                                DenseMapInfo<Type>::getEmptyKey());
+  }
+
+  static inline OverrideSignatureKey getTombstoneKey() {
+    return OverrideSignatureKey(
+        DenseMapInfo<GenericSignature *>::getTombstoneKey(),
+        DenseMapInfo<GenericSignature *>::getTombstoneKey(),
+        DenseMapInfo<Type>::getTombstoneKey());
+  }
+
+  static unsigned getHashValue(const OverrideSignatureKey &Val) {
+    return hash_combine(
+        DenseMapInfo<GenericSignature *>::getHashValue(Val.baseMethodSig),
+        DenseMapInfo<GenericSignature *>::getHashValue(Val.derivedClassSig),
+        DenseMapInfo<Type>::getHashValue(Val.superclassTy));
+  }
+};
+} // namespace llvm
+
 #endif

lib/AST/ASTContext.cpp

@@ -4340,6 +4340,107 @@ CanGenericSignature ASTContext::getExistentialSignature(CanType existential,
   return genericSig;
 }
 
+GenericSignature *ASTContext::getOverrideGenericSignature(ValueDecl *base,
+                                                          ValueDecl *derived) {
+  auto baseGenericCtx = base->getAsGenericContext();
+  auto &ctx = base->getASTContext();
+
+  if (!baseGenericCtx) {
+    return nullptr;
+  }
+
+  auto baseClass = base->getDeclContext()->getSelfClassDecl();
+
+  if (!baseClass) {
+    return nullptr;
+  }
+
+  auto derivedClass = derived->getDeclContext()->getSelfClassDecl();
+  auto *baseClassSig = baseClass->getGenericSignature();
+
+  if (!derivedClass) {
+    return nullptr;
+  }
+
+  if (derivedClass->getSuperclass().isNull()) {
+    return nullptr;
+  }
+
+  if (derivedClass->getGenericSignature() == nullptr &&
+      !baseGenericCtx->isGeneric()) {
+    return nullptr;
+  }
+
+  auto subMap = derivedClass->getSuperclass()->getContextSubstitutionMap(
+      derivedClass->getModuleContext(), baseClass);
+
+  if (baseGenericCtx->getGenericSignature() == nullptr) {
+    return nullptr;
+  }
+  unsigned derivedDepth = 0;
+
+  auto key = OverrideSignatureKey(baseGenericCtx->getGenericSignature(),
+                                  derivedClass->getGenericSignature(),
+                                  derivedClass->getSuperclass());
+
+  if (overrideSigCache.count(key) == 1) {
+    return overrideSigCache.lookup(key);
+  }
+
+  if (auto *derivedSig = derivedClass->getGenericSignature())
+    derivedDepth = derivedSig->getGenericParams().back()->getDepth() + 1;
+
+  GenericSignatureBuilder builder(ctx);
+  builder.addGenericSignature(derivedClass->getGenericSignature());
+
+  if (auto derivedGenericCtx = derived->getAsGenericContext()) {
+    if (derivedGenericCtx->isGeneric()) {
+      for (auto param : *derivedGenericCtx->getGenericParams()) {
+        builder.addGenericParameter(param);
+      }
+    }
+  }
+
+  auto source =
+      GenericSignatureBuilder::FloatingRequirementSource::forAbstract();
+
+  unsigned baseDepth = 0;
+
+  if (baseClassSig) {
+    baseDepth = baseClassSig->getGenericParams().back()->getDepth() + 1;
+  }
+
+  auto substFn = [&](SubstitutableType *type) -> Type {
+    auto *gp = cast<GenericTypeParamType>(type);
+
+    if (gp->getDepth() < baseDepth) {
+      return Type(gp).subst(subMap);
+    }
+
+    return CanGenericTypeParamType::get(
+        gp->getDepth() - baseDepth + derivedDepth, gp->getIndex(), ctx);
+  };
+
+  auto lookupConformanceFn =
+      [&](CanType depTy, Type substTy,
+          ProtocolDecl *proto) -> Optional<ProtocolConformanceRef> {
+    if (auto conf = subMap.lookupConformance(depTy, proto))
+      return conf;
+
+    return ProtocolConformanceRef(proto);
+  };
+
+  for (auto reqt : baseGenericCtx->getGenericSignature()->getRequirements()) {
+    if (auto substReqt = reqt.subst(substFn, lookupConformanceFn)) {
+      builder.addRequirement(*substReqt, source, nullptr);
+    }
+  }
+
+  auto *genericSig = std::move(builder).computeGenericSignature(SourceLoc());
+  overrideSigCache.insert(std::make_pair(key, genericSig));
+  return genericSig;
+}
+
 SILLayout *SILLayout::get(ASTContext &C,
                           CanGenericSignature Generics,
                           ArrayRef<SILField> Fields) {

lib/Sema/TypeCheckDeclOverride.cpp

@@ -931,110 +931,6 @@ static void checkOverrideAccessControl(ValueDecl *baseDecl, ValueDecl *decl,
   }
 }
 
-static GenericSignature *getOverrideGenericSignature(ValueDecl *base,
-                                                     ValueDecl *derived) {
-  auto baseGenericCtx = base->getAsGenericContext();
-  auto &ctx = base->getASTContext();
-
-  if (!baseGenericCtx) {
-    return nullptr;
-  }
-
-  auto baseClass = base->getDeclContext()->getSelfClassDecl();
-
-  if (!baseClass) {
-    return nullptr;
-  }
-
-  auto derivedClass = derived->getDeclContext()->getSelfClassDecl();
-  auto *baseClassSig = baseClass->getGenericSignature();
-
-  if (!derivedClass) {
-    return nullptr;
-  }
-
-  if (derivedClass->getSuperclass().isNull()) {
-    return nullptr;
-  }
-
-  if (derivedClass->getGenericSignature() == nullptr &&
-      !baseGenericCtx->isGeneric()) {
-    return nullptr;
-  }
-
-  auto subMap = derivedClass->getSuperclass()->getContextSubstitutionMap(
-      derivedClass->getModuleContext(), baseClass);
-
-  if (baseGenericCtx->getGenericSignature() == nullptr) {
-    return nullptr;
-  }
-  unsigned derivedDepth = 0;
-
-  auto key = swift::ASTContext::OverrideSignatureKey(
-      baseGenericCtx->getGenericSignature()->getAsString(),
-      derivedClass->getGenericSignature()->getAsString(),
-      derivedClass->getSuperclass());
-
-  for (auto pair : ctx.overrideSigCache) {
-    if (pair.first.isEqual(key)) {
-      return pair.second;
-    }
-  }
-
-  if (auto *derivedSig = derivedClass->getGenericSignature())
-    derivedDepth = derivedSig->getGenericParams().back()->getDepth() + 1;
-
-  GenericSignatureBuilder builder(ctx);
-  builder.addGenericSignature(derivedClass->getGenericSignature());
-
-  if (auto derivedGenericCtx = derived->getAsGenericContext()) {
-    if (derivedGenericCtx->isGeneric()) {
-      for (auto param : *derivedGenericCtx->getGenericParams()) {
-        builder.addGenericParameter(param);
-      }
-    }
-  }
-
-  auto source =
-      GenericSignatureBuilder::FloatingRequirementSource::forAbstract();
-
-  unsigned baseDepth = 0;
-
-  if (baseClassSig) {
-    baseDepth = baseClassSig->getGenericParams().back()->getDepth() + 1;
-  }
-
-  auto substFn = [&](SubstitutableType *type) -> Type {
-    auto *gp = cast<GenericTypeParamType>(type);
-
-    if (gp->getDepth() < baseDepth) {
-      return Type(gp).subst(subMap);
-    }
-
-    return CanGenericTypeParamType::get(
-        gp->getDepth() - baseDepth + derivedDepth, gp->getIndex(), ctx);
-  };
-
-  auto lookupConformanceFn =
-      [&](CanType depTy, Type substTy,
-          ProtocolDecl *proto) -> Optional<ProtocolConformanceRef> {
-    if (auto conf = subMap.lookupConformance(depTy, proto))
-      return conf;
-
-    return ProtocolConformanceRef(proto);
-  };
-
-  for (auto reqt : baseGenericCtx->getGenericSignature()->getRequirements()) {
-    if (auto substReqt = reqt.subst(substFn, lookupConformanceFn)) {
-      builder.addRequirement(*substReqt, source, nullptr);
-    }
-  }
-
-  auto *genericSig = std::move(builder).computeGenericSignature(SourceLoc());
-  ctx.overrideSigCache.push_back(std::make_pair(key, genericSig));
-  return genericSig;
-}
-
 bool OverrideMatcher::checkOverride(ValueDecl *baseDecl,
                                     OverrideCheckingAttempt attempt) {
   auto &diags = ctx.Diags;
@@ -1057,7 +953,7 @@ bool OverrideMatcher::checkOverride(ValueDecl *baseDecl,
 
   if (baseGenericCtx && derivedGenericCtx) {
     // If the generic signatures are different, then complain
-    if (auto newSig = getOverrideGenericSignature(baseDecl, decl)) {
+    if (auto newSig = ctx.getOverrideGenericSignature(baseDecl, decl)) {
       if (auto derivedSig = derivedGenericCtx->getGenericSignature()) {
         auto requirementsSatisfied =
             derivedSig->requirementsNotSatisfiedBy(newSig).empty();