Don't Mutate ModuleDecl's File List When Creating Synthesized Files

include/swift/AST/FileUnit.h

@@ -57,10 +57,6 @@ class FileUnit : public DeclContext, public ASTAllocated<FileUnit> {
 
   /// Returns the synthesized file for this source file, creating one and
   /// inserting it into the module if it does not exist.
-  ///
-  /// \warning Because this function mutates the parent module's list of files,
-  ///          it will invalidate the iterators of any upstream callers of
-  ///          \c ModuleDecl::getFiles().
   SynthesizedFileUnit &getOrCreateSynthesizedFile();
 
   /// Look up a (possibly overloaded) value set at top-level scope

include/swift/AST/Module.h

@@ -317,6 +317,15 @@ class ModuleDecl
     return { Files.begin(), Files.size() };
   }
 
+  /// Add the given file to this module.
+  ///
+  /// FIXME: Remove this function from public view. Modules never need to add
+  /// files once they are created.
+  ///
+  /// \warning There are very few safe points to call this function once a
+  /// \c ModuleDecl has been created. If you find yourself needing to insert
+  /// a file in the middle of e.g. semantic analysis, use a \c
+  /// SynthesizedFileUnit instead.
   void addFile(FileUnit &newFile);
 
   /// Creates a map from \c #filePath strings to corresponding \c #fileID

lib/AST/Module.cpp

@@ -320,13 +320,13 @@ SourceLookupCache::SourceLookupCache(const ModuleDecl &M) {
   FrontendStatsTracer tracer(M.getASTContext().Stats,
                              "module-populate-cache");
   for (const FileUnit *file : M.getFiles()) {
-    if (auto *SFU = dyn_cast<SynthesizedFileUnit>(file)) {
-      addToUnqualifiedLookupCache(SFU->getTopLevelDecls(), false);
-      continue;
-    }
     auto *SF = cast<SourceFile>(file);
     addToUnqualifiedLookupCache(SF->getTopLevelDecls(), false);
     addToUnqualifiedLookupCache(SF->getHoistedDecls(), false);
+
+    if (auto *SFU = file->getSynthesizedFile()) {
+      addToUnqualifiedLookupCache(SFU->getTopLevelDecls(), false);
+    }
   }
 }
 
@@ -548,9 +548,13 @@ SourceFile *CodeCompletionFileRequest::evaluate(Evaluator &evaluator,
   llvm_unreachable("Couldn't find the completion file?");
 }
 
