AST: GenericSignatures point directly to their RequirementMachine

include/swift/AST/GenericSignature.h

@@ -265,6 +265,8 @@ class alignas(1 << TypeAlignInBits) GenericSignatureImpl final
 
   GenericEnvironment *GenericEnv = nullptr;
 
+  rewriting::RequirementMachine *Machine = nullptr;
+
   // Make vanilla new/delete illegal.
   void *operator new(size_t Bytes) = delete;
   void operator delete(void *Data) = delete;

lib/AST/ASTContext.cpp

@@ -67,7 +67,6 @@
 #include <algorithm>
 #include <memory>
 
-#include "RequirementMachine/RequirementMachine.h"
 #include "RequirementMachine/RewriteContext.h"
 
 using namespace swift;
@@ -538,14 +537,6 @@ struct ASTContext::Implementation {
 
   /// Memory allocation arena for the term rewriting system.
   std::unique_ptr<rewriting::RewriteContext> TheRewriteContext;
-
-  /// Stored requirement machines for canonical generic signatures.
-  ///
-  /// This should come after TheRewriteContext above, since various destructors
-  /// compile stats in the histograms stored in our RewriteContext.
-  llvm::DenseMap<GenericSignature,
-                 std::unique_ptr<rewriting::RequirementMachine>>
-    RequirementMachines;
 };
 
 ASTContext::Implementation::Implementation()
