[AST] Add mechanism to notify module loading to other clients

This patch adds an optional callback function member to the AST Context.
The callback gets invoked when a new module (or overlay module) gets
loaded.

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

@@ -374,6 +374,10 @@ class ASTContext final {
   llvm::BumpPtrAllocator &
   getAllocator(AllocationArena arena = AllocationArena::Permanent) const;
 
+  /// An optional generic callback function invoked prior to importing a module.
+  mutable std::function<bool(llvm::StringRef ModuleName, bool IsOverlay)>
+      PreModuleImportCallback;
+
 public:
   /// Allocate - Allocate memory from the ASTContext bump pointer.
   void *Allocate(unsigned long bytes, unsigned alignment,
@@ -489,6 +493,10 @@ class ASTContext final {
   /// Set a new stats reporter.
   void setStatsReporter(UnifiedStatsReporter *stats);
 
+  /// Set the module import callback function.
+  void setPreModuleImportCallback(
+      std::function<bool(llvm::StringRef, bool)> Callback);
+
 public:
   /// getIdentifier - Return the uniqued and AST-Context-owned version of the
   /// specified string.
@@ -1107,6 +1115,9 @@ 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.
+
   /// \returns The requested module, or NULL if the module cannot be found.
   ModuleDecl *getModule(ImportPath::Module ModulePath);
 

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>
 
@@ -692,6 +693,12 @@ void ASTContext::setStatsReporter(UnifiedStatsReporter *stats) {
   Stats = stats;
 }
 
+/// Set a new module loading callback function.
+void ASTContext::setPreModuleImportCallback(
+    std::function<bool(llvm::StringRef, bool)> Callback) {
+  PreModuleImportCallback = Callback;
+}
+
 RC<syntax::SyntaxArena> ASTContext::getSyntaxArena() const {
   return getImpl().TheSyntaxArena;
 }
@@ -2210,6 +2217,8 @@ ASTContext::getModule(ImportPath::Module ModulePath) {
     return M;
 
   auto moduleID = ModulePath[0];
+  if (PreModuleImportCallback)
+    PreModuleImportCallback(moduleID.Item.str(), false /*=IsOverlay*/);
   for (auto &importer : getImpl().ModuleLoaders) {
     if (ModuleDecl *M = importer->loadModule(moduleID.Loc, ModulePath)) {
       if (LangOpts.EnableModuleLoadingRemarks) {
@@ -2237,9 +2246,14 @@ 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 (PreModuleImportCallback) {
+      SmallString<256> path;
+      ModPath.getString(path);
+      if (!path.empty())
+        PreModuleImportCallback(path.str(), true /*=IsOverlay*/);
     }
+    if (ModuleDecl *M = importer->loadModule(SourceLoc(), ModPath))
+      return M;
   }
 
   return nullptr;