Merge pull request #38868 from slavapestov/trie-overlapping-rules

RequirementMachine: Some more performance improvements

Author: slavapestov

lib/AST/RequirementMachine/RequirementMachine.cpp

@@ -389,10 +389,6 @@ void RequirementMachine::computeCompletion() {
 
     checkCompletionResult();
 
-    // Simplify right hand sides in preparation for building the
-    // property map.
-    System.simplifyRewriteSystem();
-
     // Check invariants.
     System.verify();
 

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

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();
 };
 

lib/AST/RequirementMachine/RewriteSystem.cpp

@@ -127,28 +127,6 @@ bool RewriteSystem::addRule(MutableTerm lhs, MutableTerm rhs) {
     MergedAssociatedTypes.emplace_back(lhs, rhs);
   }
 
-  // Since we added a new rule, we have to check for overlaps between the
-  // new rule and all existing rules.
-  for (unsigned j : indices(Rules)) {
-    // A rule does not overlap with itself.
-    if (i == j)
-      continue;
-
-    // We don't have to check for overlap with deleted rules.
-    if (Rules[j].isDeleted())
-      continue;
-
-    // The overlap check is not commutative so we have to check both
-    // directions.
-    Worklist.emplace_back(i, j);
-    Worklist.emplace_back(j, i);
-
-    if (DebugCompletion) {
-      llvm::dbgs() << "$ Queued up (" << i << ", " << j << ") and ";
-      llvm::dbgs() << "(" << j << ", " << i << ")\n";
-    }
-  }
-
   // Tell the caller that we added a new rule.
   return true;
 }

lib/AST/RequirementMachine/RewriteSystem.h

@@ -13,7 +13,7 @@
 #ifndef SWIFT_REWRITESYSTEM_H
 #define SWIFT_REWRITESYSTEM_H
 
-#include <algorithm>
+#include "llvm/ADT/DenseSet.h"
 
 #include "ProtocolGraph.h"
 #include "Symbol.h"
@@ -48,12 +48,6 @@ class Rule final {
   const Term &getLHS() const { return LHS; }
   const Term &getRHS() const { return RHS; }
 
-  OverlapKind checkForOverlap(const Rule &other,
-                              MutableTerm &t,
-                              MutableTerm &v) const {
-    return LHS.checkForOverlap(other.LHS, t, v);
-  }
-
   /// Returns if the rule was deleted.
   bool isDeleted() const {
     return deleted;
@@ -114,9 +108,8 @@ class RewriteSystem final {
   /// See RewriteSystem::processMergedAssociatedTypes() for details.
   std::vector<std::pair<MutableTerm, MutableTerm>> MergedAssociatedTypes;
 
-  /// A list of pending pairs for checking overlap in the completion
-  /// procedure.
-  std::deque<std::pair<unsigned, unsigned>> Worklist;
+  /// Pairs of rules which have already been checked for overlap.
+  llvm::DenseSet<std::pair<unsigned, unsigned>> CheckedOverlaps;
 
   /// Set these to true to enable debugging output.
   unsigned DebugSimplify : 1;
@@ -176,8 +169,9 @@ class RewriteSystem final {
   void dump(llvm::raw_ostream &out) const;
 
 private:
-  Optional<std::pair<MutableTerm, MutableTerm>>
-  computeCriticalPair(const Rule &lhs, const Rule &rhs) const;
+  std::pair<MutableTerm, MutableTerm>
+  computeCriticalPair(ArrayRef<Symbol>::const_iterator from,
+                      const Rule &lhs, const Rule &rhs) const;
 
   Symbol mergeAssociatedTypes(Symbol lhs, Symbol rhs) const;
   void processMergedAssociatedTypes();