RequirementMachine: Split up list of associated types in ProtocolGraph

Store the protocol's direct associated types separately from the inherited
associated types, since in a couple of places we only need the direct
associated types.

Also, factor out a new ProtocolGraph::compute() method that does all the
steps in the right order.

Author: slavapestov

lib/AST/RequirementMachine/EquivalenceClassMap.cpp

@@ -572,19 +572,7 @@ void EquivalenceClassMap::concretizeNestedTypesFromConcreteParent(
 
     auto *concrete = conformance.getConcrete();
 
-    // We might have duplicates in the list due to diamond inheritance.
-    // FIXME: Filter those out further upstream?
-    // FIXME: This should actually be outside of the loop over the conforming protos...
-    llvm::SmallDenseSet<AssociatedTypeDecl *, 4> visited;
     for (auto *assocType : assocTypes) {
-      if (!visited.insert(assocType).second)
-        continue;
-
-      // Get the actual protocol in case we inherited this associated type.
-      auto *actualProto = assocType->getProtocol();
-      if (actualProto != proto)
-        continue;
-
       if (DebugConcretizeNestedTypes) {
         llvm::dbgs() << "^^ " << "Looking up type witness for "
                      << proto->getName() << ":" << assocType->getName()

lib/AST/RequirementMachine/ProtocolGraph.cpp

@@ -109,24 +109,9 @@ void ProtocolGraph::computeInheritedAssociatedTypes() {
   for (const auto *proto : llvm::reverse(Protocols)) {
     auto &info = Info[proto];
 
-    // We might inherit the same associated type multiple times due to
-    // diamond inheritance, so make sure we only add each associated
-    // type once.
-    llvm::SmallDenseSet<const AssociatedTypeDecl *, 4> visited;
-
-    for (const auto *inherited : info.Inherited) {
-      if (inherited == proto)
-        continue;
-
+    for (const auto *inherited : info.AllInherited) {
       for (auto *inheritedType : getProtocolInfo(inherited).AssociatedTypes) {
-        if (!visited.insert(inheritedType).second)
-          continue;
-
-        // The 'if (inherited == proto)' above avoids a potential
-        // iterator invalidation here, because we're updating
-        // getProtocolInfo(proto).AssociatedTypes while iterating over
-        // getProtocolInfo(inherited).AssociatedTypes.
-        info.AssociatedTypes.push_back(inheritedType);
+        info.InheritedAssociatedTypes.push_back(inheritedType);
       }
     }
   }
@@ -190,6 +175,14 @@ unsigned ProtocolGraph::computeProtocolDepth(const ProtocolDecl *proto) {
   return depth;
 }
 
+/// Compute everything in the right order.
+void ProtocolGraph::compute() {
+  computeTransitiveClosure();
+  computeLinearOrder();
+  computeInheritedProtocols();
+  computeInheritedAssociatedTypes();
+}
+
 /// Defines a linear order with the property that if a protocol P inherits
 /// from another protocol Q, then P < Q. (The converse cannot be true, since
 /// this is a linear order.)

lib/AST/RequirementMachine/ProtocolGraph.h

@@ -34,11 +34,14 @@ struct ProtocolInfo {
   /// itself. Computed by ProtocolGraph::computeInheritedProtocols().
   llvm::TinyPtrVector<const ProtocolDecl *> AllInherited;
 
-  /// Transitive closure of inherited associated types together with all
-  /// associated types from the protocol itself. Computed by
-  /// ProtocolGraph::computeInheritedAssociatedTypes().
+  /// Associated types defined in the protocol itself.
   llvm::TinyPtrVector<AssociatedTypeDecl *> AssociatedTypes;
 
+  /// Associated types from all inherited protocols, not including duplicates or
+  /// those defined in the protocol itself. Computed by
+  /// ProtocolGraph::computeInheritedAssociatedTypes().
+  llvm::TinyPtrVector<AssociatedTypeDecl *> InheritedAssociatedTypes;
+
   /// The protocol's requirement signature.
   ArrayRef<Requirement> Requirements;
 
@@ -76,28 +79,36 @@ struct ProtocolInfo {
 /// referenced from a set of generic requirements.
 ///
 /// Out-of-line methods are documented in ProtocolGraph.cpp.
-struct ProtocolGraph {
+class ProtocolGraph {
   llvm::DenseMap<const ProtocolDecl *, ProtocolInfo> Info;
   std::vector<const ProtocolDecl *> Protocols;
   bool Debug = false;
 
+public:
   void visitRequirements(ArrayRef<Requirement> reqs);
 
   bool isKnownProtocol(const ProtocolDecl *proto) const;
 
+  /// Returns the sorted list of protocols, with the property
+  /// that (P refines Q) => P < Q. See compareProtocols()
+  /// for details.
+  ArrayRef<const ProtocolDecl *> getProtocols() const {
+    return Protocols;
+  }
+
   const ProtocolInfo &getProtocolInfo(
       const ProtocolDecl *proto) const;
 
+private:
   void addProtocol(const ProtocolDecl *proto);
-
   void computeTransitiveClosure();
-
   void computeLinearOrder();
-
   void computeInheritedAssociatedTypes();
-
   void computeInheritedProtocols();
 
+public:
+  void compute();
+
   int compareProtocols(const ProtocolDecl *lhs,
                        const ProtocolDecl *rhs) const;
 

lib/AST/RequirementMachine/RequirementMachine.cpp

@@ -86,21 +86,21 @@ void RewriteSystemBuilder::addGenericSignature(CanGenericSignature sig) {
   // Collect all protocols transitively referenced from the generic signature's
   // requirements.
   Protocols.visitRequirements(sig->getRequirements());
-  Protocols.computeTransitiveClosure();
-  Protocols.computeLinearOrder();
-  Protocols.computeInheritedProtocols();
-  Protocols.computeInheritedAssociatedTypes();
+  Protocols.compute();
 
   // Add rewrite rules for each protocol.
-  for (auto *proto : Protocols.Protocols) {
+  for (auto *proto : Protocols.getProtocols()) {
     if (Debug) {
       llvm::dbgs() << "protocol " << proto->getName() << " {\n";
     }
 
     const auto &info = Protocols.getProtocolInfo(proto);
 
-    for (auto *type : info.AssociatedTypes)
-      addAssociatedType(type, proto);
+    for (auto *assocType : info.AssociatedTypes)
+      addAssociatedType(assocType, proto);
+
+    for (auto *assocType : info.InheritedAssociatedTypes)
+      addAssociatedType(assocType, proto);
 
     for (auto req : info.Requirements)
       addRequirement(req.getCanonical(), proto);

lib/AST/RequirementMachine/RewriteSystem.cpp

@@ -1051,7 +1051,7 @@ Type getTypeForAtomRange(Iter begin, Iter end, Type root,
       //
       for (auto *proto : atom.getProtocols()) {
         const auto &info = protos.getProtocolInfo(proto);
-        for (auto *otherAssocType : info.AssociatedTypes) {
+        auto checkOtherAssocType = [&](AssociatedTypeDecl *otherAssocType) {
           otherAssocType = otherAssocType->getAssociatedTypeAnchor();
 
           if (otherAssocType->getName() == name &&
@@ -1060,6 +1060,14 @@ Type getTypeForAtomRange(Iter begin, Iter end, Type root,
                                  assocType->getProtocol()) < 0)) {
             assocType = otherAssocType;
           }
+        };
+
+        for (auto *otherAssocType : info.AssociatedTypes) {
+          checkOtherAssocType(otherAssocType);
+        }
+
+        for (auto *otherAssocType : info.InheritedAssociatedTypes) {
+          checkOtherAssocType(otherAssocType);
         }
       }
     }