[IncludeTreeFileSystem] Make dir_begin() work as expected in OverlayFileSystem

clang/include/clang/CAS/IncludeTree.h

@@ -869,6 +869,11 @@ class IncludeTreeRoot : public IncludeTreeBase<IncludeTreeRoot> {
 Expected<IntrusiveRefCntPtr<llvm::vfs::FileSystem>>
 createIncludeTreeFileSystem(IncludeTreeRoot &Root);
 
+/// Create the same IncludeTreeFileSystem but from IncludeTree::FileList.
+Expected<IntrusiveRefCntPtr<llvm::vfs::FileSystem>>
+createIncludeTreeFileSystem(llvm::cas::ObjectStore &CAS,
+                            IncludeTree::FileList &List);
+
 } // namespace cas
 } // namespace clang
 

clang/lib/CAS/IncludeTree.cpp

@@ -9,6 +9,7 @@
 #include "clang/CAS/IncludeTree.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallBitVector.h"
+#include "llvm/CAS/ObjectStore.h"
 #include "llvm/Support/EndianStream.h"
 #include "llvm/Support/Error.h"
 #include <utility>
@@ -1000,7 +1001,9 @@ class IncludeTreeFileSystem : public llvm::vfs::FileSystem {
 
   llvm::vfs::directory_iterator dir_begin(const Twine &Dir,
                                           std::error_code &EC) override {
-    EC = llvm::errc::operation_not_permitted;
+    // Return no_such_file_or_directory so llvm::vfs::OverlayFileSystem can
+    // ignore this layer when iterating directories.
+    EC = llvm::errc::no_such_file_or_directory;
     return llvm::vfs::directory_iterator();
   }
   llvm::ErrorOr<std::string> getCurrentWorkingDirectory() const override {
@@ -1033,12 +1036,18 @@ cas::createIncludeTreeFileSystem(IncludeTreeRoot &Root) {
   if (!FileList)
     return FileList.takeError();
 
+  return createIncludeTreeFileSystem(Root.getCAS(), *FileList);
+}
+
+Expected<IntrusiveRefCntPtr<llvm::vfs::FileSystem>>
+cas::createIncludeTreeFileSystem(llvm::cas::ObjectStore &CAS,
+                                 IncludeTree::FileList &FileList) {
   // Map from FilenameRef to ContentsRef.
   llvm::DenseMap<ObjectRef, ObjectRef> SeenContents;
 
   IntrusiveRefCntPtr<IncludeTreeFileSystem> IncludeTreeFS =
-      new IncludeTreeFileSystem(Root.getCAS());
-  llvm::Error E = FileList->forEachFile(
+      new IncludeTreeFileSystem(CAS);
+  llvm::Error E = FileList.forEachFile(
       [&](IncludeTree::File File,
           IncludeTree::FileList::FileSizeTy Size) -> llvm::Error {
         auto InsertPair = SeenContents.insert(

clang/unittests/CAS/IncludeTreeTest.cpp

@@ -5,8 +5,10 @@
 #include "llvm/CAS/CachingOnDiskFileSystem.h"
 #include "llvm/CAS/ObjectStore.h"
 #include "llvm/Support/Error.h"
+#include "llvm/Support/VirtualFileSystem.h"
 #include "llvm/Testing/Support/Error.h"
 #include "gtest/gtest.h"
+#include <system_error>
 
 using namespace clang;
 using namespace clang::cas;
@@ -294,3 +296,48 @@ TEST(IncludeTree, IncludeTreeFileListDuplicates) {
       llvm::Succeeded());
   EXPECT_EQ(I, Files.size());
 }
+
+TEST(IncludeTree, IncludeTreeFileSystemOverlay) {
+  std::shared_ptr<ObjectStore> DB = llvm::cas::createInMemoryCAS();
+  SmallVector<IncludeTree::FileList::FileEntry> 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({File->getRef(), I});
+  }
+  std::optional<IncludeTree::FileList> FileList;
+  ASSERT_THAT_ERROR(
+      IncludeTree::FileList::create(*DB, Files, {}).moveInto(FileList),
+      llvm::Succeeded());
+  IntrusiveRefCntPtr<llvm::vfs::FileSystem> IncludeTreeFS;
+  ASSERT_THAT_ERROR(
+      createIncludeTreeFileSystem(*DB, *FileList).moveInto(IncludeTreeFS),
+      llvm::Succeeded());
+
+  auto FS = llvm::makeIntrusiveRefCnt<llvm::vfs::InMemoryFileSystem>();
+  FS->setCurrentWorkingDirectory("/dir");
+  FS->addFile("file1", 0,  llvm::MemoryBuffer::getMemBuffer("str"));
+  FS->addFile("file2", 0,  llvm::MemoryBuffer::getMemBuffer("other"));
+
+  llvm::vfs::OverlayFileSystem OverlayFS(std::move(FS));
+  OverlayFS.pushOverlay(IncludeTreeFS);
+
+  std::error_code EC;
+  int NumFile = 0;
+  for (auto I = OverlayFS.dir_begin("/dir", EC);
+       !EC && I != llvm::vfs::directory_iterator(); I.increment(EC)) {
+    ASSERT_FALSE(EC);
+    ++NumFile;
+    std::string Path = "/dir/file" + std::to_string(NumFile);
+    ASSERT_EQ(I->path(), Path);
+  }
+  ASSERT_EQ(NumFile, 2);
+}