Fix caching bug in MemoryBufferSerializedModuleLoader

By populating the memory cache before loading the module, we can avoid a cycle
where a module is imported that is an overlay, which then triggers
ClangImporter, which then (redundantly) triggers the import of the overlay
module, which would reimport the module again, since it's import is still
underway and it hasn't been entered into the cache yet.

rdar://118846313

include/swift/AST/ASTContext.h

@@ -1250,6 +1250,13 @@ public:
   /// in this context.
   void addLoadedModule(ModuleDecl *M);
 
+  /// Remove an externally-sourced module from the set of known loaded modules
+  /// in this context. Modules are added to the cache before being loaded to
+  /// avoid loading a module twice when there is a cyclic dependency between an
+  /// overlay and its Clang module. If a module import fails, the non-imported
+  /// module should be removed from the cache again.
+  void removeLoadedModule(Identifier RealName);
+
   /// Change the behavior of all loaders to ignore swiftmodules next to
   /// swiftinterfaces.
   void setIgnoreAdjacentModules(bool value);

lib/AST/ASTContext.cpp

@@ -2267,6 +2267,10 @@ void ASTContext::addLoadedModule(ModuleDecl *M) {
   getImpl().LoadedModules[M->getRealName()] = M;
 }
 
+void ASTContext::removeLoadedModule(Identifier RealName) {
+  getImpl().LoadedModules.erase(RealName);
+}
+
 void ASTContext::setIgnoreAdjacentModules(bool value) {
   IgnoreAdjacentModules = value;
 }

lib/Serialization/SerializedModuleLoader.cpp

