RequirementMachine: Refactor RuleBuilder in preparation for rule sharing

Author: slavapestov

lib/AST/RequirementMachine/ConcreteTypeWitness.cpp

@@ -562,8 +562,8 @@ void PropertyMap::inferConditionalRequirements(
           llvm::dbgs() << "@@@ Unknown protocol: "<< proto->getName() << "\n";
         }
 
-        RuleBuilder builder(Context, System.getProtocolMap());
-        builder.addProtocol(proto, /*initialComponent=*/false);
+        RuleBuilder builder(Context, System.getReferencedProtocols());
+        builder.addReferencedProtocol(proto);
         builder.collectRulesFromReferencedProtocols();
 
         for (const auto &rule : builder.PermanentRules)

lib/AST/RequirementMachine/RequirementLowering.cpp

@@ -1051,7 +1051,7 @@ void RuleBuilder::addRequirements(ArrayRef<Requirement> requirements) {
   // Collect all protocols transitively referenced from these requirements.
   for (auto req : requirements) {
     if (req.getKind() == RequirementKind::Conformance) {
-      addProtocol(req.getProtocolDecl(), /*initialComponent=*/false);
+      addReferencedProtocol(req.getProtocolDecl());
     }
   }
 
@@ -1066,7 +1066,7 @@ void RuleBuilder::addRequirements(ArrayRef<StructuralRequirement> requirements)
   // Collect all protocols transitively referenced from these requirements.
   for (auto req : requirements) {
     if (req.req.getKind() == RequirementKind::Conformance) {
-      addProtocol(req.req.getProtocolDecl(), /*initialComponent=*/false);
+      addReferencedProtocol(req.req.getProtocolDecl());
     }
   }
 
@@ -1077,16 +1077,65 @@ void RuleBuilder::addRequirements(ArrayRef<StructuralRequirement> requirements)
     addRequirement(req, /*proto=*/nullptr);
 }
 
+/// For building a rewrite system for a protocol connected component from
+/// user-written requirements. Used when actually building requirement
+/// signatures.
 void RuleBuilder::addProtocols(ArrayRef<const ProtocolDecl *> protos) {
-  // Collect all protocols transitively referenced from this connected component
-  // of the protocol dependency graph.
-  for (auto proto : protos) {
-    addProtocol(proto, /*initialComponent=*/true);
+  for (auto *proto : protos) {
+    ReferencedProtocols.insert(proto);
+  }
+
+  for (auto *proto : protos) {
+    if (Dump) {
+      llvm::dbgs() << "protocol " << proto->getName() << " {\n";
+    }
+
+    addPermanentProtocolRules(proto);
+
+    for (auto req : proto->getStructuralRequirements())
+      addRequirement(req, proto);
+
+    for (auto req : proto->getTypeAliasRequirements())
+      addRequirement(req.getCanonical(), proto, /*requirementID=*/None);
+
+    for (auto *otherProto : proto->getProtocolDependencies())
+      addReferencedProtocol(otherProto);
+
+    if (Dump) {
+      llvm::dbgs() << "}\n";
+    }
   }
 
+  // Collect all protocols transitively referenced from this connected component
+  // of the protocol dependency graph.
   collectRulesFromReferencedProtocols();
 }
 
+/// Add permanent rules for a protocol, consisting of:
+///
+/// - The identity conformance rule [P].[P] => [P].
+/// - An associated type introduction rule for each associated type.
+/// - An inherited associated type introduction rule for each associated
+///   type of each inherited protocol.
+void RuleBuilder::addPermanentProtocolRules(const ProtocolDecl *proto) {
+  MutableTerm lhs;
+  lhs.add(Symbol::forProtocol(proto, Context));
+  lhs.add(Symbol::forProtocol(proto, Context));
+
+  MutableTerm rhs;
+  rhs.add(Symbol::forProtocol(proto, Context));
+
+  PermanentRules.emplace_back(lhs, rhs);
+
+  for (auto *assocType : proto->getAssociatedTypeMembers())
+    addAssociatedType(assocType, proto);
+
+  for (auto *inheritedProto : Context.getInheritedProtocols(proto)) {
+    for (auto *assocType : inheritedProto->getAssociatedTypeMembers())
+      addAssociatedType(assocType, proto);
+  }
+}
+
 /// For an associated type T in a protocol P, we add a rewrite rule:
 ///
 ///   [P].T => [P:T]
@@ -1264,72 +1313,44 @@ void RuleBuilder::addTypeAlias(const ProtocolTypeAlias &alias,
                                 /*requirementID=*/None);
 }
 
