[lldb][SwiftASTContext] Teach clang module loading from CAS

Teach SwiftASTContext to load clang module dependencies from CAS instead
of FileSystem.

Author: cachemeifyoucan

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

@@ -913,6 +913,51 @@ static std::string GetClangModulesCacheProperty() {
   return std::string(path);
 }
 
+static void ConfigureCASStorage(SwiftASTContext *m_ast_context,
+                                FileSpec CandidateConfigSearchPath) {
+  // Config CAS from properties.
+  llvm::cas::CASConfiguration cas_config;
+  cas_config.CASPath =
+      ModuleList::GetGlobalModuleListProperties().GetCASOnDiskPath().GetPath();
+  cas_config.PluginPath =
+      ModuleList::GetGlobalModuleListProperties().GetCASPluginPath().GetPath();
+  cas_config.PluginOptions =
+      ModuleList::GetGlobalModuleListProperties().GetCASPluginOptions();
+
+  auto &m_description = m_ast_context->GetDescription();
+
+  if (!cas_config.CASPath.empty()) {
+    if (auto maybe_cas = cas_config.createDatabases()) {
+      m_ast_context->SetCASStorage(std::move(maybe_cas->first),
+                                   std::move(maybe_cas->second));
+      m_ast_context->GetCASOptions().Config = cas_config;
+      LOG_PRINTF(GetLog(LLDBLog::Types),
+                 "Setup CAS from module list properties with cas path: %s",
+                 cas_config.CASPath.c_str());
+      return;
+    } else
+      llvm::consumeError(maybe_cas.takeError());
+  }
+
+  // Try search from candidiate path.
+  auto search_config = llvm::cas::CASConfiguration::createFromSearchConfigFile(
+      CandidateConfigSearchPath.GetPath());
+  if (!search_config)
+    return;
+
+  if (auto maybe_cas = search_config->second.createDatabases()) {
+    m_ast_context->SetCASStorage(std::move(maybe_cas->first),
+                                 std::move(maybe_cas->second));
+    m_ast_context->GetCASOptions().Config = search_config->second;
+    LOG_PRINTF(GetLog(LLDBLog::Types), "Setup CAS from cas config file: %s",
+               search_config->first.c_str());
+    return;
+  } else
+    llvm::consumeError(maybe_cas.takeError());
+
+  return;
+}
+
 SwiftASTContext::ScopedDiagnostics::ScopedDiagnostics(
     swift::DiagnosticConsumer &consumer)
     : m_consumer(consumer),
@@ -1930,41 +1975,80 @@ void SwiftASTContext::AddExtraClangCC1Args(
     return;
   }
 
-  // Clear module cache key and other CAS options to load modules from disk
-  // directly.
-  invocation.getFrontendOpts().ModuleCacheKeys.clear();
-  invocation.getCASOpts() = clang::CASOptions();
-
-  // Ignore CAS info inside modules when loading.
-  invocation.getFrontendOpts().ModuleLoadIgnoreCAS = true;
-
   // Add options to allow clang importer to do implicit module build.
   invocation.getLangOpts().ImplicitModules = true;
   invocation.getHeaderSearchOpts().ImplicitModuleMaps = true;
   invocation.getHeaderSearchOpts().ModuleCachePath =
       GetCompilerInvocation().getClangModuleCachePath().str();
 
-  // Remove non-existing modules in a systematic way.
-  auto CheckFileExists = [&](const std::string &file) -> bool {
-    if (llvm::sys::fs::exists(file))
-      return true;
-    std::string warn;
-    llvm::raw_string_ostream(warn)
-        << "Nonexistent explicit module file " << file;
-    AddDiagnostic(eSeverityWarning, warn);
-    return false;
-  };
-  for (auto it = invocation.getHeaderSearchOpts().PrebuiltModuleFiles.begin();
-       it != invocation.getHeaderSearchOpts().PrebuiltModuleFiles.end();) {
-    if (!CheckFileExists(it->second))
-      it = invocation.getHeaderSearchOpts().PrebuiltModuleFiles.erase(it);
-    else
-      ++it;
+  bool use_cas_module = m_cas && m_action_cache;
+  if (use_cas_module) {
+    // Load from CAS.
+    invocation.getCASOpts().CASPath = GetCASOptions().Config.CASPath;
+    invocation.getCASOpts().PluginPath = GetCASOptions().Config.PluginPath;
+    invocation.getCASOpts().PluginOptions =
+        GetCASOptions().Config.PluginOptions;
+
+    // Check the module availability in CAS, if not, fallback to regular load.
+    auto CheckModuleInCAS = [&](const std::string &key) {
+      auto id = m_cas->parseID(key);
+      if (!id) {
+        llvm::consumeError(id.takeError());
+        return false;
+      }
+      auto lookup = m_action_cache->get(*id);
+      if (!lookup) {
+        llvm::consumeError(lookup.takeError());
+        return false;
+      }
+      return (bool)*lookup;
+    };
+
+    use_cas_module = llvm::all_of(
+        invocation.getFrontendOpts().ModuleCacheKeys, [&](const auto &entry) {
+          auto exist = CheckModuleInCAS(entry.second);
+          if (!exist) {
+            LOG_PRINTF(GetLog(LLDBLog::Types),
+                       "module '%s' cannot be load "
+                       "from CAS using key: %s, fallback to "
+                       "load from file system",
+                       entry.first.c_str(), entry.second.c_str());
+          }
+          return exist;
+        });
+  }
+
+  if (!use_cas_module) {
+    // Clear module cache key and other CAS options to load modules from disk
+    // directly.
+    invocation.getFrontendOpts().ModuleCacheKeys.clear();
+    invocation.getCASOpts() = clang::CASOptions();
+
+    // Ignore CAS info inside modules when loading.
+    invocation.getFrontendOpts().ModuleLoadIgnoreCAS = true;
+
+    // Remove non-existing modules in a systematic way.
+    auto CheckFileExists = [&](const std::string &file) -> bool {
+      if (llvm::sys::fs::exists(file))
+        return true;
+      std::string warn;
+      llvm::raw_string_ostream(warn)
+          << "Nonexistent explicit module file " << file;
+      AddDiagnostic(eSeverityWarning, warn);
+      return false;
+    };
+    for (auto it = invocation.getHeaderSearchOpts().PrebuiltModuleFiles.begin();
+         it != invocation.getHeaderSearchOpts().PrebuiltModuleFiles.end();) {
+      if (!CheckFileExists(it->second))
+        it = invocation.getHeaderSearchOpts().PrebuiltModuleFiles.erase(it);
+      else
+        ++it;
+    }
+    invocation.getFrontendOpts().ModuleFiles.erase(
+        llvm::remove_if(invocation.getFrontendOpts().ModuleFiles,
+                        [&](const auto &mod) { return !CheckFileExists(mod); }),
+        invocation.getFrontendOpts().ModuleFiles.end());
   }
-  invocation.getFrontendOpts().ModuleFiles.erase(
-      llvm::remove_if(invocation.getFrontendOpts().ModuleFiles,
-                      [&](const auto &mod) { return !CheckFileExists(mod); }),
-      invocation.getFrontendOpts().ModuleFiles.end());
 
   invocation.generateCC1CommandLine(
       [&](const llvm::Twine &arg) { dest.push_back(arg.str()); });
@@ -2546,6 +2630,8 @@ SwiftASTContext::CreateInstance(lldb::LanguageType language, Module &module,
     }
   }
 
