[AST] Add callback to the module loading methods

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

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

@@ -1108,27 +1108,37 @@ class ASTContext final {
   /// be returned.
   ///
   /// \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,7 +2204,8 @@ 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))
@@ -2212,6 +2214,10 @@ ASTContext::getModule(ImportPath::Module ModulePath) {
   auto moduleID = ModulePath[0];
   for (auto &importer : getImpl().ModuleLoaders) {
     if (ModuleDecl *M = importer->loadModule(moduleID.Loc, ModulePath)) {
+      if (callback)
+        callback(
+            llvm::formatv("Loading module overlay for {0}", M->getNameStr())
+                .str());
       if (LangOpts.EnableModuleLoadingRemarks) {
         Diags.diagnose(ModulePath.getSourceRange().Start,
                        diag::module_loaded,
@@ -2224,7 +2230,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());
@@ -2238,30 +2246,40 @@ ModuleDecl *ASTContext::getOverlayModule(const FileUnit *FU) {
     if (importer.get() == getClangModuleLoader())
       continue;
     if (ModuleDecl *M = importer->loadModule(SourceLoc(), ModPath)) {
+      if (callback)
+        callback(
+            llvm::formatv("Loading module overlay for {0}", M->getNameStr())
+                .str());
       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);
   }