[clang][cas] Skip duplicates in forEachFile

If we see the same File or FileList more than once, skip visiting it
again. This improves performance if the same module is reached through
multiple dependency chains, or the same file is in multiple
modules/PCH/TU file lists explicitly.

(cherry picked from commit 591fa52)
(cherry picked from commit 653e01b)

Author: benlangmuir

clang/include/clang/CAS/IncludeTree.h

@@ -309,6 +309,10 @@ class IncludeTree::FileList : public IncludeTreeBase<FileList> {
     return File(std::move(*Node));
   }
 
+  llvm::Error
+  forEachFileImpl(llvm::DenseSet<ObjectRef> &Seen,
+                  llvm::function_ref<llvm::Error(File, FileSizeTy)> Callback);
+
   static bool isValid(const ObjectProxy &Node);
   static bool isValid(ObjectStore &CAS, ObjectRef Ref) {
     auto Node = CAS.getProxy(Ref);

clang/lib/CAS/IncludeTree.cpp

@@ -237,27 +237,38 @@ IncludeTree::FileList::getFileSize(size_t I) const {
       Data.data() + I * sizeof(FileSizeTy));
 }
 
-llvm::Error IncludeTree::FileList::forEachFile(
+llvm::Error IncludeTree::FileList::forEachFileImpl(
+    llvm::DenseSet<ObjectRef> &Seen,
     llvm::function_ref<llvm::Error(File, FileSizeTy)> Callback) {
-  size_t I = 0;
+  size_t Next = 0;
   size_t FileCount = getNumFilesCurrentList();
   return forEachReference([&](ObjectRef Ref) -> llvm::Error {
-    if (I < FileCount) {
+    size_t Index = Next++;
+    if (!Seen.insert(Ref).second)
+      return llvm::Error::success();
+
+    if (Index < FileCount) {
       auto Include = getFile(Ref);
       if (!Include)
         return Include.takeError();
-      return Callback(std::move(*Include), getFileSize(I++));
+      return Callback(std::move(*Include), getFileSize(Index));
     }
+
     // Otherwise, it's a chained FileList.
-    ++I;
     auto Proxy = getCAS().getProxy(Ref);
     if (!Proxy)
       return Proxy.takeError();
     FileList FL(std::move(*Proxy));
-    return FL.forEachFile(Callback);
+    return FL.forEachFileImpl(Seen, Callback);
   });
 }
 
+llvm::Error IncludeTree::FileList::forEachFile(
+    llvm::function_ref<llvm::Error(File, FileSizeTy)> Callback) {
+  llvm::DenseSet<ObjectRef> Seen;
+  return forEachFileImpl(Seen, Callback);
+}
+
 Expected<IncludeTree::FileList>
 IncludeTree::FileList::create(ObjectStore &DB, ArrayRef<FileEntry> Files,
                               ArrayRef<ObjectRef> FileLists) {

clang/unittests/CAS/IncludeTreeTest.cpp

@@ -197,6 +197,8 @@ TEST(IncludeTree, IncludeTreeFileList) {
   size_t I = 0;
   ASSERT_THAT_ERROR(
       L->forEachFile([&](IncludeTree::File F, auto Size) -> llvm::Error {
+        if (I >= Files.size())
+          return llvm::Error::success(); // diagnosed later.
         EXPECT_EQ(F.getFilenameRef(), Files[I].getFilenameRef())
             << "filename mismatch at " << I;
         EXPECT_EQ(F.getContentsRef(), Files[I].getContentsRef())
@@ -206,4 +208,55 @@ TEST(IncludeTree, IncludeTreeFileList) {
         return llvm::Error::success();
       }),
       llvm::Succeeded());
+  EXPECT_EQ(I, Files.size());
+}
+
+TEST(IncludeTree, IncludeTreeFileListDuplicates) {
+  std::shared_ptr<ObjectStore> DB = llvm::cas::createInMemoryCAS();
+  SmallVector<IncludeTree::File> Files;
+  for (unsigned I = 0; I < 10; ++I) {
+    std::optional<IncludeTree::File> File;
+    std::string Path = "/file" + std::to_string(I);
+    static constexpr StringRef Bytes = "123456789";
+    std::optional<ObjectRef> Content;
+    ASSERT_THAT_ERROR(
+        DB->storeFromString({}, Bytes.substr(0, I)).moveInto(Content),
+        llvm::Succeeded());
+    ASSERT_THAT_ERROR(
+        IncludeTree::File::create(*DB, Path, *Content).moveInto(File),
+        llvm::Succeeded());
+    Files.push_back(std::move(*File));
+  }
+
+  auto MakeFileList = [&](unsigned Begin, unsigned End,
+                          ArrayRef<ObjectRef> Lists) {
+    SmallVector<IncludeTree::FileList::FileEntry> Entries;
+    for (; Begin != End; ++Begin)
+      Entries.push_back({Files[Begin].getRef(), Begin});
+    return IncludeTree::FileList::create(*DB, Entries, Lists);
+  };
+
+  std::optional<IncludeTree::FileList> L89, L;
+  ASSERT_THAT_ERROR(MakeFileList(8, 10, {}).moveInto(L89), llvm::Succeeded());
+  EXPECT_EQ(L89->getNumReferences(), 2u);
+  ASSERT_THAT_ERROR(
+      MakeFileList(0, 9, {L89->getRef(), L89->getRef()}).moveInto(L),
+      llvm::Succeeded());
+  EXPECT_EQ(L->getNumReferences(), 11u); // 0, 1, ..., 8, {8, 9}, {8, 9}
+
+  size_t I = 0;
+  ASSERT_THAT_ERROR(
+      L->forEachFile([&](IncludeTree::File F, auto Size) -> llvm::Error {
+        if (I >= Files.size())
+          return llvm::Error::success(); // diagnosed later.
+        EXPECT_EQ(F.getFilenameRef(), Files[I].getFilenameRef())
+            << "filename mismatch at " << I;
+        EXPECT_EQ(F.getContentsRef(), Files[I].getContentsRef())
+            << "contents mismatch at " << I;
+        EXPECT_EQ(Size, I) << "size mismatch at " << I;
+        I += 1;
+        return llvm::Error::success();
+      }),
+      llvm::Succeeded());
+  EXPECT_EQ(I, Files.size());
 }