[DependencyScanning] Add ability to scan TU with multiple module imports

clang/include/clang/Frontend/FrontendActions.h

@@ -34,12 +34,15 @@ class InitOnlyAction : public FrontendAction {
 
 /// Preprocessor-based frontend action that also loads PCH files.
 class ReadPCHAndPreprocessAction : public FrontendAction {
+  ArrayRef<StringRef> ImportModules;
   void ExecuteAction() override;
 
   std::unique_ptr<ASTConsumer> CreateASTConsumer(CompilerInstance &CI,
                                                  StringRef InFile) override;
 
 public:
+  ReadPCHAndPreprocessAction(ArrayRef<StringRef> ImportModules)
+      : ImportModules(ImportModules) {}
   bool usesPreprocessorOnly() const override { return false; }
 };
 
@@ -319,15 +322,6 @@ class PrintPreprocessedAction : public PreprocessorFrontendAction {
   bool hasPCHSupport() const override { return true; }
 };
 
-class GetDependenciesByModuleNameAction : public PreprocessOnlyAction {
-  StringRef ModuleName;
-  void ExecuteAction() override;
-
-public:
-  GetDependenciesByModuleNameAction(StringRef ModuleName)
-      : ModuleName(ModuleName) {}
-};
-
 }  // end namespace clang
 
 #endif

clang/include/clang/Tooling/DependencyScanning/DependencyScanningTool.h

@@ -174,14 +174,19 @@ class DependencyScanningTool {
   /// \param LookupModuleOutput This function is called to fill in
   ///                           "-fmodule-file=", "-o" and other output
   ///                           arguments for dependencies.
+  /// \param AdditionalModules Additional modules that need to be imported.
+  /// \param TUBuffer Optional memory buffer for translation unit input. If
+  ///                 TUBuffer is nullopt, the input should be included in the
+  ///                 Commandline already.
   ///
   /// \returns a \c StringError with the diagnostic output if clang errors
   /// occurred, \c TranslationUnitDeps otherwise.
-  llvm::Expected<TranslationUnitDeps>
-  getTranslationUnitDependencies(const std::vector<std::string> &CommandLine,
-                                 StringRef CWD,
-                                 const llvm::DenseSet<ModuleID> &AlreadySeen,
-                                 LookupModuleOutputCallback LookupModuleOutput);
+  llvm::Expected<TranslationUnitDeps> getTranslationUnitDependencies(
+      const std::vector<std::string> &CommandLine, StringRef CWD,
+      const llvm::DenseSet<ModuleID> &AlreadySeen,
+      LookupModuleOutputCallback LookupModuleOutput,
+      ArrayRef<StringRef> AdditionalModules = {},
+      std::optional<llvm::MemoryBufferRef> TUBuffer = std::nullopt);
 
   /// Given a compilation context specified via the Clang driver command-line,
   /// gather modular dependencies of module with the given name, and return the

clang/include/clang/Tooling/DependencyScanning/DependencyScanningWorker.h

@@ -18,6 +18,7 @@
 #include "llvm/CAS/CachingOnDiskFileSystem.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/FileSystem.h"
+#include "llvm/Support/MemoryBufferRef.h"
 #include <optional>
 #include <string>
 
@@ -119,24 +120,27 @@ class DependencyScanningWorker {
                            llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> FS);
 
   /// Run the dependency scanning tool for a given clang driver command-line,
-  /// and report the discovered dependencies to the provided consumer. If \p
-  /// ModuleName isn't empty, this function reports the dependencies of module
-  /// \p ModuleName.
+  /// and report the discovered dependencies to the provided consumer. If
+  /// TUBuffer is not nullopt, it is used as TU input for the dependency
+  /// scanning. Otherwise, the input should be included as part of the
+  /// command-line. ModuleNames is a list of additional modules to be imported
+  /// in addition to what is referenced in TU.
   ///
   /// \returns false if clang errors occurred (with diagnostics reported to
   /// \c DiagConsumer), true otherwise.
-  bool computeDependencies(StringRef WorkingDirectory,
-                           const std::vector<std::string> &CommandLine,
-                           DependencyConsumer &DepConsumer,
-                           DependencyActionController &Controller,
-                           DiagnosticConsumer &DiagConsumer,
-                           std::optional<StringRef> ModuleName = std::nullopt);
+  bool computeDependencies(
+      StringRef WorkingDirectory, const std::vector<std::string> &CommandLine,
+      DependencyConsumer &DepConsumer, DependencyActionController &Controller,
+      DiagnosticConsumer &DiagConsumer, ArrayRef<StringRef> ModuleNames = {},
+      std::optional<llvm::MemoryBufferRef> TUBuffer = std::nullopt);
+
   /// \returns A \c StringError with the diagnostic output if clang errors
   /// occurred, success otherwise.
   llvm::Error computeDependencies(
       StringRef WorkingDirectory, const std::vector<std::string> &CommandLine,
       DependencyConsumer &Consumer, DependencyActionController &Controller,
-      std::optional<StringRef> ModuleName = std::nullopt);
+      ArrayRef<StringRef> ModuleNames = {},
+      std::optional<llvm::MemoryBufferRef> TUBuffer = std::nullopt);
 
   /// Scan from a compiler invocation.
   /// If \p DiagGenerationAsCompilation is true it will generate error

