RequirementMachine: Track known protocols in the rewrite system

Conditional requirement inference needs to be able to add rewrite rules
from the requirement signatures of hitherto-unseen protocols, so to
help with that, extract out the RuleBuilder's ProtocolMap and move it
into the RewriteSystem.

Author: slavapestov

lib/AST/RequirementMachine/RequirementLowering.h

@@ -65,7 +65,6 @@ getRuleForRequirement(const Requirement &req,
 /// appearing on the right hand side of conformance requirements.
 struct RuleBuilder {
   RewriteContext &Context;
-  bool Dump;
 
   /// 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.
@@ -81,7 +80,9 @@ struct RuleBuilder {
   ///
   /// This is what breaks the cycle in requirement signature computation for a
   /// group of interdependent protocols.
-  llvm::DenseMap<const ProtocolDecl *, bool> ProtocolMap;
+  llvm::DenseMap<const ProtocolDecl *, bool> &ProtocolMap;
+
+  /// The keys of the above map in insertion order.
   std::vector<const ProtocolDecl *> Protocols;
 
   /// New rules to add which will be marked 'permanent'. These are rules for
@@ -95,7 +96,15 @@ struct RuleBuilder {
   /// eliminated by homotopy reduction.
   std::vector<std::pair<MutableTerm, MutableTerm>> RequirementRules;
 
-  RuleBuilder(RewriteContext &ctx, bool dump) : Context(ctx), Dump(dump) {}
+  /// Enables debugging output. Controlled by the -dump-requirement-machine
+  /// frontend flag.
+  bool Dump;
+
+  RuleBuilder(RewriteContext &ctx,
+              llvm::DenseMap<const ProtocolDecl *, bool> &protocolMap)
+      : Context(ctx), ProtocolMap(protocolMap),
+        Dump(ctx.getASTContext().LangOpts.DumpRequirementMachine) {}
+
   void addRequirements(ArrayRef<Requirement> requirements);
   void addRequirements(ArrayRef<StructuralRequirement> requirements);
   void addProtocols(ArrayRef<const ProtocolDecl *> proto);

lib/AST/RequirementMachine/RequirementMachine.cpp

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

lib/AST/RequirementMachine/RewriteSystem.h

@@ -203,6 +203,17 @@ 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.
+  ///
+  /// 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;
+
   DebugOptions Debug;
 
   /// Whether we've initialized the rewrite system with a call to initialize().
@@ -233,6 +244,10 @@ class RewriteSystem final {
   /// Return the rewrite context used for allocating memory.
   RewriteContext &getRewriteContext() const { return Context; }
 
+  llvm::DenseMap<const ProtocolDecl *, bool> &getProtocolMap() {
+    return ProtocolMap;
+  }
+
   DebugOptions getDebugOptions() const { return Debug; }
 
   void initialize(bool recordLoops, ArrayRef<const ProtocolDecl *> protos,
@@ -243,6 +258,10 @@ class RewriteSystem final {
     return Protos;
   }
 
+  bool isKnownProtocol(const ProtocolDecl *proto) const {
+    return ProtocolMap.find(proto) != ProtocolMap.end();
+  }
+
   unsigned getRuleID(const Rule &rule) const {
     assert((unsigned)(&rule - &*Rules.begin()) < Rules.size());
     return (unsigned)(&rule - &*Rules.begin());