@@ -1527,12 +1527,16 @@ MemoryBufferSerializedModuleLoader::loadModule(SourceLoc importLoc,
 
   auto *M = ModuleDecl::create(moduleID.Item, Ctx);
   SWIFT_DEFER { M->setHasResolvedImports(); };
+  if (AllowMemoryCache)
+    Ctx.addLoadedModule(M);
 
   auto *file = loadAST(*M, moduleID.Loc, /*moduleInterfacePath=*/"",
                        /*moduleInterfaceSourcePath=*/"",
                        std::move(moduleInputBuffer), {}, {}, isFramework);
-  if (!file)
+  if (!file) {
+    Ctx.removeLoadedModule(moduleID.Item);
     return nullptr;
+  }
 
   // The MemoryBuffer loader is used by LLDB during debugging. Modules imported
   // from .swift_ast sections are never produced from textual interfaces. By
@@ -1540,8 +1544,6 @@ MemoryBufferSerializedModuleLoader::loadModule(SourceLoc importLoc,
   if (BypassResilience)
     M->setBypassResilience();
   M->addFile(*file);
-  if (AllowMemoryCache)
-    Ctx.addLoadedModule(M);
   return M;
 }
 

test/DWARFImporter/basic.swift

@@ -16,9 +16,13 @@
 // RUN: %lldb-moduleimport-test -verbose -dump-module %t/a.out \
 // RUN:    -dummy-dwarfimporter | %FileCheck %s --check-prefix=SWIFTONLY
 
-// CHECK: Importing basic... ok!
-// FAIL: Importing basic... ok!
-// SWIFTONLY: Importing basic... ok!
+// CHECK: Importing basic...
+// CHECK: Import successful!
+// FAIL: Importing basic... 
+// FAIL: Import successful!
+// SWIFTONLY: Importing basic...
+// SWIFTONLY: Import successful!
+
 import ObjCModule
 
 let pureSwift = Int32(42)

test/DebugInfo/ASTSection-multi.swift

@@ -42,8 +42,11 @@
 // RUN:                     %t/a3.o %t/a3-mod.o
 
 // RUN: %lldb-moduleimport-test -verbose %t/a.out | %FileCheck %s
-// CHECK: Importing a0... ok!
-// CHECK: Importing a1... ok!
-// CHECK: Importing a2... ok!
-// CHECK: Importing a3... ok!
+// CHECK-DAG: Importing a0...
+// CHECK-DAG: Import successful!
+// CHECK-DAG: Importing a1...
+// CHECK-DAG: Import successful!
+// CHECK-DAG: Importing a2...
+// CHECK-DAG: Import successful!
+// CHECK-DAG: Importing a3...
 

test/DebugInfo/ASTSection-single.swift

@@ -11,4 +11,5 @@
 
 // RUN: %lldb-moduleimport-test -verbose %t/a0-mod.o | %FileCheck %s
 // CHECK: Path: {{.*}}/Inputs, framework=false, system=false
-// CHECK: Importing a0... ok!
+// CHECK: Importing a0...
+// CHECK: Import successful!

test/DebugInfo/ASTSection.swift

@@ -26,7 +26,8 @@ func objCUser(_ obj: ObjCClass) {}
 #endif
 
 // CHECK: - Target: {{.+}}-{{.+}}-{{.+}}
-// CHECK: Importing ASTSection... ok!
+// CHECK: Importing ASTSection...
+// CHECK: Import successful!
 
 // LINETABLE-CHECK-NOT: ASTSection
 

test/DebugInfo/ASTSectionOverlay.swift

@@ -0,0 +1,16 @@
+// REQUIRES: executable_test
+// REQUIRES: objc_interop
+
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift -emit-module %S/Inputs/overlay.swift -module-name ClangModuleWithOverlay -I %S/Inputs -emit-module-path %t/ClangModuleWithOverlay.swiftmodule
+// RUN: %target-build-swift -c %S/Inputs/overlay.swift -module-name ClangModuleWithOverlay -I %S/Inputs -o %t/ClangModuleWithOverlay.o -parse-as-library
+// RUN: %target-build-swift -emit-executable %s %t/ClangModuleWithOverlay.o -I %t -g -o %t/ASTSectionOverlay -module-name ASTSectionOverlay -emit-module -Xlinker -add_ast_path -Xlinker %t/ClangModuleWithOverlay.swiftmodule
+
+// RUN: %lldb-moduleimport-test -verbose %t/ASTSectionOverlay | %FileCheck %s
+// CHECK: Loading ClangModuleWithOverlay
+// CHECK-NOT: Loading (overlay) ClangModuleWithOverlay
+// CHECK-NOT: Loading{{.*}}ClangModuleWithOverlay
+
+import ClangModuleWithOverlay
+let c = ClangType(i: 0)
+fromSwiftOverlay()

test/DebugInfo/ASTSection_ObjC.swift

@@ -14,6 +14,7 @@
 // RUN: %lldb-moduleimport-test -verbose %t/ASTSection | %FileCheck %s --allow-empty --check-prefix=LINETABLE-CHECK
 
 // CHECK: - Target: {{.+}}-{{.+}}-{{.+}}
-// CHECK: Importing ASTSection... ok!
+// CHECK: Importing ASTSection...
+// CHECK: Import successful!
 
 // LINETABLE-CHECK-NOT: ASTSection

test/DebugInfo/ASTSection_linker.swift

@@ -21,4 +21,6 @@
 // CHECK: - Target: {{.+}}-{{.+}}-{{.+}}
 // CHECK: - SDK path: {{.*}}MacOS{{.*}}.sdk
 // CHECK: - -Xcc options: -working-directory {{.+}} -DA -DB
-// CHECK: Importing ASTSection... ok!
+// CHECK: Importing ASTSection...
+// CHECK: Import successful!
+

test/DebugInfo/Inputs/ClangModuleWithOverlay.h

@@ -0,0 +1,3 @@
+struct ClangType {
+  int i;
+};

test/DebugInfo/Inputs/module.modulemap

@@ -20,3 +20,7 @@ module Macro {
   header "Macro.h"
 }
 
+module ClangModuleWithOverlay {
+  header "ClangModuleWithOverlay.h"
+  export *
+}

test/DebugInfo/Inputs/overlay.swift

@@ -0,0 +1,3 @@
+@_exported import ClangModuleWithOverlay
+
+public func fromSwiftOverlay() {}

tools/lldb-moduleimport-test/lldb-moduleimport-test.cpp

@@ -381,7 +381,25 @@ int main(int argc, char **argv) {
     auto *ClangImporter = static_cast<swift::ClangImporter *>(
         CI.getASTContext().getClangModuleLoader());
     ClangImporter->setDWARFImporterDelegate(dummyDWARFImporter);
-  }
+ }
+
+  if (Verbose)
+    CI.getASTContext().SetPreModuleImportCallback(
+        [&](llvm::StringRef module_name,
+            swift::ASTContext::ModuleImportKind kind) {
+          switch (kind) {
+          case swift::ASTContext::Module:
+            llvm::outs() << "Loading " << module_name.str() << "\n";
+            break;
+          case swift::ASTContext::Overlay:
+            llvm::outs() << "Loading (overlay) " << module_name.str() << "\n";
+            break;
+          case swift::ASTContext::BridgingHeader:
+            llvm::outs() << "Compiling bridging header: " << module_name.str()
+                         << "\n";
+            break;
+          }
+        });
 
   llvm::SmallString<0> error;
   llvm::raw_svector_ostream errs(error);
@@ -401,7 +419,7 @@ int main(int argc, char **argv) {
   // Attempt to import all modules we found.
   for (auto path : modules) {
     if (Verbose)
-      llvm::outs() << "Importing " << path << "... ";
+      llvm::outs() << "Importing " << path << "...\n";
 
     swift::ImportPath::Module::Builder modulePath;
 #ifdef SWIFT_SUPPORTS_SUBMODULES
@@ -420,7 +438,7 @@ int main(int argc, char **argv) {
       return 1;
     }
     if (Verbose)
-      llvm::outs() << "ok!\n";
+      llvm::outs() << "Import successful!\n";
     if (DumpModule) {
       llvm::SmallVector<swift::Decl*, 10> Decls;
       Module->getTopLevelDecls(Decls);