Cache the results of SwiftASTContext::GetCompileUnitImports()

Every time a scratch context is requested, a call to
GetCompileUnitImports() is made, which is both expensive and
redundant. This patch adds a cache that keeps track of compile units
for which the imports were already performed. This results in a
measurable performance improvement on the LLDB test suite.

rdar://75381959

Author: adrian-prantl

lldb/source/Plugins/TypeSystem/Swift/SwiftASTContext.cpp

@@ -8314,8 +8314,8 @@ bool SwiftASTContext::GetImplicitImports(
     llvm::SmallVectorImpl<swift::AttributedImport<swift::ImportedModule>>
         &modules,
     Status &error) {
-  if (!GetCompileUnitImports(swift_ast_context, sc, stack_frame_wp, modules,
-                             error)) {
+  if (!swift_ast_context.GetCompileUnitImports(sc, stack_frame_wp, modules,
+                                               error)) {
     return false;
   }
 
@@ -8387,22 +8387,54 @@ bool SwiftASTContext::CacheUserImports(SwiftASTContext &swift_ast_context,
 }
 
 bool SwiftASTContext::GetCompileUnitImports(
-    SwiftASTContext &swift_ast_context, SymbolContext &sc,
-    lldb::StackFrameWP &stack_frame_wp,
+    SymbolContext &sc, lldb::StackFrameWP &stack_frame_wp,
     llvm::SmallVectorImpl<swift::AttributedImport<swift::ImportedModule>>
         &modules,
     Status &error) {
+  return GetCompileUnitImportsImpl(sc, stack_frame_wp, &modules, error);
+}
+
+void SwiftASTContext::PerformCompileUnitImports(
+    SymbolContext &sc, lldb::StackFrameWP &stack_frame_wp, Status &error) {
+  GetCompileUnitImportsImpl(sc, stack_frame_wp, nullptr, error);
+}
+
+static std::pair<Module *, lldb::user_id_t>
+GetCUSignature(CompileUnit &compile_unit) {
+  return {compile_unit.GetModule().get(), compile_unit.GetID()};
+}
+
+bool SwiftASTContext::GetCompileUnitImportsImpl(
+    SymbolContext &sc, lldb::StackFrameWP &stack_frame_wp,
+    llvm::SmallVectorImpl<swift::AttributedImport<swift::ImportedModule>>
+        *modules,
+    Status &error) {
+  CompileUnit *compile_unit = sc.comp_unit;
+  if (compile_unit)
+    // Check the cache if this compile unit's imports were previously
+    // requested.  If the caller didn't request the list of imported
+    // modules then there is nothing left to do for subsequent
+    // GetCompileUnitImportsImpl() calls as the previously loaded
+    // modules should still be loaded.  The fact the we
+    // unconditionally return true does not matter because the only
+    // way to get here is through void PerformCompileUnitImports(),
+    // which discards the return value.
+    if (!m_cu_imports.insert(GetCUSignature(*compile_unit)).second)
+      // List of imports isn't requested and we already processed this CU?
+      if (!modules)
+        return true;
+
   // Import the Swift standard library and its dependencies.
   SourceModule swift_module;
   swift_module.path.emplace_back("Swift");
   auto *stdlib =
-      LoadOneModule(swift_module, swift_ast_context, stack_frame_wp, error);
+      LoadOneModule(swift_module, *this, stack_frame_wp, error);
   if (!stdlib)
     return false;
 
-  modules.emplace_back(swift::ImportedModule(stdlib));
+  if (modules)
+    modules->emplace_back(swift::ImportedModule(stdlib));
 
-  CompileUnit *compile_unit = sc.comp_unit;
   if (!compile_unit || compile_unit->GetLanguage() != lldb::eLanguageTypeSwift)
     return true;
 
@@ -8417,11 +8449,12 @@ bool SwiftASTContext::GetCompileUnitImports(
       continue;
 
     auto *loaded_module =
-        LoadOneModule(module, swift_ast_context, stack_frame_wp, error);
+        LoadOneModule(module, *this, stack_frame_wp, error);
     if (!loaded_module)
       return false;
 
-    modules.emplace_back(swift::ImportedModule(loaded_module));
+    if (modules)
+      modules->emplace_back(swift::ImportedModule(loaded_module));
   }
   return true;
 }

lldb/source/Plugins/TypeSystem/Swift/SwiftASTContext.h

@@ -773,15 +773,26 @@ class SwiftASTContext : public TypeSystemSwift {
                                lldb::StackFrameWP &stack_frame_wp,
                                swift::SourceFile &source_file, Status &error);
 
-  /// Retrieve the modules imported by the compilation unit.
-  static bool GetCompileUnitImports(
-      SwiftASTContext &swift_ast_context, SymbolContext &sc,
-      lldb::StackFrameWP &stack_frame_wp,
+  /// Retrieve/import the modules imported by the compilation
+  /// unit. Early-exists with false if there was an import failure.
+  bool GetCompileUnitImports(
+      SymbolContext &sc, lldb::StackFrameWP &stack_frame_wp,
       llvm::SmallVectorImpl<swift::AttributedImport<swift::ImportedModule>>
           &modules,
       Status &error);
 
+  /// Perform all the implicit imports for the current frame.
+  void PerformCompileUnitImports(SymbolContext &sc,
+                                 lldb::StackFrameWP &stack_frame_wp,
+                                 Status &error);
+
 protected:
+  bool GetCompileUnitImportsImpl(
+      SymbolContext &sc, lldb::StackFrameWP &stack_frame_wp,
+      llvm::SmallVectorImpl<swift::AttributedImport<swift::ImportedModule>>
+          *modules,
+      Status &error);
+
   /// This map uses the string value of ConstStrings as the key, and the
   /// TypeBase
   /// * as the value. Since the ConstString strings are uniqued, we can use
@@ -867,6 +878,8 @@ class SwiftASTContext : public TypeSystemSwift {
   /// respective result (true = loaded, false = failed to load).
   std::unordered_map<detail::SwiftLibraryLookupRequest, bool>
       library_load_cache;
+  /// A cache for GetCompileUnitImports();
+  llvm::DenseSet<std::pair<Module *, lldb::user_id_t>> m_cu_imports;
 
   typedef std::map<Module *, std::vector<lldb::DataBufferSP>> ASTFileDataMap;
   ASTFileDataMap m_ast_file_data_map;

lldb/source/Target/Target.cpp

@@ -2557,10 +2557,7 @@ llvm::Optional<SwiftASTContextReader> Target::GetScratchSwiftASTContext(
       !swift_ast_ctx->HasFatalErrors()) {
     StackFrameWP frame_wp(frame_sp);
     SymbolContext sc = frame_sp->GetSymbolContext(lldb::eSymbolContextEverything);
-    llvm::SmallVector<swift::AttributedImport<swift::ImportedModule>, 16>
-        modules;
-    swift_ast_ctx->GetCompileUnitImports(*swift_ast_ctx, sc, frame_wp, modules,
-                                         error);
+    swift_ast_ctx->PerformCompileUnitImports(sc, frame_wp, error);
   }
 
   if (!swift_ast_ctx)