Merge pull request #28246 from CodaFi/contextual-equivalence

Drop the TopLevelContext from the high-level type checker entrypoints

Author: CodaFi

include/swift/AST/ASTContext.h

@@ -277,6 +277,10 @@ class ASTContext final {
   /// The # of times we have performed typo correction.
   unsigned NumTypoCorrections = 0;
 
+  /// The next auto-closure discriminator.  This needs to be preserved
+  /// across invocations of both the parser and the type-checker.
+  unsigned NextAutoClosureDiscriminator = 0;
+
 private:
   /// The current generation number, which reflects the number of
   /// times that external modules have been loaded.

include/swift/Parse/LocalContext.h

@@ -61,12 +61,7 @@ class LocalContext {
 };
 
 /// Information associated with parsing the top-level context.
-class TopLevelContext : public LocalContext {
-public:
-  /// The next auto-closure discriminator.  This needs to be preserved
-  /// across invocations of both the parser and the type-checker.
-  unsigned NextAutoClosureDiscriminator = 0;
-};
+class TopLevelContext : public LocalContext {};
 
 } // end namespace swift
 

include/swift/Subsystems.h

@@ -168,7 +168,7 @@ namespace swift {
   /// Once parsing and name-binding are complete this optionally walks the ASTs
   /// to add calls to externally provided functions that simulate
   /// "program counter"-like debugging events.
-  void performPCMacro(SourceFile &SF, TopLevelContext &TLC);
+  void performPCMacro(SourceFile &SF);
 
   /// Creates a type checker instance on the given AST context, if it
   /// doesn't already have one.
@@ -184,8 +184,7 @@ namespace swift {
   ///
   /// \param StartElem Where to start for incremental type-checking in the main
   /// source file.
-  void performTypeChecking(SourceFile &SF, TopLevelContext &TLC,
-                           unsigned StartElem = 0);
+  void performTypeChecking(SourceFile &SF, unsigned StartElem = 0);
 
   /// Now that we have type-checked an entire module, perform any type
   /// checking that requires the full module, e.g., Objective-C method

lib/Frontend/Frontend.cpp

@@ -858,10 +858,10 @@ void CompilerInstance::parseAndCheckTypesUpTo(
       return;
     }
 
-    performTypeChecking(SF, PersistentState->getTopLevelContext());
+    performTypeChecking(SF);
 
     if (!Context->hadError() && Invocation.getFrontendOptions().PCMacro) {
-      performPCMacro(SF, PersistentState->getTopLevelContext());
+      performPCMacro(SF);
     }
 
     // Playground transform knows to look out for PCMacro's changes and not
@@ -975,8 +975,7 @@ void CompilerInstance::parseAndTypeCheckMainFileUpTo(
         performNameBinding(MainFile, CurTUElem);
         break;
       case SourceFile::TypeChecked:
-        performTypeChecking(MainFile, PersistentState->getTopLevelContext(),
-                            CurTUElem);
+        performTypeChecking(MainFile, CurTUElem);
         break;
       }
     }

lib/IDE/REPLCodeCompletion.cpp

@@ -212,8 +212,7 @@ doCodeCompletion(SourceFile &SF, StringRef EnteredCode, unsigned *BufferID,
   do {
     parseIntoSourceFile(SF, *BufferID, &Done, nullptr, &PersistentState);
   } while (!Done);
-  performTypeChecking(SF, PersistentState.getTopLevelContext(),
-                      OriginalDeclCount);
+  performTypeChecking(SF, OriginalDeclCount);
 
   performCodeCompletionSecondPass(PersistentState, *CompletionCallbacksFactory);
 

lib/Immediate/REPL.cpp

@@ -193,7 +193,7 @@ typeCheckREPLInput(ModuleDecl *MostRecentModule, StringRef Name,
         parseIntoSourceFile(REPLInputFile, BufferID, &Done, nullptr,
                             &PersistentState);
   } while (!Done);
-  performTypeChecking(REPLInputFile, PersistentState.getTopLevelContext());
+  performTypeChecking(REPLInputFile);
   return REPLModule;
 }
 

lib/Sema/PCMacro.cpp

@@ -665,14 +665,13 @@ class Instrumenter : InstrumenterBase {
 
 } // end anonymous namespace
 
-void swift::performPCMacro(SourceFile &SF, TopLevelContext &TLC) {
+void swift::performPCMacro(SourceFile &SF) {
   class ExpressionFinder : public ASTWalker {
   private:
     unsigned TmpNameIndex = 0;
-    TopLevelContext &TLC;
 
   public:
-    ExpressionFinder(TopLevelContext &TLC) : TLC(TLC) {}
+    ExpressionFinder() = default;
 
     bool walkToDeclPre(Decl *D) override {
       ASTContext &ctx = D->getASTContext();
@@ -692,7 +691,7 @@ void swift::performPCMacro(SourceFile &SF, TopLevelContext &TLC) {
             if (NewBody != Body) {
               TLCD->setBody(NewBody);
               TypeChecker::checkTopLevelErrorHandling(TLCD);
-              TypeChecker::contextualizeTopLevelCode(TLC, TLCD);
+              TypeChecker::contextualizeTopLevelCode(TLCD);
             }
             return false;
           }
@@ -702,7 +701,7 @@ void swift::performPCMacro(SourceFile &SF, TopLevelContext &TLC) {
     }
   };
 
-  ExpressionFinder EF(TLC);
+  ExpressionFinder EF;
   for (Decl *D : SF.Decls) {
     D->walk(EF);
   }

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/TypeCheckStmt.cpp

@@ -231,14 +231,14 @@ bool TypeChecker::contextualizeInitializer(Initializer *DC, Expr *E) {
   return CC.hasAutoClosures();
 }
 
-void TypeChecker::contextualizeTopLevelCode(TopLevelContext &TLC,
-                                            TopLevelCodeDecl *TLCD) {
-  unsigned nextDiscriminator = TLC.NextAutoClosureDiscriminator;
+void TypeChecker::contextualizeTopLevelCode(TopLevelCodeDecl *TLCD) {
+  auto &Context = TLCD->DeclContext::getASTContext();
+  unsigned nextDiscriminator = Context.NextAutoClosureDiscriminator;
   ContextualizeClosures CC(TLCD, nextDiscriminator);
   TLCD->getBody()->walk(CC);
-  assert(nextDiscriminator == TLC.NextAutoClosureDiscriminator &&
+  assert(nextDiscriminator == Context.NextAutoClosureDiscriminator &&
          "reentrant/concurrent invocation of contextualizeTopLevelCode?");
-  TLC.NextAutoClosureDiscriminator = CC.NextDiscriminator;
+  Context.NextAutoClosureDiscriminator = CC.NextDiscriminator;
 }
 
 /// Emits an error with a fixit for the case of unnecessary cast over a

lib/Sema/TypeChecker.cpp

@@ -315,8 +315,7 @@ static void typeCheckFunctionsAndExternalDecls(SourceFile &SF, TypeChecker &TC)
   TC.definedFunctions.clear();
 }
 
-void swift::performTypeChecking(SourceFile &SF, TopLevelContext &TLC,
-                                unsigned StartElem) {
+void swift::performTypeChecking(SourceFile &SF, unsigned StartElem) {
   if (SF.ASTStage == SourceFile::TypeChecked)
     return;
 
@@ -366,7 +365,7 @@ void swift::performTypeChecking(SourceFile &SF, TopLevelContext &TLC,
       if (auto *TLCD = dyn_cast<TopLevelCodeDecl>(D)) {
         // Immediately perform global name-binding etc.
         TypeChecker::typeCheckTopLevelCodeDecl(TLCD);
-        TypeChecker::contextualizeTopLevelCode(TLC, TLCD);
+        TypeChecker::contextualizeTopLevelCode(TLCD);
       } else {
         TypeChecker::typeCheckDecl(D);
       }
@@ -375,7 +374,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);
 
@@ -1159,8 +1157,7 @@ class TypeChecker final {
   ///
   /// \returns true if any closures were found
   static bool contextualizeInitializer(Initializer *DC, Expr *init);
-  static void contextualizeTopLevelCode(TopLevelContext &TLC,
-                                        TopLevelCodeDecl *TLCD);
+  static void contextualizeTopLevelCode(TopLevelCodeDecl *TLCD);
 
   /// Return the type-of-reference of the given value.
   ///