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

See #59144 for more on why this is a bad idea.

Patch out the synthesized file unit accessor to only clear the source cache, then patch up all the places that were assuming they could iterate over the module's file list and see synthesized files.

rdar://94164512

Author: CodaFi

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/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);
@@ -1341,9 +1342,10 @@ static void performParallelIRGeneration(IRGenDescriptor desc) {
     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);
+      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);
     }
   });
 }