+  ConfigureCASStorage(swift_ast_sp.get(), module.GetFileSpec());
+
   // The serialized triple is the triple of the last binary
   // __swiftast section that was processed. Instead of relying on
   // the section contents order, we overwrite the triple in the
@@ -3027,6 +3113,8 @@ lldb::TypeSystemSP SwiftASTContext::CreateInstance(
     }
   }
 
+  ConfigureCASStorage(swift_ast_sp.get(), sc.module_sp->GetFileSpec());
+
   std::string resource_dir = HostInfo::GetSwiftResourceDir(
       triple, swift_ast_sp->GetPlatformSDKPath());
   ConfigureResourceDirs(swift_ast_sp->GetCompilerInvocation(), resource_dir,

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

@@ -311,6 +311,12 @@ class SwiftASTContext : public TypeSystemSwift {
     m_platform_sdk_path = path.str();
   }
 
+  void SetCASStorage(std::shared_ptr<llvm::cas::ObjectStore> cas,
+                     std::shared_ptr<llvm::cas::ActionCache> action_cache) {
+    m_cas = std::move(cas);
+    m_action_cache = std::move(action_cache);
+  }
+
   /// \return the ExtraArgs of the ClangImporterOptions.
   const std::vector<std::string> &GetClangArguments();
 
@@ -987,6 +993,9 @@ class SwiftASTContext : public TypeSystemSwift {
       library_load_cache;
   /// A cache for GetCompileUnitImports();
   llvm::DenseSet<std::pair<Module *, lldb::user_id_t>> m_cu_imports;
+  /// ObjectStore and ActionCache;
+  std::shared_ptr<llvm::cas::ObjectStore> m_cas;
+  std::shared_ptr<llvm::cas::ActionCache> m_action_cache;
 
   typedef std::map<Module *, std::vector<lldb::DataBufferSP>> ASTFileDataMap;
   ASTFileDataMap m_ast_file_data_map;

lldb/test/Shell/Swift/caching.test

@@ -9,8 +9,21 @@
 # RUN:          -module-name main -o %t/main
 
 # RUN: %lldb %t/main -s %t/lldb.script 2>&1 | FileCheck %s
+
+## Setup CAS and try loading from CAS
+# RUN: sed "s|DIR|%/t|g" %t/lldb.script.template > %t/lldb_2.script
+# RUN: %lldb %t/main -s %t/lldb_2.script 2>&1 | FileCheck %s --check-prefix=CAS-LOAD --check-prefix=CHECK
+
+## Check fallback to file system.
+# RUN: rm -rf %t/cas
+# RUN: %lldb %t/main -s %t/lldb_2.script 2>&1 | FileCheck %s --check-prefix=CAS-FALLBACK --check-prefix=CHECK
+
+# CAS-FALLBACK:  operator()() -- module '{{.*}}' cannot be load from CAS using key
+# CAS-FALLBACK-SAME: fallback to load from file system
+# CAS-LOAD:      ConfigureCASStorage() -- Setup CAS from module list properties with cas path
 # CHECK:         LogConfiguration() --   Extra clang arguments
 # CHECK-COUNT-1: LogConfiguration() --     -triple
+# CAS-LOAD:      LogConfiguration() --     -fmodule-file-cache-key
 # CHECK:         (Int) ${{.*}} = 1
 
 //--- main.swift
@@ -28,3 +41,14 @@ run
 # Create a SwiftASTContext
 expr 1
 quit
+
+//--- lldb.script.template
+# Force loading from interface to simulate no binary module available.
+settings set symbols.swift-module-loading-mode prefer-interface
+settings set symbols.cas-path DIR/cas
+log enable lldb types
+b test
+run
+# Create a SwiftASTContext
+expr 1
+quit