Merge pull request swiftlang#77666 from hamishknight/lets-try-this-again

[AST] Remove `ModuleDecl::addFile`

Author: hamishknight

include/swift/AST/FileUnit.h

@@ -486,9 +486,11 @@ void simple_display(llvm::raw_ostream &out, const FileUnit *file);
 inline FileUnit &ModuleDecl::getMainFile(FileUnitKind expectedKind) const {
   assert(expectedKind != FileUnitKind::Source &&
          "must use specific source kind; see getMainSourceFile");
-  assert(!Files.empty() && "No files added yet");
-  assert(Files.front()->getKind() == expectedKind);
-  return *Files.front();
+
+  auto files = getFiles();
+  assert(!files.empty() && "No files in module");
+  assert(files.front()->getKind() == expectedKind);
+  return *files.front();
 }
 
 } // end namespace swift

include/swift/AST/Module.h

@@ -314,7 +314,9 @@ class ModuleDecl
   // FIXME: Do we really need to bloat all modules with this?
   DebuggerClient *DebugClient = nullptr;
 
-  SmallVector<FileUnit *, 2> Files;
+  /// The list of files in the module. This is guaranteed to be set once module
+  /// construction has completed. It must not be mutated afterwards.
+  std::optional<SmallVector<FileUnit *, 2>> Files;
 
   llvm::SmallDenseMap<Identifier, SmallVector<OverlayFile *, 1>>
     declaredCrossImports;
@@ -359,24 +361,55 @@ class ModuleDecl
   /// Used by the debugger to bypass resilient access to fields.
   bool BypassResilience = false;
 
-  ModuleDecl(Identifier name, ASTContext &ctx, ImplicitImportInfo importInfo);
+public:
+  using PopulateFilesFn = llvm::function_ref<void(
+      ModuleDecl *, llvm::function_ref<void(FileUnit *)>)>;
+
+private:
+  ModuleDecl(Identifier name, ASTContext &ctx, ImplicitImportInfo importInfo,
+             PopulateFilesFn populateFiles, bool isMainModule);
 
 public:
   /// Creates a new module with a given \p name.
   ///
   /// \param importInfo Information about which modules should be implicitly
   /// imported by each file of this module.
