RequirementMachine: Compute strongly connected components from the protocol dependency graph

Author: slavapestov

include/swift/AST/ASTContext.h

@@ -1198,6 +1198,11 @@ class ASTContext final {
   bool isRecursivelyConstructingRequirementMachine(
       CanGenericSignature sig);
 
+  /// Retrieve or create a term rewriting system for answering queries on
+  /// type parameters written against the given protocol requirement signature.
+  rewriting::RequirementMachine *getOrCreateRequirementMachine(
+      const ProtocolDecl *proto);
+
   /// Retrieve a generic signature with a single unconstrained type parameter,
   /// like `<T>`.
   CanGenericSignature getSingleGenericParameterSignature() const;

lib/AST/ASTContext.cpp

@@ -1985,6 +1985,15 @@ bool ASTContext::isRecursivelyConstructingRequirementMachine(
   return rewriteCtx->isRecursivelyConstructingRequirementMachine(sig);
 }
 
+rewriting::RequirementMachine *
+ASTContext::getOrCreateRequirementMachine(const ProtocolDecl *proto) {
+  auto &rewriteCtx = getImpl().TheRewriteContext;
+  if (!rewriteCtx)
+    rewriteCtx.reset(new rewriting::RewriteContext(*this));
+
+  return rewriteCtx->getRequirementMachine(proto);
+}
+
 Optional<llvm::TinyPtrVector<ValueDecl *>>
 OverriddenDeclsRequest::getCachedResult() const {
   auto decl = std::get<0>(getStorage());

lib/AST/RequirementMachine/Debug.h

@@ -43,6 +43,9 @@ enum class DebugFlags : unsigned {
 
   /// Print debug output from the generating conformances algorithm.
   GeneratingConformances = (1<<7),
+
+  /// Print debug output from the protocol dependency graph.
+  ProtocolDependencies = (1<<8)
 };
 
 using DebugOptions = OptionSet<DebugFlags>;

lib/AST/RequirementMachine/ProtocolGraph.cpp

@@ -28,6 +28,14 @@ void ProtocolGraph::visitRequirements(ArrayRef<Requirement> reqs) {
   }
 }
 
+/// Adds information about all protocols transitvely referenced from
+/// \p protos.
+void ProtocolGraph::visitProtocols(ArrayRef<const ProtocolDecl *> protos) {
+  for (auto proto : protos) {
+    addProtocol(proto);
+  }
+}
+
 /// Return true if we know about this protocol.
 bool ProtocolGraph::isKnownProtocol(const ProtocolDecl *proto) const {
   return Info.count(proto) > 0;

lib/AST/RequirementMachine/ProtocolGraph.h

@@ -85,6 +85,7 @@ class ProtocolGraph {
   bool Debug = false;
 
 public:
+  void visitProtocols(ArrayRef<const ProtocolDecl *> protos);
   void visitRequirements(ArrayRef<Requirement> reqs);
 
   bool isKnownProtocol(const ProtocolDecl *proto) const;

lib/AST/RequirementMachine/RequirementMachine.cpp

@@ -43,10 +43,12 @@ struct RewriteSystemBuilder {
   RewriteSystemBuilder(RewriteContext &ctx, bool dump)
     : Context(ctx), Dump(dump) {}
   void addGenericSignature(CanGenericSignature sig);
+  void addProtocols(ArrayRef<const ProtocolDecl *> proto);
   void addAssociatedType(const AssociatedTypeDecl *type,
                          const ProtocolDecl *proto);
   void addRequirement(const Requirement &req,
                       const ProtocolDecl *proto);
+  void processProtocolDependencies();
 };
 
 } // end namespace
@@ -85,33 +87,22 @@ void RewriteSystemBuilder::addGenericSignature(CanGenericSignature sig) {
   Protocols.visitRequirements(sig.getRequirements());
   Protocols.compute();
 
-  // Add rewrite rules for each protocol.
-  for (auto *proto : Protocols.getProtocols()) {
-    if (Dump) {
-      llvm::dbgs() << "protocol " << proto->getName() << " {\n";
-    }
-
-    const auto &info = Protocols.getProtocolInfo(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);
-
-    if (Dump) {
-      llvm::dbgs() << "}\n";
-    }
-  }
+  processProtocolDependencies();
 
   // Add rewrite rules for all requirements in the top-level signature.
   for (const auto &req : sig.getRequirements())
     addRequirement(req, /*proto=*/nullptr);
 }
 
+void RewriteSystemBuilder::addProtocols(ArrayRef<const ProtocolDecl *> protos) {
+  // Collect all protocols transitively referenced from this connected component
+  // of the protocol dependency graph.
+  Protocols.visitProtocols(protos);
+  Protocols.compute();
+
+  processProtocolDependencies();
+}
+
 /// For an associated type T in a protocol P, we add a rewrite rule:
 ///
 ///   [P].T => [P:T]
@@ -237,6 +228,30 @@ void RewriteSystemBuilder::addRequirement(const Requirement &req,
   RequirementRules.emplace_back(subjectTerm, constraintTerm);
 }
 
+void RewriteSystemBuilder::processProtocolDependencies() {
+  // Add rewrite rules for each protocol.
+  for (auto *proto : Protocols.getProtocols()) {
+    if (Dump) {
+      llvm::dbgs() << "protocol " << proto->getName() << " {\n";
+    }
+
+    const auto &info = Protocols.getProtocolInfo(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);
+
+    if (Dump) {
+      llvm::dbgs() << "}\n";
+    }
+  }
+}
+
 void RequirementMachine::verify(const MutableTerm &term) const {
 #ifndef NDEBUG
   // If the term is in the generic parameter domain, ensure we have a valid
@@ -378,6 +393,40 @@ void RequirementMachine::initWithGenericSignature(CanGenericSignature sig) {
   }
 }
 
+void RequirementMachine::initWithProtocols(ArrayRef<const ProtocolDecl *> protos) {
+  auto &ctx = Context.getASTContext();
+  auto *Stats = ctx.Stats;
+
+  if (Stats)
+    ++Stats->getFrontendCounters().NumRequirementMachines;
+
+  FrontendStatsTracer tracer(Stats, "build-rewrite-system");
+
+  if (Dump) {
+    llvm::dbgs() << "Adding protocols";
+    for (auto *proto : protos) {
+      llvm::dbgs() << " " << proto->getName();
+    }
+    llvm::dbgs() << " {\n";
+  }
+
+  RewriteSystemBuilder builder(Context, Dump);
+  builder.addProtocols(protos);
+
+  // Add the initial set of rewrite rules to the rewrite system, also
+  // providing the protocol graph to use for the linear order on terms.
+  System.initialize(std::move(builder.AssociatedTypeRules),
+                    std::move(builder.RequirementRules),
+                    std::move(builder.Protocols));
+
+  // FIXME: Only if the protocols were written in source, though.
+  computeCompletion(RewriteSystem::AllowInvalidRequirements);
+
+  if (Dump) {
+    llvm::dbgs() << "}\n";
+  }
+}
+
 /// Attempt to obtain a confluent rewrite system using the completion
 /// procedure.
 void RequirementMachine::computeCompletion(RewriteSystem::ValidityPolicy policy) {

lib/AST/RequirementMachine/RequirementMachine.h

@@ -79,6 +79,7 @@ class RequirementMachine final {
   RequirementMachine &operator=(RequirementMachine &&) = delete;
 
   void initWithGenericSignature(CanGenericSignature sig);
+  void initWithProtocols(ArrayRef<const ProtocolDecl *> protos);
 
   bool isComplete() const;
 

lib/AST/RequirementMachine/RewriteContext.cpp

@@ -16,6 +16,7 @@
 #include "RequirementMachine.h"
 #include "RewriteSystem.h"
 #include "RewriteContext.h"
+#include "RequirementMachine.h"
 
 using namespace swift;
 using namespace rewriting;
@@ -36,6 +37,7 @@ static DebugOptions parseDebugFlags(StringRef debugFlags) {
       .Case("concretize-nested-types", DebugFlags::ConcretizeNestedTypes)
       .Case("homotopy-reduction", DebugFlags::HomotopyReduction)
       .Case("generating-conformances", DebugFlags::GeneratingConformances)
+      .Case("protocol-dependencies", DebugFlags::ProtocolDependencies)
       .Default(None);
     if (!flag) {
       llvm::errs() << "Unknown debug flag in -debug-requirement-machine "
@@ -379,6 +381,110 @@ bool RewriteContext::isRecursivelyConstructingRequirementMachine(
   return !found->second->isComplete();
 }
 
+void RewriteContext::getRequirementMachineRec(
+    const ProtocolDecl *proto,
+    SmallVectorImpl<const ProtocolDecl *> &stack) {
+  assert(Protos.count(proto) == 0);
+
+  // Initialize the next component index and push the entry
+  // on the stack
+  {
+    auto &entry = Protos[proto];
+    entry.Index = NextComponentIndex;
+    entry.LowLink = NextComponentIndex;
+    entry.OnStack = 1;
+  }
+
+  NextComponentIndex++;
+  stack.push_back(proto);
+
+  // Look at each successor.
+  for (auto *depProto : proto->getProtocolDependencies()) {
+    auto found = Protos.find(depProto);
+    if (found == Protos.end()) {
+      // Successor has not yet been visited. Recurse.
+      getRequirementMachineRec(depProto, stack);
+
+      auto &entry = Protos[proto];
+      assert(Protos.count(depProto) != 0);
+      entry.LowLink = std::min(entry.LowLink, Protos[depProto].LowLink);
+    } else if (found->second.OnStack) {
+      // Successor is on the stack and hence in the current SCC.
+      auto &entry = Protos[proto];
+      entry.LowLink = std::min(entry.LowLink, found->second.Index);
+    }
+  }
+
+  auto &entry = Protos[proto];
+
+  // If this a root node, pop the stack and generate an SCC.
+  if (entry.LowLink == entry.Index) {
+    unsigned id = Components.size();
+    SmallVector<const ProtocolDecl *, 3> protos;
+
+    const ProtocolDecl *depProto = nullptr;
+    do {
+      depProto = stack.back();
+      stack.pop_back();
+
+      assert(Protos.count(depProto) != 0);
+      Protos[depProto].OnStack = false;
+      Protos[depProto].ComponentID = id;
+
+      protos.push_back(depProto);
+    } while (depProto != proto);
+
+    if (Debug.contains(DebugFlags::ProtocolDependencies)) {
+      llvm::dbgs() << "Connected component: [";
+      bool first = true;
+      for (auto *depProto : protos) {
+        if (!first) {
+          llvm::dbgs() << ", ";
+        } else {
+          first = false;
+        }
+        llvm::dbgs() << depProto->getName();
+      }
+      llvm::dbgs() << "]\n";
+    }
+
+    Components[id] = {Context.AllocateCopy(protos), nullptr};
+  }
+}
+
+RequirementMachine *RewriteContext::getRequirementMachine(
+    const ProtocolDecl *proto) {
+  auto found = Protos.find(proto);
+  if (found == Protos.end()) {
+    SmallVector<const ProtocolDecl *, 3> stack;
+    getRequirementMachineRec(proto, stack);
+    assert(stack.empty());
+
+    found = Protos.find(proto);
+    assert(found != Protos.end());
+  }
+
+  assert(Components.count(found->second.ComponentID) != 0);
+  auto &component = Components[found->second.ComponentID];
+
+  auto *&machine = component.Machine;
+
+  if (machine) {
+    if (!machine->isComplete()) {
+      llvm::errs() << "Re-entrant construction of requirement "
+                   << "machine for:";
+      for (auto *proto : component.Protos)
+        llvm::errs() << " " << proto->getName();
+      abort();
+    }
+  } else {
+    machine = new RequirementMachine(*this);
+    machine->initWithProtocols(component.Protos);
+  }
+
+  return machine;
+}
+
 /// We print stats in the destructor, which should get executed at the end of
 /// a compilation job.
 RewriteContext::~RewriteContext() {
@@ -402,4 +508,7 @@ RewriteContext::~RewriteContext() {
     delete pair.second;
 
   Machines.clear();
+
+  for (const auto &pair : Components)
+    delete pair.second.Machine;
 }