RequirementMachine: Cache the mapping from associated type symbols to associated type declarations

slavapestov · slavapestov · commit eec233507be0 · 2021-08-13T17:27:49.000-04:00
diff --git a/lib/AST/RequirementMachine/RewriteContext.cpp b/lib/AST/RequirementMachine/RewriteContext.cpp
@@ -126,6 +126,79 @@ MutableTerm RewriteContext::getMutableTermForType(CanType paramType,
   return MutableTerm(symbols);
 }
 
+/// Map an associated type symbol to an associated type declaration.
+///
+/// Note that the protocol graph is not part of the caching key; each
+/// protocol graph is a subgraph of the global inheritance graph, so
+/// the specific choice of subgraph does not change the result.
+AssociatedTypeDecl *RewriteContext::getAssociatedTypeForSymbol(
+    Symbol symbol, const ProtocolGraph &protos) {
+  auto found = AssocTypes.find(symbol);
+  if (found != AssocTypes.end())
+    return found->second;
+
+  assert(symbol.getKind() == Symbol::Kind::AssociatedType);
+  auto *proto = symbol.getProtocols()[0];
+  auto name = symbol.getName();
+
+  AssociatedTypeDecl *assocType = nullptr;
+
+  // Special case: handle unknown protocols, since they can appear in the
+  // invalid types that getCanonicalTypeInContext() must handle via
+  // concrete substitution; see the definition of getCanonicalTypeInContext()
+  // below for details.
+  if (!protos.isKnownProtocol(proto)) {
+    assert(symbol.getProtocols().size() == 1 &&
+           "Unknown associated type symbol must have a single protocol");
+    assocType = proto->getAssociatedType(name)->getAssociatedTypeAnchor();
+  } else {
+    // An associated type symbol [P1&P1&...&Pn:A] has one or more protocols
+    // P0...Pn and an identifier 'A'.
+    //
+    // We map it back to a AssociatedTypeDecl as follows:
+    //
+    // - For each protocol Pn, look for associated types A in Pn itself,
+    //   and all protocols that Pn refines.
+    //
+    // - For each candidate associated type An in protocol Qn where
+    //   Pn refines Qn, get the associated type anchor An' defined in
+    //   protocol Qn', where Qn refines Qn'.
+    //
+    // - Out of all the candidiate pairs (Qn', An'), pick the one where
+    //   the protocol Qn' is the lowest element according to the linear
+    //   order defined by TypeDecl::compare().
+    //
+    // The associated type An' is then the canonical associated type
+    // representative of the associated type symbol [P0&...&Pn:A].
+    //
+    for (auto *proto : symbol.getProtocols()) {
+      const auto &info = protos.getProtocolInfo(proto);
+      auto checkOtherAssocType = [&](AssociatedTypeDecl *otherAssocType) {
+        otherAssocType = otherAssocType->getAssociatedTypeAnchor();
+
+        if (otherAssocType->getName() == name &&
+            (assocType == nullptr ||
+             TypeDecl::compare(otherAssocType->getProtocol(),
+                               assocType->getProtocol()) < 0)) {
+          assocType = otherAssocType;
+        }
+      };
+
+      for (auto *otherAssocType : info.AssociatedTypes) {
+        checkOtherAssocType(otherAssocType);
+      }
+
+      for (auto *otherAssocType : info.InheritedAssociatedTypes) {
+        checkOtherAssocType(otherAssocType);
+      }
+    }
+  }
+
+  assert(assocType && "Need to look harder");
+  AssocTypes[symbol] = assocType;
+  return assocType;
+}
+
 /// Compute the interface type for a range of symbols, with an optional
 /// root type.
 ///