-  static ModuleDecl *
-  create(Identifier name, ASTContext &ctx,
-         ImplicitImportInfo importInfo = ImplicitImportInfo()) {
-    return new (ctx) ModuleDecl(name, ctx, importInfo);
+  /// \param populateFiles A function which populates the files for the module.
+  /// Once called, the module's list of files may not change.
+  static ModuleDecl *create(Identifier name, ASTContext &ctx,
+                            ImplicitImportInfo importInfo,
+                            PopulateFilesFn populateFiles) {
+    return new (ctx) ModuleDecl(name, ctx, importInfo, populateFiles,
+                                /*isMainModule*/ false);
+  }
+
+  /// Creates a new module with a given \p name.
+  ///
+  /// \param populateFiles A function which populates the files for the module.
+  /// Once called, the module's list of files may not change.
+  static ModuleDecl *create(Identifier name, ASTContext &ctx,
+                            PopulateFilesFn populateFiles) {
+    return new (ctx) ModuleDecl(name, ctx, ImplicitImportInfo(), populateFiles,
+                                /*isMainModule*/ false);
   }
 
+  /// Creates a new main module with a given \p name. The main module is the
+  /// module being built by the compiler, containing the primary source files.
+  ///
+  /// \param importInfo Information about which modules should be implicitly
+  /// imported by each file of this module.
+  /// \param populateFiles A function which populates the files for the module.
+  /// Once called, the module's list of files may not change.
   static ModuleDecl *createMainModule(ASTContext &ctx, Identifier name,
-                                      ImplicitImportInfo iinfo) {
-    auto *Mod = ModuleDecl::create(name, ctx, iinfo);
-    Mod->Bits.ModuleDecl.IsMainModule = true;
-    return Mod;
+                                      ImplicitImportInfo iinfo,
+                                      PopulateFilesFn populateFiles) {
+    return new (ctx) ModuleDecl(name, ctx, iinfo, populateFiles,
+                                /*isMainModule*/ true);
+  }
+
+  /// Creates an empty module with a given \p name.
+  static ModuleDecl *createEmpty(Identifier name, ASTContext &ctx) {
+    return create(name, ctx, ImplicitImportInfo(), [](auto, auto) {});
   }
 
   using Decl::getASTContext;
@@ -398,24 +431,11 @@ class ModuleDecl
   /// Only to be called by MemoryBufferSerializedModuleLoader.
   void setBypassResilience() { BypassResilience = true; }
 
-  ArrayRef<FileUnit *> getFiles() {
-    ASSERT(!Files.empty() || failedToLoad());
-    return Files;
+  ArrayRef<FileUnit *> getFiles() const {
+    ASSERT(Files.has_value() &&
+           "Attempting to query files before setting them");
+    return *Files;
   }
-  ArrayRef<const FileUnit *> getFiles() const {
-    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);
 
   /// Produces the source file that contains the given source location, or
   /// \c nullptr if the source location isn't in this module.

include/swift/AST/SourceFile.h

@@ -823,8 +823,8 @@ inline SourceFile::ParsingOptions operator|(SourceFile::ParsingFlags lhs,
 }
 
 inline SourceFile &ModuleDecl::getMainSourceFile() const {
-  assert(!Files.empty() && "No files added yet");
-  return *cast<SourceFile>(Files.front());
+  assert(!getFiles().empty() && "No files in module");
+  return *cast<SourceFile>(getFiles().front());
 }
 
 inline FileUnit *ModuleDecl::EntryPointInfoTy::getEntryPointFile() const {

include/swift/IDETool/SyntacticMacroExpansion.h

@@ -45,35 +45,32 @@ class SyntacticMacroExpansionInstance {
   DiagnosticEngine Diags{SourceMgr};
   std::unique_ptr<ASTContext> Ctx;
   ModuleDecl *TheModule = nullptr;
+  SourceFile *SF = nullptr;
   llvm::StringMap<MacroDecl *> MacroDecls;
 
-  /// Create 'SourceFile' using the buffer.
-  swift::SourceFile *getSourceFile(llvm::MemoryBuffer *inputBuf);
-
   /// Synthesize 'MacroDecl' AST object to use the expansion.
   swift::MacroDecl *
   getSynthesizedMacroDecl(swift::Identifier name,
                           const MacroExpansionSpecifier &expansion);
 
-  /// Expand single 'expansion' in SF.
-  void expand(swift::SourceFile *SF,
-                    const MacroExpansionSpecifier &expansion,
-                    SourceEditConsumer &consumer);
+  /// Expand single 'expansion'.
+  void expand(const MacroExpansionSpecifier &expansion,
+              SourceEditConsumer &consumer);
 
 public:
   SyntacticMacroExpansionInstance() {}
 
-  /// Setup the instance with \p args .
+  /// Setup the instance with \p args and a given \p inputBuf.
   bool setup(StringRef SwiftExecutablePath, ArrayRef<const char *> args,
+             llvm::MemoryBuffer *inputBuf,
              std::shared_ptr<PluginRegistry> plugins, std::string &error);
 
   ASTContext &getASTContext() { return *Ctx; }
 
-  /// Expand all macros in \p inputBuf and send the edit results to \p consumer.
-  /// Expansions are specified by \p expansions .
-  void expandAll(llvm::MemoryBuffer *inputBuf,
-                     ArrayRef<MacroExpansionSpecifier> expansions,
-                     SourceEditConsumer &consumer);
+  /// Expand all macros and send the edit results to \p consumer. Expansions are
+  /// specified by \p expansions .
+  void expandAll(ArrayRef<MacroExpansionSpecifier> expansions,
+                 SourceEditConsumer &consumer);
 };
 
 /// Manager object to vend 'SyntacticMacroExpansionInstance'.
@@ -86,9 +83,11 @@ class SyntacticMacroExpansion {
                           std::shared_ptr<PluginRegistry> Plugins)
       : SwiftExecutablePath(SwiftExecutablePath), Plugins(Plugins) {}
 
-  /// Get instance configured with the specified compiler arguments.
+  /// Get instance configured with the specified compiler arguments and
+  /// input buffer.
   std::shared_ptr<SyntacticMacroExpansionInstance>
-  getInstance(ArrayRef<const char *> args, std::string &error);
+  getInstance(ArrayRef<const char *> args, llvm::MemoryBuffer *inputBuf,
+              std::string &error);
 };
 
 } // namespace ide

include/swift/Subsystems.h

@@ -117,6 +117,9 @@ namespace swift {
                               bool TokenizeInterpolatedString = true,
                               ArrayRef<Token> SplitTokens = ArrayRef<Token>());
 
+  /// Perform import resolution for the module.
+  void performImportResolution(ModuleDecl *M);
+
   /// This walks the AST to resolve imports.
   void performImportResolution(SourceFile &SF);
 

lib/AST/ASTContext.cpp

@@ -734,8 +734,10 @@ ConstraintCheckerArenaRAII::~ConstraintCheckerArenaRAII() {
 }
 
 static ModuleDecl *createBuiltinModule(ASTContext &ctx) {
-  auto M = ModuleDecl::create(ctx.getIdentifier(BUILTIN_NAME), ctx);
-  M->addFile(*new (ctx) BuiltinUnit(*M));
+  auto *M = ModuleDecl::create(ctx.getIdentifier(BUILTIN_NAME), ctx,
+                               [&](ModuleDecl *M, auto addFile) {
+    addFile(new (ctx) BuiltinUnit(*M));
+  });
   M->setHasResolvedImports();
   return M;
 }

lib/AST/ASTDemangler.cpp

@@ -1135,7 +1135,7 @@ ASTBuilder::getAcceptableTypeDeclCandidate(ValueDecl *decl,
 
 DeclContext *ASTBuilder::getNotionalDC() {
   if (!NotionalDC) {
-    NotionalDC = ModuleDecl::create(Ctx.getIdentifier(".RemoteAST"), Ctx);
+    NotionalDC = ModuleDecl::createEmpty(Ctx.getIdentifier(".RemoteAST"), Ctx);
     NotionalDC = new (Ctx) TopLevelCodeDecl(NotionalDC);
   }
   return NotionalDC;

lib/AST/Module.cpp

@@ -737,7 +737,9 @@ void SourceLookupCache::lookupClassMember(ImportPath::Access accessPath,
 //===----------------------------------------------------------------------===//
 
 ModuleDecl::ModuleDecl(Identifier name, ASTContext &ctx,
-                       ImplicitImportInfo importInfo)
+                       ImplicitImportInfo importInfo,
+                       PopulateFilesFn populateFiles,
+                       bool isMainModule)
     : DeclContext(DeclContextKind::Module, nullptr),
       TypeDecl(DeclKind::Module, &ctx, name, SourceLoc(), {}),
       ImportInfo(importInfo) {
@@ -756,7 +758,7 @@ ModuleDecl::ModuleDecl(Identifier name, ASTContext &ctx,
   Bits.ModuleDecl.ImplicitDynamicEnabled = 0;
   Bits.ModuleDecl.IsSystemModule = 0;
   Bits.ModuleDecl.IsNonSwiftModule = 0;
-  Bits.ModuleDecl.IsMainModule = 0;
+  Bits.ModuleDecl.IsMainModule = isMainModule;
   Bits.ModuleDecl.HasIncrementalInfo = 0;
   Bits.ModuleDecl.HasHermeticSealAtLink = 0;
   Bits.ModuleDecl.IsEmbeddedSwiftModule = 0;
@@ -766,6 +768,22 @@ ModuleDecl::ModuleDecl(Identifier name, ASTContext &ctx,
   Bits.ModuleDecl.CXXStdlibKind = 0;
   Bits.ModuleDecl.AllowNonResilientAccess = 0;
   Bits.ModuleDecl.SerializePackageEnabled = 0;
+
+  // Populate the module's files.
+  SmallVector<FileUnit *, 2> files;
+  populateFiles(this, [&](FileUnit *file) {
+    // If this is a LoadedFile, make sure it loaded without error.
+    assert(!(isa<LoadedFile>(file) &&
+             cast<LoadedFile>(file)->hadLoadError()));
+
+    // Require Main and REPL files to be the first file added.
+    assert(files.empty() ||
+           !isa<SourceFile>(file) ||
+           cast<SourceFile>(file)->Kind == SourceFileKind::Library ||
+           cast<SourceFile>(file)->Kind == SourceFileKind::SIL);
+    files.push_back(file);
+  });
+  Files.emplace(std::move(files));
 }
 
 void ModuleDecl::setIsSystemModule(bool flag) {
@@ -788,20 +806,6 @@ ImplicitImportList ModuleDecl::getImplicitImports() const {
                            {});
 }
 
-void ModuleDecl::addFile(FileUnit &newFile) {
-  // If this is a LoadedFile, make sure it loaded without error.
-  assert(!(isa<LoadedFile>(newFile) &&
-           cast<LoadedFile>(newFile).hadLoadError()));
-
-  // Require Main and REPL files to be the first file added.
-  assert(Files.empty() ||
-         !isa<SourceFile>(newFile) ||
-         cast<SourceFile>(newFile).Kind == SourceFileKind::Library ||
-         cast<SourceFile>(newFile).Kind == SourceFileKind::SIL);
-  Files.push_back(&newFile);
-  clearLookupCache();
-}
-
 SourceFile *ModuleDecl::getSourceFileContainingLocation(SourceLoc loc) {
   if (loc.isInvalid())
     return nullptr;
@@ -1210,7 +1214,7 @@ void SourceFile::lookupObjCMethods(
 }
 
 bool ModuleDecl::shouldCollectDisplayDecls() const {
-  for (const FileUnit *file : Files) {
+  for (const FileUnit *file : getFiles()) {
     if (!file->shouldCollectDisplayDecls())
       return false;
   }
@@ -1791,6 +1795,15 @@ void SourceFile::dumpSeparatelyImportedOverlays() const {
 
 void ModuleDecl::getImportedModulesForLookup(
     SmallVectorImpl<ImportedModule> &modules) const {
+  // FIXME: We can unfortunately end up in a cycle where we attempt to query
+  // the files for a serialized module while attempting to load its LoadedFile
+  // unit. This is because both SerializedModuleLoader and ClangImporter
+  // kick the computation of transitive module dependencies, which can end up
+  // pointing back to the original module in cases where it's an overlay. We
+  // ought to be more lazy there. This is also why
+  // `SourceFile::getImportedModules` cannot assume imports have been resolved.
+  if (!Files.has_value())
+    return;
   FORWARD(getImportedModulesForLookup, (modules));
 }
 

lib/ClangImporter/ClangImporter.cpp

@@ -1583,11 +1583,14 @@ ClangImporter::create(ASTContext &ctx,
     = clangContext.Selectors.getSelector(2, setObjectForKeyedSubscriptIdents);
 
   // Set up the imported header module.
-  auto *importedHeaderModule =
-      ModuleDecl::create(ctx.getIdentifier(CLANG_HEADER_MODULE_NAME), ctx);
-  importer->Impl.ImportedHeaderUnit =
-    new (ctx) ClangModuleUnit(*importedHeaderModule, importer->Impl, nullptr);
-  importedHeaderModule->addFile(*importer->Impl.ImportedHeaderUnit);
+  auto *importedHeaderModule = ModuleDecl::create(
+      ctx.getIdentifier(CLANG_HEADER_MODULE_NAME), ctx,
+      [&](ModuleDecl *importedHeaderModule, auto addFile) {
+        importer->Impl.ImportedHeaderUnit = new (ctx)
+            ClangModuleUnit(*importedHeaderModule, importer->Impl, nullptr);
+        addFile(importer->Impl.ImportedHeaderUnit);
+      });
+
   importedHeaderModule->setHasResolvedImports();
   importedHeaderModule->setIsNonSwiftModule(true);
 
@@ -2807,17 +2810,19 @@ ClangModuleUnit *ClangImporter::Implementation::getWrapperForModule(
   if (auto mainModule = SwiftContext.MainModule) {
     implicitImportInfo = mainModule->getImplicitImportInfo();
   }
-  auto wrapper = ModuleDecl::create(name, SwiftContext, implicitImportInfo);
+  ClangModuleUnit *file = nullptr;
+  auto wrapper = ModuleDecl::create(name, SwiftContext, implicitImportInfo,
+                                    [&](ModuleDecl *wrapper, auto addFile) {
+    file = new (SwiftContext) ClangModuleUnit(*wrapper, *this, underlying);
+    addFile(file);
+  });
   wrapper->setIsSystemModule(underlying->IsSystem);
   wrapper->setIsNonSwiftModule();
   wrapper->setHasResolvedImports();
   if (!underlying->ExportAsModule.empty())
     wrapper->setExportAsName(
         SwiftContext.getIdentifier(underlying->ExportAsModule));
 
-  auto file = new (SwiftContext) ClangModuleUnit(*wrapper, *this,
-                                                 underlying);
-  wrapper->addFile(*file);
   SwiftContext.getClangModuleLoader()->findOverlayFiles(diagLoc, wrapper, file);
   cacheEntry.setPointer(file);
 

lib/ClangImporter/DWARFImporter.cpp

@@ -116,23 +116,25 @@ ModuleDecl *ClangImporter::Implementation::loadModuleDWARF(
   if (it != DWARFModuleUnits.end())
     return it->second->getParentModule();
 
-  auto *decl = ModuleDecl::create(name, SwiftContext);
-  decl->setIsNonSwiftModule();
-  decl->setHasResolvedImports();
-  auto *wrapperUnit = new (SwiftContext) DWARFModuleUnit(*decl, *this);
-  DWARFModuleUnits.insert({name, wrapperUnit});
-  decl->addFile(*wrapperUnit);
+  auto *M = ModuleDecl::create(name, SwiftContext,
+                               [&](ModuleDecl *M, auto addFile) {
+    auto *wrapperUnit = new (SwiftContext) DWARFModuleUnit(*M, *this);
+    DWARFModuleUnits.insert({name, wrapperUnit});
+    addFile(wrapperUnit);
+  });
+  M->setIsNonSwiftModule();
+  M->setHasResolvedImports();
 
   // Force load overlay modules for all imported modules.
-  assert(namelookup::getAllImports(decl).size() == 1 &&
-         namelookup::getAllImports(decl).front().importedModule == decl &&
+  assert(namelookup::getAllImports(M).size() == 1 &&
+         namelookup::getAllImports(M).front().importedModule == M &&
          "DWARF module depends on additional modules?");
 
   // Register the module with the ASTContext so it is available for lookups.
   if (!SwiftContext.getLoadedModule(name))
-    SwiftContext.addLoadedModule(decl);
+    SwiftContext.addLoadedModule(M);
 
-  return decl;
+  return M;
 }
 
 // This function exists to defeat the lazy member importing mechanism. The