diff --git a/llvm/include/llvm/Support/FileCollector.h b/llvm/include/llvm/Support/FileCollector.h --- a/llvm/include/llvm/Support/FileCollector.h +++ b/llvm/include/llvm/Support/FileCollector.h @@ -24,8 +24,10 @@ class FileCollector { public: FileCollector(std::string Root, std::string OverlayRoot); + virtual ~FileCollector() = default; - void addFile(const Twine &file); + virtual void addFile(const Twine &file); + virtual void addDirectory(const Twine &dir); /// Write the yaml mapping (for the VFS) to the given file. std::error_code writeMapping(StringRef mapping_file); diff --git a/llvm/lib/Support/FileCollector.cpp b/llvm/lib/Support/FileCollector.cpp --- a/llvm/lib/Support/FileCollector.cpp +++ b/llvm/lib/Support/FileCollector.cpp @@ -68,6 +68,18 @@ addFileImpl(FileStr); } +void FileCollector::addDirectory(const Twine &dir) { + assert(sys::fs::is_directory(dir)); + std::error_code EC; + sys::fs::recursive_directory_iterator Iter(dir, EC); + sys::fs::recursive_directory_iterator End; + addFile(dir); // Add the directory in case it's empty. + for (; Iter != End && !EC; Iter.increment(EC)) { + if (sys::fs::is_regular_file(Iter->path())) + addFile(Iter->path()); + } +} + void FileCollector::addFileImpl(StringRef SrcPath) { // We need an absolute src path to append to the root. SmallString<256> AbsoluteSrc = SrcPath; @@ -213,20 +225,9 @@ llvm::vfs::directory_iterator dir_begin(const llvm::Twine &Dir, std::error_code &EC) override { - auto It = FS->dir_begin(Dir, EC); - if (EC) - return It; // Collect everything that's listed in case the user needs it. - Collector->addFile(Dir); - for (; !EC && It != llvm::vfs::directory_iterator(); It.increment(EC)) { - if (It->type() == sys::fs::file_type::regular_file || - It->type() == sys::fs::file_type::directory_file || - It->type() == sys::fs::file_type::symlink_file) { - Collector->addFile(It->path()); - } - } - if (EC) - return It; + Collector->addDirectory(Dir); + // Return a new iterator. return FS->dir_begin(Dir, EC); } diff --git a/llvm/unittests/Support/FileCollectorTest.cpp b/llvm/unittests/Support/FileCollectorTest.cpp --- a/llvm/unittests/Support/FileCollectorTest.cpp +++ b/llvm/unittests/Support/FileCollectorTest.cpp @@ -120,6 +120,29 @@ EXPECT_FALSE(FileCollector.hasSeen("/path/to/d")); } +TEST(FileCollectorTest, addDirectory) { + ScopedDir file_root("file_root", true); + ScopedFile a(file_root + "/aaa"); + ScopedFile b(file_root + "/bbb"); + ScopedFile c(file_root + "/ccc"); + + std::string root_fs = std::string(file_root.Path.str()); + TestingFileCollector FileCollector(root_fs, root_fs); + + FileCollector.addDirectory(file_root.Path); + + // Make sure the root is correct. + EXPECT_EQ(FileCollector.Root, root_fs); + + // Make sure we've seen all the added files. + EXPECT_TRUE(FileCollector.hasSeen(a.Path)); + EXPECT_TRUE(FileCollector.hasSeen(b.Path)); + EXPECT_TRUE(FileCollector.hasSeen(c.Path)); + + // Make sure we've only seen the added files. + EXPECT_FALSE(FileCollector.hasSeen("/file_root/ddd")); +} + TEST(FileCollectorTest, copyFiles) { ScopedDir file_root("file_root", true); ScopedFile a(file_root + "/aaa");