clang/lib/Frontend/FrontendActions.cpp

@@ -67,7 +67,8 @@ InitOnlyAction::CreateASTConsumer(CompilerInstance &CI, StringRef InFile) {
 void InitOnlyAction::ExecuteAction() {
 }
 
-// Basically PreprocessOnlyAction::ExecuteAction.
+// DependencyScanningPreprocessAction imports additional modules first, then
+// preprocess the source file.
 void ReadPCHAndPreprocessAction::ExecuteAction() {
   Preprocessor &PP = getCompilerInstance().getPreprocessor();
 
@@ -77,6 +78,23 @@ void ReadPCHAndPreprocessAction::ExecuteAction() {
   Token Tok;
   // Start parsing the specified input file.
   PP.EnterMainSourceFile();
+
+  // Import additional modules first.
+  CompilerInstance &CI = getCompilerInstance();
+  SourceManager &SM = PP.getSourceManager();
+  FileID MainFileID = SM.getMainFileID();
+  PPCallbacks *CB = PP.getPPCallbacks();
+  unsigned Offset = -ImportModules.size() - 1;
+  for (auto Module : ImportModules) {
+    SourceLocation IncludeLoc =
+        SM.getLocForEndOfFile(MainFileID).getLocWithOffset(Offset++);
+    IdentifierInfo *ModuleID = PP.getIdentifierInfo(Module);
+    SmallVector<std::pair<IdentifierInfo *, SourceLocation>, 2> Path;
+    Path.push_back(std::make_pair(ModuleID, IncludeLoc));
+    auto ModResult = CI.loadModule(IncludeLoc, Path, Module::Hidden, false);
+    CB->moduleImport(IncludeLoc, Path, ModResult);
+  }
+
   do {
     PP.Lex(Tok);
   } while (Tok.isNot(tok::eof));
@@ -1200,18 +1218,3 @@ void PrintDependencyDirectivesSourceMinimizerAction::ExecuteAction() {
   printDependencyDirectivesAsSource(FromFile.getBuffer(), Directives,
                                     llvm::outs());
 }
-
-void GetDependenciesByModuleNameAction::ExecuteAction() {
-  CompilerInstance &CI = getCompilerInstance();
-  Preprocessor &PP = CI.getPreprocessor();
-  SourceManager &SM = PP.getSourceManager();
-  FileID MainFileID = SM.getMainFileID();
-  PP.EnterSourceFile(MainFileID, nullptr, SourceLocation());
-  SourceLocation FileStart = SM.getLocForStartOfFile(MainFileID);
-  SmallVector<std::pair<IdentifierInfo *, SourceLocation>, 2> Path;
-  IdentifierInfo *ModuleID = PP.getIdentifierInfo(ModuleName);
-  Path.push_back(std::make_pair(ModuleID, FileStart));
-  auto ModResult = CI.loadModule(FileStart, Path, Module::Hidden, false);
-  PPCallbacks *CB = PP.getPPCallbacks();
-  CB->moduleImport(SourceLocation(), Path, ModResult);
-}

clang/lib/Tooling/DependencyScanning/DependencyScanningTool.cpp

@@ -268,11 +268,13 @@ llvm::Expected<TranslationUnitDeps>
 DependencyScanningTool::getTranslationUnitDependencies(
     const std::vector<std::string> &CommandLine, StringRef CWD,
     const llvm::DenseSet<ModuleID> &AlreadySeen,
-    LookupModuleOutputCallback LookupModuleOutput) {
+    LookupModuleOutputCallback LookupModuleOutput,
+    ArrayRef<StringRef> AddedModules,
+    std::optional<llvm::MemoryBufferRef> TUBuffer) {
   FullDependencyConsumer Consumer(AlreadySeen);
   auto Controller = createActionController(LookupModuleOutput);
-  llvm::Error Result =
-      Worker.computeDependencies(CWD, CommandLine, Consumer, *Controller);
+  llvm::Error Result = Worker.computeDependencies(
+      CWD, CommandLine, Consumer, *Controller, AddedModules, TUBuffer);
   if (Result)
     return std::move(Result);
   return Consumer.takeTranslationUnitDeps();

clang/lib/Tooling/DependencyScanning/DependencyScanningWorker.cpp

@@ -34,6 +34,8 @@
 #include "llvm/CAS/ObjectStore.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/Error.h"
+#include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/Path.h"
 #include "llvm/Support/PrefixMapper.h"
 #include "llvm/TargetParser/Host.h"
 #include <optional>
@@ -400,7 +402,7 @@ class DependencyScanningAction : public tooling::ToolAction {
       ScanningOutputFormat Format, ScanningOptimizations OptimizeArgs,
       bool EagerLoadModules, bool DisableFree, bool EmitDependencyFile,
       bool DiagGenerationAsCompilation, const CASOptions &CASOpts,
-      std::optional<StringRef> ModuleName = std::nullopt,
+      ArrayRef<StringRef> ModuleNames = {},
       raw_ostream *VerboseOS = nullptr)
       : WorkingDirectory(WorkingDirectory), Consumer(Consumer),
         Controller(Controller), DepFS(std::move(DepFS)),
@@ -409,7 +411,7 @@ class DependencyScanningAction : public tooling::ToolAction {
         EagerLoadModules(EagerLoadModules), DisableFree(DisableFree),
         CASOpts(CASOpts), EmitDependencyFile(EmitDependencyFile),
         DiagGenerationAsCompilation(DiagGenerationAsCompilation),
-        ModuleName(ModuleName), VerboseOS(VerboseOS) {
+        ModuleNames(ModuleNames), VerboseOS(VerboseOS) {
     // The FullIncludeTree output format completely subsumes header search and
     // VFS optimizations due to how it works. Disable these optimizations so
     // we're not doing unneeded work.
@@ -602,12 +604,8 @@ class DependencyScanningAction : public tooling::ToolAction {
     // Avoid some checks and module map parsing when loading PCM files.
     ScanInstance.getPreprocessorOpts().ModulesCheckRelocated = false;
 
-    std::unique_ptr<FrontendAction> Action;
-
-    if (ModuleName)
-      Action = std::make_unique<GetDependenciesByModuleNameAction>(*ModuleName);
-    else
-      Action = std::make_unique<ReadPCHAndPreprocessAction>();
+    std::unique_ptr<FrontendAction> Action =
+        std::make_unique<ReadPCHAndPreprocessAction>(ModuleNames);
 
     // Normally this would be handled by GeneratePCHAction
     if (ScanInstance.getFrontendOpts().ProgramAction == frontend::GeneratePCH)
@@ -692,7 +690,7 @@ class DependencyScanningAction : public tooling::ToolAction {
   const CASOptions &CASOpts;
   bool EmitDependencyFile = false;
   bool DiagGenerationAsCompilation;
-  std::optional<StringRef> ModuleName;
+  ArrayRef<StringRef> ModuleNames;
   std::optional<CompilerInstance> ScanInstanceStorage;
   std::shared_ptr<ModuleDepCollector> MDC;
   std::vector<std::string> LastCC1Arguments;
@@ -747,7 +745,8 @@ DependencyScanningWorker::getOrCreateFileManager() const {
 llvm::Error DependencyScanningWorker::computeDependencies(
     StringRef WorkingDirectory, const std::vector<std::string> &CommandLine,
     DependencyConsumer &Consumer, DependencyActionController &Controller,
-    std::optional<StringRef> ModuleName) {
+    ArrayRef<StringRef> ModuleNames,
+    std::optional<llvm::MemoryBufferRef> TUBuffer) {
   std::vector<const char *> CLI;
   for (const std::string &Arg : CommandLine)
     CLI.push_back(Arg.c_str());
@@ -761,7 +760,7 @@ llvm::Error DependencyScanningWorker::computeDependencies(
   TextDiagnosticPrinter DiagPrinter(DiagnosticsOS, DiagOpts.release());
 
   if (computeDependencies(WorkingDirectory, CommandLine, Consumer, Controller,
-                          DiagPrinter, ModuleName))
+                          DiagPrinter, ModuleNames, TUBuffer))
     return llvm::Error::success();
   return llvm::make_error<llvm::StringError>(DiagnosticsOS.str(),
                                              llvm::inconvertibleErrorCode());
@@ -831,29 +830,41 @@ static bool createAndRunToolInvocation(
 bool DependencyScanningWorker::computeDependencies(
     StringRef WorkingDirectory, const std::vector<std::string> &CommandLine,
     DependencyConsumer &Consumer, DependencyActionController &Controller,
-    DiagnosticConsumer &DC, std::optional<StringRef> ModuleName) {
+    DiagnosticConsumer &DC, ArrayRef<StringRef> ModuleNames,
+    std::optional<llvm::MemoryBufferRef> TUBuffer) {
   // Reset what might have been modified in the previous worker invocation.
   BaseFS->setCurrentWorkingDirectory(WorkingDirectory);
 
   std::optional<std::vector<std::string>> ModifiedCommandLine;
   llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> ModifiedFS;
 
-  // If we're scanning based on a module name alone, we don't expect the client
-  // to provide us with an input file. However, the driver really wants to have
-  // one. Let's just make it up to make the driver happy.
-  if (ModuleName) {
+  // Create an empty buffer if module names are provided.
+  std::unique_ptr<llvm::MemoryBuffer> SrcBuffer;
+  if (!ModuleNames.empty() && !TUBuffer) {
+    SrcBuffer = llvm::MemoryBuffer::getMemBuffer("");
+    TUBuffer = SrcBuffer->getMemBufferRef();
+  }
+
+  // If the TU is passed as a buffer to the dependency scanner, overlay that
+  // with InMemoryFS and setup the input correctly in command-line.
+  if (TUBuffer) {
     auto OverlayFS =
         llvm::makeIntrusiveRefCnt<llvm::vfs::OverlayFileSystem>(BaseFS);
     auto InMemoryFS =
         llvm::makeIntrusiveRefCnt<llvm::vfs::InMemoryFileSystem>();
     InMemoryFS->setCurrentWorkingDirectory(WorkingDirectory);
 
-    SmallString<128> FakeInputPath;
-    // TODO: We should retry the creation if the path already exists.
-    llvm::sys::fs::createUniquePath(*ModuleName + "-%%%%%%%%.input",
-                                    FakeInputPath,
-                                    /*MakeAbsolute=*/false);
-    InMemoryFS->addFile(FakeInputPath, 0, llvm::MemoryBuffer::getMemBuffer(""));
+    // Use a stable name "<compiler-generated>" in root directory.
+    llvm::SmallString<256> FakeInputPath(
+        llvm::sys::path::root_path(WorkingDirectory));
+    llvm::sys::path::append(FakeInputPath, "<compiler-generated>");
+
+    // Extend the buffer so it starts with some white space reserved for
+    // importing additional modules.
+    std::string FileContent(ModuleNames.size(), ' ');
+    FileContent.append(TUBuffer->getBuffer().str());
+    InMemoryFS->addFile(FakeInputPath, 0,
+                        llvm::MemoryBuffer::getMemBufferCopy(FileContent));
 
     llvm::IntrusiveRefCntPtr<llvm::vfs::FileSystem> InMemoryOverlay =
         InMemoryFS;
@@ -900,7 +911,7 @@ bool DependencyScanningWorker::computeDependencies(
                                   DisableFree,
                                   /*EmitDependencyFile=*/false,
                                   /*DiagGenerationAsCompilation=*/false, getCASOpts(),
-                                  ModuleName);
+                                  ModuleNames);
   bool Success = false;
   if (FinalCommandLine[1] == "-cc1") {
     Success = createAndRunToolInvocation(FinalCommandLine, Action, *FileMgr,
@@ -976,7 +987,7 @@ void DependencyScanningWorker::computeDependenciesFromCompilerInvocation(
       ScanningOptimizations::Default, /*DisableFree=*/false, EagerLoadModules,
       /*EmitDependencyFile=*/!DepFile.empty(), DiagGenerationAsCompilation,
       getCASOpts(),
-      /*ModuleName=*/std::nullopt, VerboseOS);
+      /*ModuleNames=*/{}, VerboseOS);
 
   // Ignore result; we're just collecting dependencies.
   //