@@ -1965,32 +1956,7 @@ ASTContext::getOrCreateRequirementMachine(CanGenericSignature sig) {
   if (!rewriteCtx)
     rewriteCtx.reset(new rewriting::RewriteContext(*this));
 
-  // Check whether we already have a requirement machine for this
-  // signature.
-  auto &machines = getImpl().RequirementMachines;
-
-  auto &machinePtr = machines[sig];
-  if (machinePtr) {
-    auto *machine = machinePtr.get();
-    if (!machine->isComplete()) {
-      llvm::errs() << "Re-entrant construction of requirement "
-                   << "machine for " << sig << "\n";
-      abort();
-    }
-
-    return machine;
-  }
-
-  auto *machine = new rewriting::RequirementMachine(*rewriteCtx);
-
-  // Store this requirement machine before adding the signature,
-  // to catch re-entrant construction via addGenericSignature()
-  // below.
-  machinePtr.reset(machine);
-
-  machine->addGenericSignature(sig);
-
-  return machine;
+  return rewriteCtx->getRequirementMachine(sig);
 }
 
 bool ASTContext::isRecursivelyConstructingRequirementMachine(
@@ -1999,13 +1965,7 @@ bool ASTContext::isRecursivelyConstructingRequirementMachine(
   if (!rewriteCtx)
     return false;
 
-  auto &machines = getImpl().RequirementMachines;
-
-  auto found = machines.find(sig);
-  if (found == machines.end())
-    return false;
-
-  return !found->second->isComplete();
+  return rewriteCtx->isRecursivelyConstructingRequirementMachine(sig);
 }
 
 Optional<llvm::TinyPtrVector<ValueDecl *>>

lib/AST/GenericSignature.cpp

@@ -194,13 +194,13 @@ GenericSignatureImpl::getGenericSignatureBuilder() const {
 
 rewriting::RequirementMachine *
 GenericSignatureImpl::getRequirementMachine() const {
-  // The requirement machine is associated with the canonical signature.
-  if (!isCanonical())
-    return getCanonicalSignature()->getRequirementMachine();
+  if (Machine)
+    return Machine;
 
-  // Requirement machines are stored on the ASTContext.
-  return getASTContext().getOrCreateRequirementMachine(
-                                             CanGenericSignature(this));
+  const_cast<GenericSignatureImpl *>(this)->Machine
+      = getASTContext().getOrCreateRequirementMachine(
+          getCanonicalSignature());
+  return Machine;
 }
 
 bool GenericSignatureImpl::isEqual(GenericSignature Other) const {

lib/AST/RequirementMachine/RequirementMachine.cpp

@@ -342,7 +342,7 @@ RequirementMachine::RequirementMachine(RewriteContext &ctx)
 
 RequirementMachine::~RequirementMachine() {}
 
-void RequirementMachine::addGenericSignature(CanGenericSignature sig) {
+void RequirementMachine::initWithGenericSignature(CanGenericSignature sig) {
   Sig = sig;
 
   PrettyStackTraceGenericSignature debugStack("building rewrite system for", sig);
@@ -359,7 +359,6 @@ void RequirementMachine::addGenericSignature(CanGenericSignature sig) {
     llvm::dbgs() << "Adding generic signature " << sig << " {\n";
   }
 
-
   // Collect the top-level requirements, and all transtively-referenced
   // protocol requirement signatures.
   RewriteSystemBuilder builder(Context, Dump);

lib/AST/RequirementMachine/RequirementMachine.h

@@ -45,6 +45,7 @@ class RewriteContext;
 /// generic signatures and interface types.
 class RequirementMachine final {
   friend class swift::ASTContext;
+  friend class swift::rewriting::RewriteContext;
 
   CanGenericSignature Sig;
 
@@ -77,7 +78,7 @@ class RequirementMachine final {
   RequirementMachine &operator=(const RequirementMachine &) = delete;
   RequirementMachine &operator=(RequirementMachine &&) = delete;
 
-  void addGenericSignature(CanGenericSignature sig);
+  void initWithGenericSignature(CanGenericSignature sig);
 
   bool isComplete() const;
   void computeCompletion();
@@ -117,4 +118,4 @@ class RequirementMachine final {
 
 } // end namespace swift
 
-#endif
+#endif

lib/AST/RequirementMachine/RewriteContext.cpp

@@ -13,6 +13,7 @@
 #include "swift/AST/Decl.h"
 #include "swift/AST/Types.h"
 #include "ProtocolGraph.h"
+#include "RequirementMachine.h"
 #include "RewriteSystem.h"
 #include "RewriteContext.h"
 
@@ -346,6 +347,37 @@ Type RewriteContext::getRelativeTypeForTerm(
       { }, protos, *this);
 }
 
+RequirementMachine *RewriteContext::getRequirementMachine(
+    CanGenericSignature sig) {
+  auto &machine = Machines[sig];
+  if (machine) {
+    if (!machine->isComplete()) {
+      llvm::errs() << "Re-entrant construction of requirement "
+                   << "machine for " << sig << "\n";
+      abort();
+    }
+
+    return machine;
+  }
+
+  // Store this requirement machine before adding the signature,
+  // to catch re-entrant construction via initWithGenericSignature()
+  // below.
+  machine = new rewriting::RequirementMachine(*this);
+  machine->initWithGenericSignature(sig);
+
+  return machine;
+}
+
+bool RewriteContext::isRecursivelyConstructingRequirementMachine(
+    CanGenericSignature sig) {
+  auto found = Machines.find(sig);
+  if (found == Machines.end())
+    return false;
+
+  return !found->second->isComplete();
+}
+
 /// We print stats in the destructor, which should get executed at the end of
 /// a compilation job.
 RewriteContext::~RewriteContext() {

lib/AST/RequirementMachine/RewriteContext.h

@@ -52,6 +52,9 @@ class RewriteContext final {
   /// Cache for merged associated type symbols.
   llvm::DenseMap<std::pair<Symbol, Symbol>, Symbol> MergedAssocTypes;
 
+  /// Requirement machines built from generic signatures.
+  llvm::DenseMap<GenericSignature, RequirementMachine *> Machines;
+
   ASTContext &Context;
 
   DebugOptions Debug;
@@ -102,6 +105,9 @@ class RewriteContext final {
   Symbol mergeAssociatedTypes(Symbol lhs, Symbol rhs,
                               const ProtocolGraph &protos);
 
+  RequirementMachine *getRequirementMachine(CanGenericSignature sig);
+  bool isRecursivelyConstructingRequirementMachine(CanGenericSignature sig);
+
   ~RewriteContext();
 };