[clang][cas] include-tree support for submodule visibility

benlangmuir · benlangmuir · commit bd61c2c82c59 · 2023-03-24T14:03:13.000-07:00
Track the submodule of each modular header and enter the submodule via
PPCachedActions. This fixes declaration and macro header visibility.
diff --git a/clang/include/clang/CAS/IncludeTree.h b/clang/include/clang/CAS/IncludeTree.h
@@ -69,10 +69,31 @@ class IncludeTree : public IncludeTreeBase<IncludeTree> {
   Expected<FileInfo> getBaseFileInfo();
 
   SrcMgr::CharacteristicKind getFileCharacteristic() const {
-    return (SrcMgr::CharacteristicKind)dataSkippingIncludes().front();
+    return (SrcMgr::CharacteristicKind)(getData().front() & ~IsSubmoduleBit);
   }
 
-  size_t getNumIncludes() const { return getNumReferences() - 1; }
+  bool isSubmodule() const { return (getData().front() & IsSubmoduleBit) != 0; }
+
+  std::optional<ObjectRef> getSubmoduleNameRef() const {
+    if (isSubmodule())
+      return getReference(getNumReferences() - 1);
+    return std::nullopt;
+  }
+
+  /// If \c getBaseFile() is a modular header, get its submodule name.
+  Expected<std::optional<StringRef>> getSubmoduleName() {
+    auto Ref = getSubmoduleNameRef();
+    if (!Ref)
+      return std::nullopt;
+    auto Node = getCAS().getProxy(*Ref);
+    if (!Node)
+      return Node.takeError();
+    return Node->getData();
+  }
+
+  size_t getNumIncludes() const {
+    return getNumReferences() - (isSubmodule() ? 2 : 1);
+  };
 
   ObjectRef getIncludeRef(size_t I) const {
     assert(I < getNumIncludes());
@@ -131,7 +152,7 @@ class IncludeTree : public IncludeTreeBase<IncludeTree> {
   static Expected<IncludeTree>
   create(ObjectStore &DB, SrcMgr::CharacteristicKind FileCharacteristic,
          ObjectRef BaseFile, ArrayRef<IncludeInfo> Includes,
-         llvm::SmallBitVector Checks);
+         std::optional<ObjectRef> SubmoduleName, llvm::SmallBitVector Checks);
 
   static Expected<IncludeTree> get(ObjectStore &DB, ObjectRef Ref);
 
@@ -141,6 +162,8 @@ class IncludeTree : public IncludeTreeBase<IncludeTree> {
   friend class IncludeTreeBase<IncludeTree>;
   friend class IncludeTreeRoot;
 
+  static constexpr char IsSubmoduleBit = 0x80;
+
   explicit IncludeTree(ObjectProxy Node) : IncludeTreeBase(std::move(Node)) {
     assert(isValid(*this));
   }
@@ -154,8 +177,10 @@ class IncludeTree : public IncludeTreeBase<IncludeTree> {
 
   Expected<Node> getIncludeNode(ObjectRef Ref, NodeKind Kind);
 
+  StringRef dataSkippingFlags() const { return getData().drop_front(); }
   StringRef dataSkippingIncludes() const {
-    return getData().drop_front(getNumIncludes() * (sizeof(uint32_t) + 1));
+    return dataSkippingFlags().drop_front(getNumIncludes() *
+                                          (sizeof(uint32_t) + 1));
   }
 
   static bool isValid(const ObjectProxy &Node);
diff --git a/clang/include/clang/Lex/PPCachedActions.h b/clang/include/clang/Lex/PPCachedActions.h
@@ -20,6 +20,7 @@
 namespace clang {
 
 class IdentifierInfo;
+class Module;
 class Preprocessor;
 
 /// This interface provides a way to override the actions of the preprocessor as
@@ -33,6 +34,7 @@ class PPCachedActions {
   /// The file that is included by an \c #include directive.
   struct IncludeFile {
     FileID FID;
+    Module *Submodule;
   };
   /// The module that is imported by an \c #include directive or \c @import.
   struct IncludeModule {
diff --git a/clang/lib/CAS/IncludeTree.cpp b/clang/lib/CAS/IncludeTree.cpp
@@ -60,12 +60,15 @@ Expected<IncludeTree::FileInfo> IncludeTree::getBaseFileInfo() {
 llvm::Error IncludeTree::forEachInclude(
     llvm::function_ref<llvm::Error(std::pair<Node, uint32_t>)> Callback) {
   size_t RefI = 0;
+  const size_t IncludeEnd = getNumIncludes();
   return forEachReference([&](ObjectRef Ref) -> llvm::Error {
     if (RefI == 0) {
       ++RefI;
       return llvm::Error::success();
     }
     size_t IncludeI = RefI - 1;
+    if (IncludeI >= IncludeEnd)
+      return llvm::Error::success();
     ++RefI;
     auto Include = getIncludeNode(Ref, getIncludeKind(IncludeI));
     if (!Include)
@@ -74,25 +77,30 @@ llvm::Error IncludeTree::forEachInclude(
   });
 }
 
-Expected<IncludeTree>
-IncludeTree::create(ObjectStore &DB,
-                    SrcMgr::CharacteristicKind FileCharacteristic,
-                    ObjectRef BaseFile, ArrayRef<IncludeInfo> Includes,
-                    llvm::SmallBitVector Checks) {
+Expected<IncludeTree> IncludeTree::create(
+    ObjectStore &DB, SrcMgr::CharacteristicKind FileCharacteristic,
+    ObjectRef BaseFile, ArrayRef<IncludeInfo> Includes,
+    std::optional<ObjectRef> SubmoduleName, llvm::SmallBitVector Checks) {
   // The data buffer is composed of
-  // 1. `uint32_t` offset and `uint8_t` kind for each includes
-  // 2. 1 byte for `CharacteristicKind`
+  // 1. 1 byte for `CharacteristicKind` and IsSubmodule
+  // 2. `uint32_t` offset and `uint8_t` kind for each includes
   // 3. variable number of bitset bytes for `Checks`.
 
   char Kind = FileCharacteristic;
-  assert(Kind == FileCharacteristic && "SrcMgr::CharacteristicKind too big!");
+  assert(Kind == FileCharacteristic && (Kind & IsSubmoduleBit) == 0 &&
+         "SrcMgr::CharacteristicKind too big!");
+  if (SubmoduleName)
+    Kind |= IsSubmoduleBit;
+
   assert(File::isValid(DB, BaseFile));
   SmallVector<ObjectRef, 16> Refs;
-  Refs.reserve(Includes.size() + 1);
+  Refs.reserve(Includes.size() + 2);
   Refs.push_back(BaseFile);
   SmallString<64> Buffer;
   Buffer.reserve(Includes.size() * sizeof(uint32_t) + 1);
 
+  Buffer += Kind;
+
   llvm::raw_svector_ostream BufOS(Buffer);
   llvm::support::endian::Writer Writer(BufOS, llvm::support::little);
 
@@ -107,7 +115,8 @@ IncludeTree::create(ObjectStore &DB,
     Writer.write(static_cast<uint8_t>(Include.Kind));
   }
 
-  Buffer += Kind;
+  if (SubmoduleName)
+    Refs.push_back(*SubmoduleName);
 
   uintptr_t Store;
   ArrayRef<uintptr_t> BitWords = Checks.getData(Store);
@@ -146,15 +155,15 @@ Expected<IncludeTree> IncludeTree::get(ObjectStore &DB, ObjectRef Ref) {
 
 IncludeTree::NodeKind IncludeTree::getIncludeKind(size_t I) const {
   assert(I < getNumIncludes());
-  StringRef Data = getData();
+  StringRef Data = dataSkippingFlags();
   assert(Data.size() >= (I + 1) * (sizeof(uint32_t) + 1));
   uint8_t K = *(Data.data() + I * (sizeof(uint32_t) + 1) + sizeof(uint32_t));
   return NodeKind(K);
 }
 
 uint32_t IncludeTree::getIncludeOffset(size_t I) const {
   assert(I < getNumIncludes());
-  StringRef Data = getData();
+  StringRef Data = dataSkippingFlags();
   assert(Data.size() >= (I + 1) * sizeof(uint32_t));
   uint32_t Offset =
       llvm::support::endian::read<uint32_t, llvm::support::little>(
@@ -164,7 +173,7 @@ uint32_t IncludeTree::getIncludeOffset(size_t I) const {
 
 bool IncludeTree::getCheckResult(size_t I) const {
   // Skip include offsets and CharacteristicKind.
-  StringRef Data = dataSkippingIncludes().drop_front();
+  StringRef Data = dataSkippingIncludes();
   unsigned ByteIndex = I / CHAR_BIT;
   size_t RemainingIndex = I % CHAR_BIT;
   uint8_t Bits = Data[ByteIndex];
@@ -187,9 +196,11 @@ bool IncludeTree::isValid(const ObjectProxy &Node) {
   if (!IncludeTreeBase::isValid(Node))
     return false;
   IncludeTreeBase Base(Node);
-  if (Base.getNumReferences() == 0)
+  if (Base.getNumReferences() == 0 || Base.getData().empty())
     return false;
   unsigned NumIncludes = Base.getNumReferences() - 1;
+  if (Base.getData().front() & IsSubmoduleBit)
+    NumIncludes -= 1;
   return Base.getData().size() >= NumIncludes * (sizeof(uint32_t) + 1) + 1;
 }
 
@@ -321,6 +332,12 @@ llvm::Error IncludeTree::print(llvm::raw_ostream &OS, unsigned Indent) {
   if (llvm::Error E = IncludeBy->print(OS))
     return E;
 
+  auto Submodule = getSubmoduleName();
+  if (!Submodule)
+    return Submodule.takeError();
+  if (*Submodule)
+    OS.indent(Indent) << "Submodule: " << **Submodule << '\n';
+
   llvm::SourceMgr SM;
   auto Blob = IncludeBy->getContents();
   if (!Blob)
diff --git a/clang/lib/Frontend/IncludeTreePPActions.cpp b/clang/lib/Frontend/IncludeTreePPActions.cpp
@@ -106,6 +106,11 @@ class IncludeTreePPActions final : public PPCachedActions {
       return {};
     };
 
+    auto reportErrorTwine = [&](const llvm::Twine &T) -> std::monostate {
+      return reportError(
+          llvm::createStringError(llvm::inconvertibleErrorCode(), T));
+    };
+
     Expected<cas::IncludeTree::Node> Node =
         IncludeInfo.Tree.getIncludeNode(IncludeInfo.CurIncludeIndex++);
     if (!Node)
@@ -142,7 +147,29 @@ class IncludeTreePPActions final : public PPCachedActions {
     PP.markIncluded(*FE);
     IncludeStack.push_back(
         {std::move(EnteredTree), SM.getLocForStartOfFile(FID)});
-    return IncludeFile{FID};
+
+    Module *M = nullptr;
+    auto SubmoduleName = EnteredTree.getSubmoduleName();
+    if (!SubmoduleName)
+      return reportError(SubmoduleName.takeError());
+    if (*SubmoduleName) {
+      SmallVector<StringRef> ModuleComponents;
+      (*SubmoduleName)->split(ModuleComponents, '.');
+      M = PP.getHeaderSearchInfo().lookupModule(
+          ModuleComponents[0], IncludeLoc,
+          /*AllowSearch=*/false, /*AllowExtraModuleMapSearch=*/false);
+      if (!M)
+        return reportErrorTwine(llvm::Twine("failed to find module '") +
+                                ModuleComponents[0] + "'");
+      for (StringRef Sub : ArrayRef(ModuleComponents).drop_front()) {
+        M = M->findOrInferSubmodule(Sub);
+        if (!M)
+          return reportErrorTwine(
+              llvm::Twine("failed to find or infer submodule '") + Sub + "'");
+      }
+    }
+
+    return IncludeFile{FID, M};
   }
 
   void exitedFile(Preprocessor &PP, FileID FID) override {
diff --git a/clang/lib/Lex/PPDirectives.cpp b/clang/lib/Lex/PPDirectives.cpp
@@ -2081,6 +2081,15 @@ void Preprocessor::HandleIncludeDirective(SourceLocation HashLoc,
       }
       EnterSourceFile(File->FID, nullptr, FilenameTok.getLocation(),
                       /*IsFirstIncludeOfFile*/ true);
+
+      if (Module *SM = File->Submodule) {
+        assert(!CurLexerSubmodule &&
+               "should not have marked this as a module yet");
+        CurLexerSubmodule = SM;
+        EnterSubmodule(SM, EndLoc, /*ForPragma=*/false);
+        EnterAnnotationToken(SourceRange(HashLoc, EndLoc),
+                             tok::annot_module_begin, SM);
+      }
       return;
     }
     if (auto *Import = std::get_if<PPCachedActions::IncludeModule>(&Include)) {
diff --git a/clang/lib/Tooling/DependencyScanning/IncludeTreeActionController.cpp b/clang/lib/Tooling/DependencyScanning/IncludeTreeActionController.cpp
@@ -85,11 +85,17 @@ class IncludeTreeBuilder {
 
   void moduleImport(Preprocessor &PP, const Module *M, SourceLocation EndLoc);
 
+  void enteredSubmodule(Preprocessor &PP, Module *M, SourceLocation ImportLoc,
+                        bool ForPragma);
+  void exitedSubmodule(Preprocessor &PP, Module *M, SourceLocation ImportLoc,
+                       bool ForPragma);
+
 private:
   struct FilePPState {
     SrcMgr::CharacteristicKind FileCharacteristic;
     cas::ObjectRef File;
     SmallVector<cas::IncludeTree::IncludeInfo, 6> Includes;
+    std::optional<cas::ObjectRef> SubmoduleName;
     llvm::SmallBitVector HasIncludeChecks;
   };
 
@@ -201,6 +207,15 @@ struct IncludeTreePPCallbacks : public PPCallbacks {
 
     Builder.moduleImport(PP, Imported, EndLoc);
   }
+
+  void EnteredSubmodule(Module *M, SourceLocation ImportLoc,
+                        bool ForPragma) override {
+    Builder.enteredSubmodule(PP, M, ImportLoc, ForPragma);
+  }
+  void LeftSubmodule(Module *M, SourceLocation ImportLoc,
+                     bool ForPragma) override {
+    Builder.exitedSubmodule(PP, M, ImportLoc, ForPragma);
+  }
 };
 } // namespace
 
@@ -354,7 +369,7 @@ void IncludeTreeBuilder::enteredInclude(Preprocessor &PP, FileID FID) {
     return;
   const SrcMgr::FileInfo &FI =
       PP.getSourceManager().getSLocEntry(FID).getFile();
-  IncludeStack.push_back({FI.getFileCharacteristic(), *FileRef, {}, {}});
+  IncludeStack.push_back({FI.getFileCharacteristic(), *FileRef, {}, {}, {}});
 }
 
 void IncludeTreeBuilder::exitedInclude(Preprocessor &PP, FileID IncludedBy,
@@ -395,6 +410,23 @@ void IncludeTreeBuilder::moduleImport(Preprocessor &PP, const Module *M,
        cas::IncludeTree::NodeKind::ModuleImport});
 }
 
+void IncludeTreeBuilder::enteredSubmodule(Preprocessor &PP, Module *M,
+                                          SourceLocation ImportLoc,
+                                          bool ForPragma) {
+  if (ForPragma)
+    return; // Will be parsed as normal.
+  if (hasErrorOccurred())
+    return;
+  assert(!IncludeStack.back().SubmoduleName && "repeated enteredSubmodule");
+  auto Ref = check(DB.storeFromString({}, M->getFullModuleName()));
+  IncludeStack.back().SubmoduleName = Ref;
+}
+void IncludeTreeBuilder::exitedSubmodule(Preprocessor &PP, Module *M,
+                                         SourceLocation ImportLoc,
+                                         bool ForPragma) {
+  // Submodule exit is handled automatically when leaving a modular file.
+}
+
 Expected<cas::IncludeTreeRoot>
 IncludeTreeBuilder::finishIncludeTree(CompilerInstance &ScanInstance,
                                       CompilerInvocation &NewInvocation) {
@@ -615,7 +647,8 @@ IncludeTreeBuilder::addToFileList(FileManager &FM, const FileEntry *FE) {
 Expected<cas::IncludeTree>
 IncludeTreeBuilder::getCASTreeForFileIncludes(FilePPState &&PPState) {
   return cas::IncludeTree::create(DB, PPState.FileCharacteristic, PPState.File,
-                                  PPState.Includes, PPState.HasIncludeChecks);
+                                  PPState.Includes, PPState.SubmoduleName,
+                                  PPState.HasIncludeChecks);
 }
 
 Expected<cas::IncludeTree::File>
diff --git a/clang/test/ClangScanDeps/modules-include-tree-submodules.c b/clang/test/ClangScanDeps/modules-include-tree-submodules.c
