RequirementMachine: Move some code out of RuleBuilder and into RequirementSignatureRequest

Author: slavapestov

lib/AST/RequirementMachine/RequirementMachine.cpp

@@ -167,22 +167,24 @@ RequirementMachine::initWithGenericSignature(CanGenericSignature sig) {
 /// Returns failure if completion fails within the configured number of steps.
 std::pair<CompletionResult, unsigned>
 RequirementMachine::initWithProtocolWrittenRequirements(
-    ArrayRef<const ProtocolDecl *> protos) {
+    ArrayRef<const ProtocolDecl *> component,
+    const llvm::DenseMap<const ProtocolDecl *,
+                         SmallVector<StructuralRequirement, 4>> protos) {
   FrontendStatsTracer tracer(Stats, "build-rewrite-system");
 
   if (Dump) {
     llvm::dbgs() << "Adding protocols";
-    for (auto *proto : protos) {
+    for (auto *proto : component) {
       llvm::dbgs() << " " << proto->getName();
     }
     llvm::dbgs() << " {\n";
   }
 
   RuleBuilder builder(Context, System.getReferencedProtocols());
-  builder.initWithProtocolWrittenRequirements(protos);
+  builder.initWithProtocolWrittenRequirements(component, protos);
 
   // Add the initial set of rewrite rules to the rewrite system.
-  System.initialize(/*recordLoops=*/true, protos,
+  System.initialize(/*recordLoops=*/true, component,
                     std::move(builder.WrittenRequirements),
                     std::move(builder.ImportedRules),
                     std::move(builder.PermanentRules),

lib/AST/RequirementMachine/RequirementMachine.h

@@ -99,7 +99,9 @@ class RequirementMachine final {
 
   std::pair<CompletionResult, unsigned>
   initWithProtocolWrittenRequirements(
-      ArrayRef<const ProtocolDecl *> protos);
+      ArrayRef<const ProtocolDecl *> component,
+      const llvm::DenseMap<const ProtocolDecl *,
+                           SmallVector<StructuralRequirement, 4>> protos);
 
   std::pair<CompletionResult, unsigned>
   initWithWrittenRequirements(

lib/AST/RequirementMachine/RequirementMachineRequests.cpp

@@ -185,13 +185,23 @@ RequirementSignatureRequestRQM::evaluate(Evaluator &evaluator,
   // component at the same time.
   auto component = ctx.getRewriteContext().getProtocolComponent(proto);
 
+  // Collect user-written requirements from the protocols in this connected
+  // component.
+  llvm::DenseMap<const ProtocolDecl *,
+                 SmallVector<StructuralRequirement, 4>> protos;
+  for (const auto *proto : component) {
+    auto &requirements = protos[proto];
+    for (auto req : proto->getStructuralRequirements())
+      requirements.push_back(req);
+    for (auto req : proto->getTypeAliasRequirements())
+      requirements.push_back({req, SourceLoc(), /*inferred=*/false});
+  }
+
   // Heap-allocate the requirement machine to save stack space.
   std::unique_ptr<RequirementMachine> machine(new RequirementMachine(
       ctx.getRewriteContext()));
 
-  SmallVector<RequirementError, 4> errors;
-
-  auto status = machine->initWithProtocolWrittenRequirements(component);
+  auto status = machine->initWithProtocolWrittenRequirements(component, protos);
   if (status.first != CompletionResult::Success) {
     // All we can do at this point is diagnose and give each protocol an empty
     // requirement signature.
@@ -260,6 +270,7 @@ RequirementSignatureRequestRQM::evaluate(Evaluator &evaluator,
 
   if (ctx.LangOpts.RequirementMachineProtocolSignatures ==
       RequirementMachineMode::Enabled) {
+    SmallVector<RequirementError, 4> errors;
     machine->System.computeRedundantRequirementDiagnostics(errors);
     diagnoseRequirementErrors(ctx, errors,
                               /*allowConcreteGenericParams=*/false);

lib/AST/RequirementMachine/RuleBuilder.cpp

@@ -123,28 +123,31 @@ void RuleBuilder::initWithProtocolSignatureRequirements(
 /// user-written requirements. Used when actually building requirement
 /// signatures.
 void RuleBuilder::initWithProtocolWrittenRequirements(
-    ArrayRef<const ProtocolDecl *> protos) {
+    ArrayRef<const ProtocolDecl *> component,
+    const llvm::DenseMap<const ProtocolDecl *,
+                         SmallVector<StructuralRequirement, 4>> protos) {
   assert(!Initialized);
   Initialized = 1;
 
   // Add all protocols to the referenced set, so that subsequent calls
   // to addReferencedProtocol() with one of these protocols don't add
   // them to the import list.
-  for (auto *proto : protos) {
+  for (const auto *proto : component)
     ReferencedProtocols.insert(proto);
-  }
 
-  for (auto *proto : protos) {
+  for (const auto *proto : component) {
+    auto found = protos.find(proto);
+    assert(found != protos.end());
+    const auto &reqs = found->second;
+
     if (Dump) {
       llvm::dbgs() << "protocol " << proto->getName() << " {\n";
     }
 
     addPermanentProtocolRules(proto);
 
-    for (auto req : proto->getStructuralRequirements())
+    for (auto req : reqs)
       addRequirement(req, proto);
-    for (auto req : proto->getTypeAliasRequirements())
-      addRequirement(req.getCanonical(), proto);
 
     for (auto *otherProto : proto->getProtocolDependencies())
       addReferencedProtocol(otherProto);

lib/AST/RequirementMachine/RuleBuilder.h

@@ -98,7 +98,10 @@ struct RuleBuilder {
   void initWithGenericSignatureRequirements(ArrayRef<Requirement> requirements);
   void initWithWrittenRequirements(ArrayRef<StructuralRequirement> requirements);
   void initWithProtocolSignatureRequirements(ArrayRef<const ProtocolDecl *> proto);
-  void initWithProtocolWrittenRequirements(ArrayRef<const ProtocolDecl *> proto);
+  void initWithProtocolWrittenRequirements(
+      ArrayRef<const ProtocolDecl *> component,
+      const llvm::DenseMap<const ProtocolDecl *,
+                           SmallVector<StructuralRequirement, 4>> protos);
   void initWithConditionalRequirements(ArrayRef<Requirement> requirements,
                                        ArrayRef<Term> substitutions);
   void addReferencedProtocol(const ProtocolDecl *proto);