Strip the REPLChecker of its TopLevelContext

It's really not important where it manages to find this unique value, but it's gotta get it from somewhere.  Shameless steal DiscriminatorFinder from DebuggerTransform to scrobble for the next available discriminator value.

Author: CodaFi

lib/Sema/TypeCheckREPL.cpp

@@ -17,6 +17,7 @@
 #include "TypeChecker.h"
 #include "swift/AST/ASTContext.h"
 #include "swift/AST/ASTVisitor.h"
+#include "swift/AST/ASTWalker.h"
 #include "swift/AST/DiagnosticsFrontend.h"
 #include "swift/AST/Expr.h"
 #include "swift/AST/NameLookup.h"
@@ -30,6 +31,36 @@
 using namespace swift;
 
 namespace {
+
+/// Find available closure discriminators.
+///
+/// The parser typically takes care of assigning unique discriminators to
+/// closures, but the parser is unavailable to this transform.
+class DiscriminatorFinder : public ASTWalker {
+  unsigned NextDiscriminator = 0;
+
+public:
+  Expr *walkToExprPost(Expr *E) override {
+    auto *ACE = dyn_cast<AbstractClosureExpr>(E);
+    if (!ACE)
+      return E;
+
+    unsigned Discriminator = ACE->getDiscriminator();
+    assert(Discriminator != AbstractClosureExpr::InvalidDiscriminator &&
+           "Existing closures should have valid discriminators");
+    if (Discriminator >= NextDiscriminator)
+      NextDiscriminator = Discriminator + 1;
+    return E;
+  }
+
+  // Get the next available closure discriminator.
+  unsigned getNextDiscriminator() {
+    if (NextDiscriminator == AbstractClosureExpr::InvalidDiscriminator)
+      llvm::report_fatal_error("Out of valid closure discriminators");
+    return NextDiscriminator++;
+  }
+};
+
 struct REPLContext {
   ASTContext &Context;
   SourceFile &SF;
@@ -182,14 +213,14 @@ struct PatternBindingPrintLHS : public ASTVisitor<PatternBindingPrintLHS> {
 
 namespace {
   class REPLChecker : public REPLContext {
-    TopLevelContext &TLC;
+    DiscriminatorFinder &DF;
 
     /// The index of the next response metavariable to bind to a REPL result.
     unsigned NextResponseVariableIndex = 0;
 
   public:
-    REPLChecker(SourceFile &SF, TopLevelContext &TLC)
-      : REPLContext(SF), TLC(TLC) {}
+    REPLChecker(SourceFile &SF, DiscriminatorFinder &DF)
+        : REPLContext(SF), DF(DF) {}
 
     void processREPLTopLevelExpr(Expr *E);
     void processREPLTopLevelPatternBinding(PatternBindingDecl *PBD);
@@ -222,7 +253,7 @@ void REPLChecker::generatePrintOfExpression(StringRef NameStr, Expr *E) {
   Arg->setSpecifier(ParamSpecifier::Default);
   auto params = ParameterList::createWithoutLoc(Arg);
 
-  unsigned discriminator = TLC.claimNextClosureDiscriminator();
+  unsigned discriminator = DF.getNextDiscriminator();
 
   ClosureExpr *CE =
       new (Context) ClosureExpr(params, SourceLoc(), SourceLoc(), SourceLoc(),
@@ -429,13 +460,18 @@ Identifier REPLChecker::getNextResponseVariableName(DeclContext *DC) {
 /// processREPLTopLevel - This is called after we've parsed and typechecked some
 /// new decls at the top level.  We inject code to print out expressions and
 /// pattern bindings the are evaluated.
-void TypeChecker::processREPLTopLevel(SourceFile &SF, TopLevelContext &TLC,
-                                      unsigned FirstDecl) {
+void TypeChecker::processREPLTopLevel(SourceFile &SF, unsigned FirstDecl) {
+  // Walk over all decls in the file to find the next available closure
+  // discriminator.
+  DiscriminatorFinder DF;
+  for (Decl *D : SF.Decls)
+    D->walk(DF);
+
   // Move new declarations out.
   std::vector<Decl *> NewDecls(SF.Decls.begin()+FirstDecl, SF.Decls.end());
   SF.Decls.resize(FirstDecl);
 
-  REPLChecker RC(SF, TLC);
+  REPLChecker RC(SF, DF);
 
   // Loop over each of the new decls, processing them, adding them back to
   // the Decls list.
@@ -455,7 +491,7 @@ void TypeChecker::processREPLTopLevel(SourceFile &SF, TopLevelContext &TLC,
       if (auto *PBD = dyn_cast<PatternBindingDecl>(D))
         RC.processREPLTopLevelPatternBinding(PBD);
 
-    contextualizeTopLevelCode(TLC, TLCD);
+    TypeChecker::contextualizeTopLevelCode(TLCD);
   }
 
   SF.clearLookupCache();

lib/Sema/TypeChecker.cpp

@@ -375,7 +375,7 @@ void swift::performTypeChecking(SourceFile &SF, TopLevelContext &TLC,
     // If we're in REPL mode, inject temporary result variables and other stuff
     // that the REPL needs to synthesize.
     if (SF.Kind == SourceFileKind::REPL && !Ctx.hadError())
-      TypeChecker::processREPLTopLevel(SF, TLC, StartElem);
+      TypeChecker::processREPLTopLevel(SF, StartElem);
 
     typeCheckFunctionsAndExternalDecls(SF, TC);
   }

lib/Sema/TypeChecker.h

@@ -38,7 +38,6 @@ namespace swift {
 class GenericSignatureBuilder;
 class NominalTypeDecl;
 class NormalProtocolConformance;
-class TopLevelContext;
 class TypeChecker;
 class TypeResolution;
 class TypeResolutionOptions;
@@ -819,8 +818,7 @@ class TypeChecker final {
 
   static void typeCheckTopLevelCodeDecl(TopLevelCodeDecl *TLCD);
 
-  static void processREPLTopLevel(SourceFile &SF, TopLevelContext &TLC,
-                                  unsigned StartElem);
+  static void processREPLTopLevel(SourceFile &SF, unsigned StartElem);
 
   static void typeCheckDecl(Decl *D);