-/// Record information about a protocol if we have no seen it yet.
-void RuleBuilder::addProtocol(const ProtocolDecl *proto,
-                              bool initialComponent) {
-  if (ProtocolMap.count(proto) > 0)
-    return;
-
-  ProtocolMap[proto] = initialComponent;
-  Protocols.push_back(proto);
+/// If we haven't seen this protocol yet, save it for later so that we can
+/// import the rewrite rules from its connected component.
+void RuleBuilder::addReferencedProtocol(const ProtocolDecl *proto) {
+  if (ReferencedProtocols.insert(proto).second)
+    ProtocolsToImport.push_back(proto);
 }
 
 /// Compute the transitive closure of the set of all protocols referenced from
 /// the right hand sides of conformance requirements, and convert their
 /// requirements to rewrite rules.
 void RuleBuilder::collectRulesFromReferencedProtocols() {
+  // Compute the transitive closure.
   unsigned i = 0;
-  while (i < Protocols.size()) {
-    auto *proto = Protocols[i++];
+  while (i < ProtocolsToImport.size()) {
+    auto *proto = ProtocolsToImport[i++];
     for (auto *depProto : proto->getProtocolDependencies()) {
-      addProtocol(depProto, /*initialComponent=*/false);
+      addReferencedProtocol(depProto);
     }
   }
 
-  // Add rewrite rules for each protocol.
-  for (auto *proto : Protocols) {
+  // If this is a rewrite system for a generic signature, add rewrite rules for
+  // each referenced protocol.
+  //
+  // if this is a rewrite system for a connected component of the protocol
+  // dependency graph, add rewrite rules for each referenced protocol not part
+  // of this connected component.
+  for (auto *proto : ProtocolsToImport) {
     if (Dump) {
       llvm::dbgs() << "protocol " << proto->getName() << " {\n";
     }
 
-    // Add the identity conformance rule [P].[P] => [P].
-    MutableTerm lhs;
-    lhs.add(Symbol::forProtocol(proto, Context));
-    lhs.add(Symbol::forProtocol(proto, Context));
-
-    MutableTerm rhs;
-    rhs.add(Symbol::forProtocol(proto, Context));
-
-    PermanentRules.emplace_back(lhs, rhs);
-
-    for (auto *assocType : proto->getAssociatedTypeMembers())
-      addAssociatedType(assocType, proto);
-
-    for (auto *inheritedProto : Context.getInheritedProtocols(proto)) {
-      for (auto *assocType : inheritedProto->getAssociatedTypeMembers())
-        addAssociatedType(assocType, proto);
-    }
+    addPermanentProtocolRules(proto);
 
-    // If this protocol is part of the initial connected component, we're
-    // building requirement signatures for all protocols in this component,
-    // and so we must start with the structural requirements.
-    //
-    // Otherwise, we should either already have a requirement signature, or
-    // we can trigger the computation of the requirement signatures of the
-    // next component recursively.
-    if (ProtocolMap[proto]) {
-      for (auto req : proto->getStructuralRequirements())
-        addRequirement(req, proto);
-
-      for (auto req : proto->getTypeAliasRequirements())
-        addRequirement(req.getCanonical(), proto, /*requirementID=*/None);
-    } else {
-      auto reqs = proto->getRequirementSignature();
-      for (auto req : reqs.getRequirements())
-        addRequirement(req.getCanonical(), proto, /*requirementID=*/None);
-      for (auto alias : reqs.getTypeAliases())
-        addTypeAlias(alias, proto);
-    }
+    auto reqs = proto->getRequirementSignature();
+    for (auto req : reqs.getRequirements())
+      addRequirement(req.getCanonical(), proto, /*requirementID=*/None);
+    for (auto alias : reqs.getTypeAliases())
+      addTypeAlias(alias, proto);
 
     if (Dump) {
       llvm::dbgs() << "}\n";

lib/AST/RequirementMachine/RequirementLowering.h

@@ -15,7 +15,7 @@
 
 #include "swift/AST/Type.h"
 #include "llvm/ADT/ArrayRef.h"
-#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include <vector>
 #include "Diagnostics.h"
@@ -76,24 +76,22 @@ getRuleForRequirement(const Requirement &req,
 struct RuleBuilder {
   RewriteContext &Context;
 
-  /// The keys are the unique protocols we've added so far. The value indicates
-  /// whether the protocol's SCC is an initial component for the rewrite system.
-  ///
-  /// A rewrite system built from a generic signature does not have any initial
-  /// protocols.
-  ///
-  /// A rewrite system built from a protocol SCC has the protocols of the SCC
-  /// itself as initial protocols.
+  /// The transitive closure of all protocols appearing on the right hand
+  /// side of conformance requirements.
+  llvm::DenseSet<const ProtocolDecl *> &ReferencedProtocols;
+
+  /// A subset of the above in insertion order, consisting of the protocols
+  /// whose rules we are going to import.
   ///
-  /// If a protocol is an initial protocol, we use its structural requirements
-  /// instead of its requirement signature as the basis of its rewrite rules.
+  /// If this is a rewrite system built from a generic signature, this vector
+  /// contains all elements in the above set.
   ///
-  /// This is what breaks the cycle in requirement signature computation for a
-  /// group of interdependent protocols.
-  llvm::DenseMap<const ProtocolDecl *, bool> &ProtocolMap;
-
-  /// The keys of the above map in insertion order.
-  std::vector<const ProtocolDecl *> Protocols;
+  /// If this is a rewrite system built from a strongly connected component
+  /// of the protocol, this vector contains all elements in the above set
+  /// except for the protocols belonging to the component representing the
+  /// rewrite system itself; those protocols are added directly instead of
+  /// being imported.
+  std::vector<const ProtocolDecl *> ProtocolsToImport;
 
   /// New rules to add which will be marked 'permanent'. These are rules for
   /// introducing associated types, and relationships between layout,
@@ -116,18 +114,18 @@ struct RuleBuilder {
   bool Dump;
 
   RuleBuilder(RewriteContext &ctx,
-              llvm::DenseMap<const ProtocolDecl *, bool> &protocolMap)
-      : Context(ctx), ProtocolMap(protocolMap),
+              llvm::DenseSet<const ProtocolDecl *> &referencedProtocols)
+      : Context(ctx), ReferencedProtocols(referencedProtocols),
         Dump(ctx.getASTContext().LangOpts.DumpRequirementMachine) {}
 
   void addRequirements(ArrayRef<Requirement> requirements);
   void addRequirements(ArrayRef<StructuralRequirement> requirements);
   void addProtocols(ArrayRef<const ProtocolDecl *> proto);
-  void addProtocol(const ProtocolDecl *proto,
-                   bool initialComponent);
+  void addReferencedProtocol(const ProtocolDecl *proto);
   void collectRulesFromReferencedProtocols();
 
 private:
+  void addPermanentProtocolRules(const ProtocolDecl *proto);
   void addAssociatedType(const AssociatedTypeDecl *type,
                          const ProtocolDecl *proto);
   void addRequirement(const Requirement &req,

lib/AST/RequirementMachine/RequirementMachine.cpp

@@ -90,7 +90,7 @@ RequirementMachine::initWithGenericSignature(CanGenericSignature sig) {
 
   // Collect the top-level requirements, and all transtively-referenced
   // protocol requirement signatures.
-  RuleBuilder builder(Context, System.getProtocolMap());
+  RuleBuilder builder(Context, System.getReferencedProtocols());
   builder.addRequirements(sig.getRequirements());
 
   // Add the initial set of rewrite rules to the rewrite system.
@@ -134,7 +134,7 @@ RequirementMachine::initWithProtocols(ArrayRef<const ProtocolDecl *> protos) {
     llvm::dbgs() << " {\n";
   }
 
-  RuleBuilder builder(Context, System.getProtocolMap());
+  RuleBuilder builder(Context, System.getReferencedProtocols());
   builder.addProtocols(protos);
 
   // Add the initial set of rewrite rules to the rewrite system.
@@ -181,7 +181,7 @@ RequirementMachine::initWithWrittenRequirements(
 
   // Collect the top-level requirements, and all transtively-referenced
   // protocol requirement signatures.
-  RuleBuilder builder(Context, System.getProtocolMap());
+  RuleBuilder builder(Context, System.getReferencedProtocols());
   builder.addRequirements(requirements);
 
   // Add the initial set of rewrite rules to the rewrite system.

lib/AST/RequirementMachine/RewriteSystem.h

@@ -76,16 +76,15 @@ class RewriteSystem final {
   /// type is an index into the Rules array defined above.
   Trie<unsigned, MatchKind::Shortest> Trie;
 
-  /// The set of protocols known to this rewrite system. The boolean associated
-  /// with each key is true if the protocol is part of the 'Protos' set above,
-  /// otherwies it is false.
+  /// The set of protocols known to this rewrite system.
   ///
-  /// See RuleBuilder::ProtocolMap for a more complete explanation. For the most
-  /// part, this is only used while building the rewrite system, but conditional
-  /// requirement inference forces us to be able to add new protocols to the
-  /// rewrite system after the fact, so this little bit of RuleBuilder state
-  /// outlives the initialization phase.
-  llvm::DenseMap<const ProtocolDecl *, bool> ProtocolMap;
+  /// See RuleBuilder::ReferencedProtocols for a more complete explanation.
+  ///
+  /// For the most part, this is only used while building the rewrite system,
+  /// but conditional requirement inference forces us to be able to add new
+  /// protocols to the rewrite system after the fact, so this little bit of
+  /// RuleBuilder state outlives the initialization phase.
+  llvm::DenseSet<const ProtocolDecl *> ReferencedProtocols;
 
   DebugOptions Debug;
 
@@ -117,8 +116,8 @@ class RewriteSystem final {
   /// Return the rewrite context used for allocating memory.
   RewriteContext &getRewriteContext() const { return Context; }
 
-  llvm::DenseMap<const ProtocolDecl *, bool> &getProtocolMap() {
-    return ProtocolMap;
+  llvm::DenseSet<const ProtocolDecl *> &getReferencedProtocols() {
+    return ReferencedProtocols;
   }
 
   DebugOptions getDebugOptions() const { return Debug; }
@@ -133,7 +132,7 @@ class RewriteSystem final {
   }
 
   bool isKnownProtocol(const ProtocolDecl *proto) const {
-    return ProtocolMap.find(proto) != ProtocolMap.end();
+    return ReferencedProtocols.count(proto) > 0;
   }
 
   unsigned getRuleID(const Rule &rule) const {