@@ -136,7 +209,7 @@ template<typename Iter>
 Type getTypeForSymbolRange(Iter begin, Iter end, Type root,
                            TypeArrayView<GenericTypeParamType> genericParams,
                            const ProtocolGraph &protos,
-                           ASTContext &ctx) {
+                           const RewriteContext &ctx) {
   Type result = root;
 
   auto handleRoot = [&](GenericTypeParamType *genericParam) {
@@ -166,11 +239,11 @@ Type getTypeForSymbolRange(Iter begin, Iter end, Type root,
         continue;
 
       case Symbol::Kind::Protocol:
-        handleRoot(GenericTypeParamType::get(0, 0, ctx));
+        handleRoot(GenericTypeParamType::get(0, 0, ctx.getASTContext()));
         continue;
 
       case Symbol::Kind::AssociatedType:
-        handleRoot(GenericTypeParamType::get(0, 0, ctx));
+        handleRoot(GenericTypeParamType::get(0, 0, ctx.getASTContext()));
 
         // An associated type term at the root means we have a dependent
         // member type rooted at Self; handle the associated type below.
@@ -191,68 +264,9 @@ Type getTypeForSymbolRange(Iter begin, Iter end, Type root,
     }
 
     // We should have a resolved type at this point.
-    assert(symbol.getKind() == Symbol::Kind::AssociatedType);
-    auto *proto = symbol.getProtocols()[0];
-    auto name = symbol.getName();
-
-    AssociatedTypeDecl *assocType = nullptr;
-
-    // Special case: handle unknown protocols, since they can appear in the
-    // invalid types that getCanonicalTypeInContext() must handle via
-    // concrete substitution; see the definition of getCanonicalTypeInContext()
-    // below for details.
-    if (!protos.isKnownProtocol(proto)) {
-      assert(root &&
-             "We only allow unknown protocols in getRelativeTypeForTerm()");
-      assert(symbol.getProtocols().size() == 1 &&
-             "Unknown associated type symbol must have a single protocol");
-      assocType = proto->getAssociatedType(name)->getAssociatedTypeAnchor();
-    } else {
-      // FIXME: Cache this
-      //
-      // An associated type symbol [P1&P1&...&Pn:A] has one or more protocols
-      // P0...Pn and an identifier 'A'.
-      //
-      // We map it back to a AssociatedTypeDecl as follows:
-      //
-      // - For each protocol Pn, look for associated types A in Pn itself,
-      //   and all protocols that Pn refines.
-      //
-      // - For each candidate associated type An in protocol Qn where
-      //   Pn refines Qn, get the associated type anchor An' defined in
-      //   protocol Qn', where Qn refines Qn'.
-      //
-      // - Out of all the candidiate pairs (Qn', An'), pick the one where
-      //   the protocol Qn' is the lowest element according to the linear
-      //   order defined by TypeDecl::compare().
-      //
-      // The associated type An' is then the canonical associated type
-      // representative of the associated type symbol [P0&...&Pn:A].
-      //
-      for (auto *proto : symbol.getProtocols()) {
-        const auto &info = protos.getProtocolInfo(proto);
-        auto checkOtherAssocType = [&](AssociatedTypeDecl *otherAssocType) {
-          otherAssocType = otherAssocType->getAssociatedTypeAnchor();
-
-          if (otherAssocType->getName() == name &&
-              (assocType == nullptr ||
-               TypeDecl::compare(otherAssocType->getProtocol(),
-                                 assocType->getProtocol()) < 0)) {
-            assocType = otherAssocType;
-          }
-        };
-
-        for (auto *otherAssocType : info.AssociatedTypes) {
-          checkOtherAssocType(otherAssocType);
-        }
-
-        for (auto *otherAssocType : info.InheritedAssociatedTypes) {
-          checkOtherAssocType(otherAssocType);
-        }
-      }
-    }
-
-    assert(assocType && "Need to look harder");
+    auto *assocType =
+        const_cast<RewriteContext &>(ctx)
+            .getAssociatedTypeForSymbol(symbol, protos);
     result = DependentMemberType::get(result, assocType);
   }
 
@@ -263,14 +277,14 @@ Type RewriteContext::getTypeForTerm(Term term,
                       TypeArrayView<GenericTypeParamType> genericParams,
                       const ProtocolGraph &protos) const {
   return getTypeForSymbolRange(term.begin(), term.end(), Type(),
-                               genericParams, protos, Context);
+                               genericParams, protos, *this);
 }
 
 Type RewriteContext::getTypeForTerm(const MutableTerm &term,
                       TypeArrayView<GenericTypeParamType> genericParams,
                       const ProtocolGraph &protos) const {
   return getTypeForSymbolRange(term.begin(), term.end(), Type(),
-                               genericParams, protos, Context);
+                               genericParams, protos, *this);
 }
 
 Type RewriteContext::getRelativeTypeForTerm(
@@ -281,7 +295,7 @@ Type RewriteContext::getRelativeTypeForTerm(
   auto genericParam = CanGenericTypeParamType::get(0, 0, Context);
   return getTypeForSymbolRange(
       term.begin() + prefix.size(), term.end(), genericParam,
-      { }, protos, Context);
+      { }, protos, *this);
 }
 
 /// We print stats in the destructor, which should get executed at the end of
diff --git a/lib/AST/RequirementMachine/RewriteContext.h b/lib/AST/RequirementMachine/RewriteContext.h
@@ -16,6 +16,7 @@
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/Types.h"
 #include "swift/Basic/Statistic.h"
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/FoldingSet.h"
 #include "llvm/Support/Allocator.h"
 #include "Histogram.h"
@@ -44,6 +45,9 @@ class RewriteContext final {
   /// Folding set for uniquing terms.
   llvm::FoldingSet<Term::Storage> Terms;
 
+  /// Cache for associated type declarations.
+  llvm::DenseMap<Symbol, AssociatedTypeDecl *> AssocTypes;
+
   RewriteContext(const RewriteContext &) = delete;
   RewriteContext(RewriteContext &&) = delete;
   RewriteContext &operator=(const RewriteContext &) = delete;
@@ -70,7 +74,7 @@ class RewriteContext final {
   MutableTerm getMutableTermForType(CanType paramType,
                                     const ProtocolDecl *proto);
 
-  ASTContext &getASTContext() { return Context; }
+  ASTContext &getASTContext() const { return Context; }
 
   Type getTypeForTerm(Term term,
                       TypeArrayView<GenericTypeParamType> genericParams,
@@ -84,6 +88,9 @@ class RewriteContext final {
                       const MutableTerm &term, const MutableTerm &prefix,
                       const ProtocolGraph &protos) const;
 
+  AssociatedTypeDecl *getAssociatedTypeForSymbol(Symbol symbol,
+                                                 const ProtocolGraph &protos);
+
   ~RewriteContext();
 };
 
