RequirementMachine: Re-use a single global RewriteContext

Author: slavapestov

include/swift/AST/RequirementMachine.h

@@ -21,6 +21,10 @@ class raw_ostream;
 
 namespace swift {
 
+namespace rewriting {
+class RewriteContext;
+}
+
 class ASTContext;
 class AssociatedTypeDecl;
 class CanType;
@@ -40,7 +44,7 @@ class RequirementMachine final {
   ASTContext &Context;
   Implementation *Impl;
 
-  explicit RequirementMachine(ASTContext &ctx);
+  explicit RequirementMachine(rewriting::RewriteContext &rewriteCtx);
 
   RequirementMachine(const RequirementMachine &) = delete;
   RequirementMachine(RequirementMachine &&) = delete;

lib/AST/ASTContext.cpp

@@ -67,6 +67,8 @@
 #include <algorithm>
 #include <memory>
 
+#include "RequirementMachine/RewriteContext.h"
+
 using namespace swift;
 
 #define DEBUG_TYPE "ASTContext"
@@ -531,6 +533,9 @@ struct ASTContext::Implementation {
   /// The scratch context used to allocate intrinsic data on behalf of \c swift::IntrinsicInfo
   std::unique_ptr<llvm::LLVMContext> IntrinsicScratchContext;
 #endif
+
+  /// Memory allocation arena for the term rewriting system.
+  std::unique_ptr<rewriting::RewriteContext> TheRewriteContext;
 };
 
 ASTContext::Implementation::Implementation()
@@ -1771,15 +1776,20 @@ static AllocationArena getArena(GenericSignature genericSig) {
   if (!genericSig)
     return AllocationArena::Permanent;
 
-  if (genericSig->hasTypeVariable())
+  if (genericSig->hasTypeVariable()) {
+    assert(false && "What's going on");
     return AllocationArena::ConstraintSolver;
+  }
 
   return AllocationArena::Permanent;
 }
 
 void ASTContext::registerGenericSignatureBuilder(
                                        GenericSignature sig,
                                        GenericSignatureBuilder &&builder) {
+  if (LangOpts.EnableRequirementMachine == RequirementMachineMode::Enabled)
+    return;
+
   auto canSig = sig.getCanonicalSignature();
   auto arena = getArena(sig);
   auto &genericSignatureBuilders =
@@ -1885,6 +1895,12 @@ GenericSignatureBuilder *ASTContext::getOrCreateGenericSignatureBuilder(
 
 RequirementMachine *ASTContext::getOrCreateRequirementMachine(
     CanGenericSignature sig) {
+  assert(!sig->hasTypeVariable());
+
+  auto &rewriteCtx = getImpl().TheRewriteContext;
+  if (!rewriteCtx)
+    rewriteCtx.reset(new rewriting::RewriteContext(*this));
+
   // Check whether we already have a requirement machine for this
   // signature.
   auto arena = getArena(sig);
@@ -1902,12 +1918,12 @@ RequirementMachine *ASTContext::getOrCreateRequirementMachine(
     return machine;
   }
 
-  auto *machine = new RequirementMachine(*this);
+  auto *machine = new RequirementMachine(*rewriteCtx);
 
   // Store this requirement machine before adding the signature,
   // to catch re-entrant construction via addGenericSignature()
   // below.
-  machinePtr = std::unique_ptr<RequirementMachine>(machine);
+  machinePtr.reset(machine);
 
   machine->addGenericSignature(sig);
 

lib/AST/RequirementMachine/RequirementMachine.cpp

@@ -295,7 +295,8 @@ void RequirementMachine::Implementation::dump(llvm::raw_ostream &out) {
   Map.dump(out);
 }
 
-RequirementMachine::RequirementMachine(ASTContext &ctx) : Context(ctx) {
+RequirementMachine::RequirementMachine(RewriteContext &ctx)
+    : Context(ctx.getASTContext()) {
   Impl = new Implementation(ctx);
 }
 

lib/AST/RequirementMachine/RequirementMachineImpl.h

@@ -28,7 +28,7 @@ namespace swift {
 
 /// We use the PIMPL pattern to avoid creeping header dependencies.
 struct RequirementMachine::Implementation {
-  rewriting::RewriteContext Context;
+  rewriting::RewriteContext &Context;
   rewriting::RewriteSystem System;
   rewriting::EquivalenceClassMap Map;
   CanGenericSignature Sig;
@@ -45,7 +45,7 @@ struct RequirementMachine::Implementation {
   std::vector<std::pair<CanType, ConformanceAccessPath>>
       CurrentConformanceAccessPaths;
 
-  explicit Implementation(ASTContext &ctx)
+  explicit Implementation(rewriting::RewriteContext &ctx)
       : Context(ctx),
         System(Context),
         Map(Context, System.getProtocols()) {}