Sema: Only run extension binding once

There's no need to walk all imports of all source files, and bind the
extensions defined therein. The only time that a non-main module
contains source files is when using -enable-source-import, and we
can just explicitly call bindExtensions() in the right place.

Author: slavapestov

include/swift/Subsystems.h

@@ -140,7 +140,7 @@ namespace swift {
   void performPCMacro(SourceFile &SF);
 
   /// Bind all 'extension' visible from \p SF to the extended nominal.
-  void bindExtensions(SourceFile &SF);
+  void bindExtensions(ModuleDecl &mod);
 
   /// Once import resolution is complete, this walks the AST to resolve types
   /// and diagnose problems therein.

lib/Frontend/Frontend.cpp

@@ -806,12 +806,13 @@ void CompilerInstance::performSemaUpTo(SourceFile::ASTStage_t LimitStage) {
   }) && "some files have not yet had their imports resolved");
   MainModule->setHasResolvedImports();
 
-  forEachFileToTypeCheck([&](SourceFile &SF) {
-    if (LimitStage == SourceFile::ImportsResolved) {
-      bindExtensions(SF);
-      return;
-    }
+  bindExtensions(*MainModule);
+
+  // If the limiting AST stage is import resolution, we're done.
+  if (LimitStage == SourceFile::ImportsResolved)
+    return;
 
+  forEachFileToTypeCheck([&](SourceFile &SF) {
     performTypeChecking(SF);
 
     // Parse the SIL decls if needed.
@@ -822,11 +823,6 @@ void CompilerInstance::performSemaUpTo(SourceFile::ASTStage_t LimitStage) {
     }
   });
 
-  // If the limiting AST stage is import resolution, we're done.
-  if (LimitStage <= SourceFile::ImportsResolved) {
-    return;
-  }
-
   finishTypeChecking();
 }
 

lib/IDE/CompletionInstance.cpp

@@ -446,7 +446,7 @@ bool CompletionInstance::performCachedOperationIfPossible(
     // Re-process the whole file (parsing will be lazily triggered). Still
     // re-use imported modules.
     performImportResolution(*newSF);
-    bindExtensions(*newSF);
+    bindExtensions(*newM);
 
     CurrentModule = newM;
     traceDC = newM;

lib/IDE/REPLCodeCompletion.cpp

@@ -239,7 +239,7 @@ doCodeCompletion(SourceFile &SF, StringRef EnteredCode, unsigned *BufferID,
   newModule->addFile(newSF);
 
   performImportResolution(newSF);
-  bindExtensions(newSF);
+  bindExtensions(*newModule);
 
   performCodeCompletionSecondPass(newSF, *CompletionCallbacksFactory);
 

lib/Immediate/REPL.cpp

@@ -184,6 +184,8 @@ typeCheckREPLInput(ModuleDecl *MostRecentModule, StringRef Name,
   auto &REPLInputFile =
       *new (Ctx) SourceFile(*REPLModule, SourceFileKind::REPL, BufferID);
   REPLModule->addFile(REPLInputFile);
+  performImportResolution(REPLInputFile);
+  bindExtensions(*REPLModule);
   performTypeChecking(REPLInputFile);
   return REPLModule;
 }

lib/Sema/SourceLoader.cpp

@@ -131,6 +131,7 @@ ModuleDecl *SourceLoader::loadModule(SourceLoc importLoc,
   importMod->addFile(*importFile);
   performImportResolution(*importFile);
   importMod->setHasResolvedImports();
+  bindExtensions(*importMod);
   return importMod;
 }
 

lib/Sema/TypeChecker.cpp

@@ -227,7 +227,7 @@ static void bindExtensionToNominal(ExtensionDecl *ext,
   nominal->addExtension(ext);
 }
 
-static void bindExtensions(SourceFile &SF) {
+void swift::bindExtensions(ModuleDecl &mod) {
   // Utility function to try and resolve the extended type without diagnosing.
   // If we succeed, we go ahead and bind the extension. Otherwise, return false.
   auto tryBindExtension = [&](ExtensionDecl *ext) -> bool {
@@ -245,20 +245,15 @@ static void bindExtensions(SourceFile &SF) {
   // resolved to a worklist.
   SmallVector<ExtensionDecl *, 8> worklist;
 
-  // FIXME: The current source file needs to be handled specially, because of
-  // private extensions.
-  for (auto import : namelookup::getAllImports(&SF)) {
-    // FIXME: Respect the access path?
-    for (auto file : import.importedModule->getFiles()) {
-      auto SF = dyn_cast<SourceFile>(file);
-      if (!SF)
-        continue;
+  for (auto file : mod.getFiles()) {
+    auto *SF = dyn_cast<SourceFile>(file);
+    if (!SF)
+      continue;
 
-      for (auto D : SF->getTopLevelDecls()) {
-        if (auto ED = dyn_cast<ExtensionDecl>(D))
-          if (!tryBindExtension(ED))
-            worklist.push_back(ED);
-      }
+    for (auto D : SF->getTopLevelDecls()) {
+      if (auto ED = dyn_cast<ExtensionDecl>(D))
+        if (!tryBindExtension(ED))
+          worklist.push_back(ED);
     }
   }
 
@@ -354,11 +349,6 @@ TypeCheckSourceFileRequest::evaluate(Evaluator &eval, SourceFile *SF) const {
       TypeChecker::buildTypeRefinementContextHierarchy(*SF);
     }
 
-    // Resolve extensions. This has to occur first during type checking,
-    // because the extensions need to be wired into the AST for name lookup
-    // to work.
-    ::bindExtensions(*SF);
-
     // Type check the top-level elements of the source file.
     for (auto D : SF->getTopLevelDecls()) {
       if (auto *TLCD = dyn_cast<TopLevelCodeDecl>(D)) {
@@ -702,10 +692,6 @@ TypeChecker::getDeclTypeCheckingSemantics(ValueDecl *decl) {
   return DeclTypeCheckingSemantics::Normal;
 }
 
-void swift::bindExtensions(SourceFile &SF) {
-  ::bindExtensions(SF);
-}
-
 LookupResult
 swift::lookupSemanticMember(DeclContext *DC, Type ty, DeclName name) {
   return TypeChecker::lookupMember(DC, ty, DeclNameRef(name), None);