[AST] Add callback to the module loading methods

This patch adds an optional callback function argument to each AST
module loading method.

This can be used for instance by other clients, such as lldb, to perform
actions when a module gets loaded, like showing progress to the end-user.

rdar://94165195

Signed-off-by: Med Ismail Bennani <[email protected]>

Author: medismailben

include/swift/AST/ASTContext.h

@@ -1107,28 +1107,44 @@ class ASTContext final {
   /// If a module by this name has already been loaded, the existing module will
   /// be returned.
   ///
+  /// \param ModulePath The module's \c ImportPath which describes
+  /// the name of the module being loaded, possibly including submodules.
+  ///
+  /// \param callback An optional generic function invoked prior to loading the
+  /// module.
+  ///
   /// \returns The requested module, or NULL if the module cannot be found.
-  ModuleDecl *getModule(ImportPath::Module ModulePath);
+  ModuleDecl *getModule(ImportPath::Module ModulePath,
+                        std::function<bool(llvm::StringRef)> callback = {});
 
   /// Attempts to load the matching overlay module for the given clang
   /// module into this ASTContext.
   ///
   /// \returns The Swift overlay module corresponding to the given Clang module,
   /// or NULL if the overlay module cannot be found.
-  ModuleDecl *getOverlayModule(const FileUnit *ClangModule);
+  ModuleDecl *
+  getOverlayModule(const FileUnit *ClangModule,
+                   std::function<bool(llvm::StringRef)> callback = {});
 
-  ModuleDecl *getModuleByName(StringRef ModuleName);
+  ModuleDecl *
+  getModuleByName(StringRef ModuleName,
+                  std::function<bool(llvm::StringRef)> callback = {});
 
-  ModuleDecl *getModuleByIdentifier(Identifier ModuleID);
+  ModuleDecl *
+  getModuleByIdentifier(Identifier ModuleID,
+                        std::function<bool(llvm::StringRef)> callback = {});
 
   /// Returns the standard library module, or null if the library isn't present.
   ///
   /// If \p loadIfAbsent is true, the ASTContext will attempt to load the module
   /// if it hasn't been set yet.
-  ModuleDecl *getStdlibModule(bool loadIfAbsent = false);
+  ModuleDecl *
+  getStdlibModule(bool loadIfAbsent = false,
+                  std::function<bool(llvm::StringRef)> callback = {});
 
-  ModuleDecl *getStdlibModule() const {
-    return const_cast<ASTContext *>(this)->getStdlibModule(false);
+  ModuleDecl *
+  getStdlibModule(std::function<bool(llvm::StringRef)> callback = {}) const {
+    return const_cast<ASTContext *>(this)->getStdlibModule(false, callback);
   }
 
   /// Insert an externally-sourced module into the set of known loaded modules

lib/AST/ASTContext.cpp

@@ -42,8 +42,8 @@
 #include "swift/AST/PropertyWrappers.h"
 #include "swift/AST/ProtocolConformance.h"
 #include "swift/AST/RawComment.h"
-#include "swift/AST/SearchPathOptions.h"
 #include "swift/AST/SILLayout.h"
+#include "swift/AST/SearchPathOptions.h"
 #include "swift/AST/SemanticAttrs.h"
 #include "swift/AST/SourceFile.h"
 #include "swift/AST/SubstitutionMap.h"
@@ -65,6 +65,7 @@
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/Compiler.h"
+#include "llvm/Support/FormatVariadic.h"
 #include <algorithm>
 #include <memory>
 
@@ -2203,14 +2204,17 @@ bool ASTContext::canImportModule(ImportPath::Module ModuleName,
 }
 
 ModuleDecl *
-ASTContext::getModule(ImportPath::Module ModulePath) {
+ASTContext::getModule(ImportPath::Module ModulePath,
+                      std::function<bool(llvm::StringRef)> callback) {
   assert(!ModulePath.empty());
 
   if (auto *M = getLoadedModule(ModulePath))
     return M;
 
   auto moduleID = ModulePath[0];
   for (auto &importer : getImpl().ModuleLoaders) {
+    if (callback)
+      callback(llvm::formatv("Loading module {0}", moduleID.Item.str()).str());
     if (ModuleDecl *M = importer->loadModule(moduleID.Loc, ModulePath)) {
       if (LangOpts.EnableModuleLoadingRemarks) {
         Diags.diagnose(ModulePath.getSourceRange().Start,
@@ -2224,7 +2228,9 @@ ASTContext::getModule(ImportPath::Module ModulePath) {
   return nullptr;
 }
 
-ModuleDecl *ASTContext::getOverlayModule(const FileUnit *FU) {
+ModuleDecl *
+ASTContext::getOverlayModule(const FileUnit *FU,
+                             std::function<bool(llvm::StringRef)> callback) {
   assert(FU && FU->getKind() == FileUnitKind::ClangModule &&
          "Overlays can only be retrieved for clang modules!");
   ImportPath::Module::Builder builder(FU->getParentModule()->getName());
@@ -2237,31 +2243,42 @@ ModuleDecl *ASTContext::getOverlayModule(const FileUnit *FU) {
   for (auto &importer : getImpl().ModuleLoaders) {
     if (importer.get() == getClangModuleLoader())
       continue;
-    if (ModuleDecl *M = importer->loadModule(SourceLoc(), ModPath)) {
-      return M;
+    if (callback) {
+      SmallString<256> path;
+      ModPath.getString(path);
+      if (!path.empty())
+        callback(llvm::formatv("Loading overlay module {0}", path).str());
     }
+    if (ModuleDecl *M = importer->loadModule(SourceLoc(), ModPath))
+      return M;
   }
 
   return nullptr;
 }
 
-ModuleDecl *ASTContext::getModuleByName(StringRef ModuleName) {
+ModuleDecl *
+ASTContext::getModuleByName(StringRef ModuleName,
+                            std::function<bool(llvm::StringRef)> callback) {
   ImportPath::Module::Builder builder(*this, ModuleName, /*separator=*/'.');
-  return getModule(builder.get());
+  return getModule(builder.get(), callback);
 }
 
-ModuleDecl *ASTContext::getModuleByIdentifier(Identifier ModuleID) {
+ModuleDecl *ASTContext::getModuleByIdentifier(
+    Identifier ModuleID, std::function<bool(llvm::StringRef)> callback) {
   ImportPath::Module::Builder builder(ModuleID);
-  return getModule(builder.get());
+  return getModule(builder.get(), callback);
 }
 
-ModuleDecl *ASTContext::getStdlibModule(bool loadIfAbsent) {
+ModuleDecl *
+ASTContext::getStdlibModule(bool loadIfAbsent,
+                            std::function<bool(llvm::StringRef)> callback) {
   if (TheStdlibModule)
     return TheStdlibModule;
 
   if (loadIfAbsent) {
     auto mutableThis = const_cast<ASTContext*>(this);
-    TheStdlibModule = mutableThis->getModuleByIdentifier(StdlibModuleName);
+    TheStdlibModule =
+        mutableThis->getModuleByIdentifier(StdlibModuleName, callback);
   } else {
     TheStdlibModule = getLoadedModule(StdlibModuleName);
   }