-#define FORWARD(name, args) \
-  for (const FileUnit *file : getFiles()) \
-    file->name args;
+#define FORWARD(name, args)                                                    \
+  for (const FileUnit *file : getFiles()) {                                    \
+    file->name args;                                                           \
+    if (auto *synth = file->getSynthesizedFile()) {                            \
+      synth->name args;                                                        \
+    }                                                                          \
+  }
 
 SourceLookupCache &ModuleDecl::getSourceLookupCache() const {
   if (!Cache) {
@@ -3064,16 +3068,9 @@ SynthesizedFileUnit &FileUnit::getOrCreateSynthesizedFile() {
       return *thisSynth;
     SynthesizedFile = new (getASTContext()) SynthesizedFileUnit(*this);
     SynthesizedFileAndKind.setPointer(SynthesizedFile);
-    // FIXME: Mutating the module in-flight is not a good idea. Any
-    // callers above us in the stack that are iterating over
-    // the module's files will have their iterators invalidated. There's
-    // a strong chance that whatever analysis led to this function being
-    // called is doing just that!
-    //
-    // Instead we ought to just call ModuleDecl::clearLookupCache() here
-    // and patch out the places looking for synthesized files hanging off of
-    // source files.
-    getParentModule()->addFile(*SynthesizedFile);
+    // Rebuild the source lookup caches now that we have a synthesized file
+    // full of declarations to look into.
+    getParentModule()->clearLookupCache();
   }
   return *SynthesizedFile;
 }

lib/Frontend/Frontend.cpp

@@ -1168,19 +1168,13 @@ bool CompilerInstance::loadPartialModulesAndImplicitImports(
 bool CompilerInstance::forEachFileToTypeCheck(
     llvm::function_ref<bool(SourceFile &)> fn) {
   if (isWholeModuleCompilation()) {
-    // FIXME: Do not refactor this to use an iterator as long as
-    // ModuleDecl::addFile is called during Sema. Synthesized files pushed
-    // during semantic analysis will cause iterator invalidation here.
-    // See notes in SourceFile::getOrCreateSynthesizedFile() for more.
-    unsigned i = 0;
-    while (i < getMainModule()->getFiles().size()) {
-      auto *SF = dyn_cast<SourceFile>(getMainModule()->getFiles()[i++]);
+    for (auto fileName : getMainModule()->getFiles()) {
+      auto *SF = dyn_cast<SourceFile>(fileName);
       if (!SF) {
         continue;
       }
       if (fn(*SF))
         return true;
-      ;
     }
   } else {
     for (auto *SF : getPrimarySourceFiles()) {

lib/IRGen/IRGen.cpp

@@ -1096,8 +1096,9 @@ GeneratedModule IRGenRequest::evaluate(Evaluator &evaluator,
     for (auto *file : filesToEmit) {
       if (auto *nextSF = dyn_cast<SourceFile>(file)) {
         IGM.emitSourceFile(*nextSF);
-      } else if (auto *nextSFU = dyn_cast<SynthesizedFileUnit>(file)) {
-        IGM.emitSynthesizedFileUnit(*nextSFU);
+        if (auto *synthSFU = file->getSynthesizedFile()) {
+          IGM.emitSynthesizedFileUnit(*synthSFU);
+        }
       } else {
         file->collectLinkLibraries([&IGM](LinkLibrary LinkLib) {
           IGM.addLinkLibrary(LinkLib);
@@ -1339,11 +1340,15 @@ static void performParallelIRGeneration(IRGenDescriptor desc) {
 
   for (auto *File : M->getFiles()) {
     if (auto *SF = dyn_cast<SourceFile>(File)) {
-      CurrentIGMPtr IGM = irgen.getGenModule(SF);
-      IGM->emitSourceFile(*SF);
-    } else if (auto *nextSFU = dyn_cast<SynthesizedFileUnit>(File)) {
-      CurrentIGMPtr IGM = irgen.getGenModule(nextSFU);
-      IGM->emitSynthesizedFileUnit(*nextSFU);
+      {
+        CurrentIGMPtr IGM = irgen.getGenModule(SF);
+        IGM->emitSourceFile(*SF);
+      }
+
+      if (auto *synthSFU = File->getSynthesizedFile()) {
+        CurrentIGMPtr IGM = irgen.getGenModule(synthSFU);
+        IGM->emitSynthesizedFileUnit(*synthSFU);
+      }
     } else {
       File->collectLinkLibraries([&](LinkLibrary LinkLib) {
         irgen.getPrimaryIGM()->addLinkLibrary(LinkLib);

lib/IRGen/IRGenRequests.cpp

@@ -68,9 +68,6 @@ TinyPtrVector<FileUnit *> IRGenDescriptor::getFilesToEmit() const {
   TinyPtrVector<FileUnit *> files;
   files.push_back(primary);
 
-  if (auto *synthesizedFile = primary->getSynthesizedFile())
-    files.push_back(synthesizedFile);
-
   return files;
 }
 

lib/TBDGen/TBDGen.cpp

@@ -1234,6 +1234,10 @@ void TBDGenVisitor::visit(const TBDGenDescriptor &desc) {
   llvm::for_each(Modules, [&](ModuleDecl *M) {
     for (auto *file : M->getFiles()) {
       visitFile(file);
+
+      // Visit synthesized file, if it exists.
+      if (auto *synthesizedFile = file->getSynthesizedFile())
+        visitFile(synthesizedFile);
     }
   });
 }