[clang][cas] Extend fix for private modules to module builds

In 6034ccd we fixed an issue where if a public module is imported
via PCH we cannot import the private module in a TU because the
modulemap is not parsed. This is also an issue when building a module
that imports the private module and the TU used a PCH, because even
though we don't import the PCH itself, we still use the
prebuilt/explicit module found via the PCH instead of building an
implicit scanner module.

This commit generalizes the original fix to handle all prebuilt/explicit
modules.

rdar://107446573

Author: benlangmuir

clang/lib/Tooling/DependencyScanning/IncludeTreeActionController.cpp

@@ -142,30 +142,33 @@ class IncludeTreeBuilder {
   std::optional<llvm::Error> ErrorToReport;
 };
 
-/// A utility for adding \c PPCallbacks to a compiler instance at the
-/// appropriate time.
-struct PPCallbacksDependencyCollector : public DependencyCollector {
-  using MakeCB =
+/// A utility for adding \c PPCallbacks and/or \cASTReaderListener to a compiler
+/// instance at the appropriate time.
+struct AttachOnlyDependencyCollector : public DependencyCollector {
+  using MakePPCB =
       llvm::unique_function<std::unique_ptr<PPCallbacks>(Preprocessor &)>;
-  MakeCB Create;
-  PPCallbacksDependencyCollector(MakeCB Create) : Create(std::move(Create)) {}
+  using MakeASTReaderL =
+      llvm::unique_function<std::unique_ptr<ASTReaderListener>(ASTReader &R)>;
+  MakePPCB CreatePPCB;
+  MakeASTReaderL CreateASTReaderL;
+  AttachOnlyDependencyCollector(MakePPCB CreatePPCB, MakeASTReaderL CreateL)
+      : CreatePPCB(std::move(CreatePPCB)),
+        CreateASTReaderL(std::move(CreateL)) {}
+
   void attachToPreprocessor(Preprocessor &PP) final {
-    std::unique_ptr<PPCallbacks> CB = Create(PP);
-    assert(CB);
-    PP.addPPCallbacks(std::move(CB));
+    if (CreatePPCB) {
+      std::unique_ptr<PPCallbacks> CB = CreatePPCB(PP);
+      assert(CB);
+      PP.addPPCallbacks(std::move(CB));
+    }
   }
-};
-/// A utility for adding \c ASTReaderListener to a compiler instance at the
-/// appropriate time.
-struct ASTReaderListenerDependencyCollector : public DependencyCollector {
-  using MakeL =
-      llvm::unique_function<std::unique_ptr<ASTReaderListener>(ASTReader &R)>;
-  MakeL Create;
-  ASTReaderListenerDependencyCollector(MakeL Create) : Create(std::move(Create)) {}
+
   void attachToASTReader(ASTReader &R) final {
-    std::unique_ptr<ASTReaderListener> L = Create(R);
-    assert(L);
-    R.addListener(std::move(L));
+    if (CreateASTReaderL) {
+      std::unique_ptr<ASTReaderListener> L = CreateASTReaderL(R);
+      assert(L);
+      R.addListener(std::move(L));
+    }
   }
 };
 
@@ -244,27 +247,29 @@ struct IncludeTreePPCallbacks : public PPCallbacks {
 /// MODULE_DIRECTORY record, and ignores the result.
 class LookupPCHModulesListener : public ASTReaderListener {
 public:
-  LookupPCHModulesListener(Preprocessor &PP) : PP(PP) {}
+  LookupPCHModulesListener(ASTReader &R) : Reader(R) {}
 
 private:
-  void ReadModuleName(StringRef ModuleName) final {
-    if (PCHFinished)
-      return;
-    // Match MODULE_DIRECTORY: allow full search and ignore failure to find
-    // the module.
-    (void)PP.getHeaderSearchInfo().lookupModule(
-        ModuleName, SourceLocation(), /*AllowSearch=*/true,
-        /*AllowExtraModuleMapSearch=*/true);
-  }
   void visitModuleFile(StringRef Filename,
                        serialization::ModuleKind Kind) final {
-    if (Kind == serialization::MK_PCH)
-      PCHFinished = true;
+    // Any prebuilt or explicit modules seen during scanning are "full" modules
+    // rather than implicitly built scanner modules.
+    if (Kind == serialization::MK_PrebuiltModule ||
+        Kind == serialization::MK_ExplicitModule) {
+      serialization::ModuleManager &Manager = Reader.getModuleManager();
+      serialization::ModuleFile *MF = Manager.lookupByFileName(Filename);
+      assert(MF && "module file missing in visitModuleFile");
+      // Match MODULE_DIRECTORY: allow full search and ignore failure to find
+      // the module.
+      HeaderSearch &HS = Reader.getPreprocessor().getHeaderSearchInfo();
+      (void)HS.lookupModule(MF->ModuleName, SourceLocation(),
+                            /*AllowSearch=*/true,
+                            /*AllowExtraModuleMapSearch=*/true);
+    }
   }
 
 private:
-  Preprocessor &PP;
-  bool PCHFinished = false;
+  ASTReader &Reader;
 };
 } // namespace
 
@@ -317,21 +322,15 @@ Error IncludeTreeActionController::initialize(
 
   // Attach callbacks for the IncludeTree of the TU. The preprocessor
   // does not exist yet, so we need to indirect this via DependencyCollector.
-  auto DC = std::make_shared<PPCallbacksDependencyCollector>(
+  auto DC = std::make_shared<AttachOnlyDependencyCollector>(
       [&Builder = current()](Preprocessor &PP) {
         return std::make_unique<IncludeTreePPCallbacks>(Builder, PP);
+      },
+      [](ASTReader &R) {
+        return std::make_unique<LookupPCHModulesListener>(R);
       });
   ScanInstance.addDependencyCollector(std::move(DC));
 
-  // Attach callback for modules loaded via PCH.
-  if (!ScanInstance.getPreprocessorOpts().ImplicitPCHInclude.empty()) {
-    auto DC = std::make_shared<ASTReaderListenerDependencyCollector>(
-      [&](ASTReader &R) {
-        return std::make_unique<LookupPCHModulesListener>(R.getPreprocessor());
-      });
-    ScanInstance.addDependencyCollector(std::move(DC));
-  }
-
   // Enable caching in the resulting commands.
   ScanInstance.getFrontendOpts().CacheCompileJob = true;
   CASOpts = ScanInstance.getCASOpts();
@@ -367,9 +366,12 @@ Error IncludeTreeActionController::initializeModuleBuild(
 
   // Attach callbacks for the IncludeTree of the module. The preprocessor
   // does not exist yet, so we need to indirect this via DependencyCollector.
-  auto DC = std::make_shared<PPCallbacksDependencyCollector>(
+  auto DC = std::make_shared<AttachOnlyDependencyCollector>(
       [&Builder = current()](Preprocessor &PP) {
         return std::make_unique<IncludeTreePPCallbacks>(Builder, PP);
+      },
+      [](ASTReader &R) {
+        return std::make_unique<LookupPCHModulesListener>(R);
       });
   ModuleScanInstance.addDependencyCollector(std::move(DC));
 

clang/test/ClangScanDeps/modules-include-tree-pch-with-private.c

@@ -14,8 +14,10 @@
 // RUN:   > %t/deps_pch.json
 
 // RUN: %deps-to-rsp %t/deps_pch.json --module-name Mod > %t/Mod.rsp
+// RUN: %deps-to-rsp %t/deps_pch.json --module-name Indirect1 > %t/Indirect1.rsp
 // RUN: %deps-to-rsp %t/deps_pch.json --tu-index 0 > %t/pch.rsp
 // RUN: %clang @%t/Mod.rsp
+// RUN: %clang @%t/Indirect1.rsp
 // RUN: %clang @%t/pch.rsp
 
 // RUN: clang-scan-deps -compilation-database %t/cdb.json \
@@ -24,10 +26,39 @@
 // RUN:   > %t/deps.json
 
 // RUN: %deps-to-rsp %t/deps.json --module-name Mod_Private > %t/Mod_Private.rsp
+// RUN: %deps-to-rsp %t/deps.json --module-name Indirect2 > %t/Indirect2.rsp
 // RUN: %deps-to-rsp %t/deps.json --tu-index 0 > %t/tu.rsp
 // RUN: %clang @%t/Mod_Private.rsp
+// RUN: %clang @%t/Indirect2.rsp
 // RUN: %clang @%t/tu.rsp
 
+// Extract include-tree casids
+// RUN: cat %t/Indirect2.rsp | sed -E 's|.*"-fcas-include-tree" "(llvmcas://[[:xdigit:]]+)".*|\1|' > %t/Indirect.casid
+// RUN: cat %t/tu.rsp | sed -E 's|.*"-fcas-include-tree" "(llvmcas://[[:xdigit:]]+)".*|\1|' > %t/tu.casid
+
+// RUN: echo "MODULE Indirect2" > %t/result.txt
+// RUN: clang-cas-test -cas %t/cas -print-include-tree @%t/Indirect.casid >> %t/result.txt
+// RUN: echo "TRANSLATION UNIT" >> %t/result.txt
+// RUN: clang-cas-test -cas %t/cas -print-include-tree @%t/tu.casid >> %t/result.txt
+
+// Explicitly check that Mod_Private is imported as a module and not a header.
+// RUN: FileCheck %s -DPREFIX=%/t -input-file %t/result.txt
+
+// CHECK-LABEL: MODULE Indirect2
+// CHECK: <module-includes> llvmcas://
+// CHECK: 1:1 <built-in> llvmcas://
+// CHECK: 2:1 [[PREFIX]]/indirect2.h llvmcas://
+// CHECK:   Submodule: Indirect2
+// CHECK:   2:1 (Module) Indirect1
+// CHECK:   3:1 (Module) Mod_Private
+
+// CHECK-LABEL: TRANSLATION UNIT
+// CHECK: (PCH) llvmcas://
+// CHECK: [[PREFIX]]/tu.m llvmcas://
+// CHECK: 1:1 <built-in> llvmcas://
+// CHECK: 2:1 (Module) Mod_Private
+// CHECK: 3:1 (Module) Indirect2
+
 //--- cdb.json.template
 [{
   "file": "DIR/tu.m",
@@ -54,12 +85,38 @@ void pub(void);
 //--- Mod.framework/PrivateHeaders/Priv.h
 void priv(void);
 
+//--- module.modulemap
+module Indirect1 {
+  header "indirect1.h"
+  export *
+}
+module Indirect2 {
+  header "indirect2.h"
+  export *
+}
+
+//--- indirect1.h
+#import <Mod/Mod.h>
+
+//--- indirect2.h
+#import "indirect1.h"
+#import <Mod/Priv.h>
+
+static inline void indirect(void) {
+  pub();
+  priv();
+}
+
 //--- prefix.h
 #import <Mod/Mod.h>
+#import "indirect1.h"
 
 //--- tu.m
 #import <Mod/Priv.h>
+#import "indirect2.h"
+
 void tu(void) {
   pub();
   priv();
+  indirect();
 }