Add vfs::recursive_directory_iterator

For now, this is only used by its unit tests.  It is similar to the API
in llvm::sys::fs::recursive_directory_iterator, but without some of the
more complex features like requesting that the iterator not recurse into
the next directory, for example.

git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@211732 91177308-0d34-0410-b5e6-96231b3b80d8

3 files changed, 168 insertions, 8 deletions

diff --git a/include/clang/Basic/VirtualFileSystem.h b/include/clang/Basic/VirtualFileSystem.h
index 3b43d1e75f..882ebdb12b 100644
--- a/include/clang/Basic/VirtualFileSystem.h
+++ b/include/clang/Basic/VirtualFileSystem.h
@@ -149,6 +149,37 @@ public:
   }
 };
 
+class FileSystem;
+
+/// \brief An input iterator over the recursive contents of a virtual path,
+/// similar to llvm::sys::fs::recursive_directory_iterator.
+class recursive_directory_iterator {
+  typedef std::stack<directory_iterator, std::vector<directory_iterator>>
+      IterState;
+
+  FileSystem *FS;
+  std::shared_ptr<IterState> State; // Input iterator semantics on copy.
+
+public:
+  recursive_directory_iterator(FileSystem &FS, const Twine &Path,
+                               std::error_code &EC);
+  /// \brief Construct an 'end' iterator.
+  recursive_directory_iterator() { }
+
+  /// \brief Equivalent to operator++, with an error code.
+  recursive_directory_iterator &increment(std::error_code &EC);
+
+  const Status &operator*() const { return *State->top(); }
+  const Status *operator->() const { return &*State->top(); }
+
+  bool operator==(const recursive_directory_iterator &Other) const {
+    return State == Other.State; // identity
+  }
+  bool operator!=(const recursive_directory_iterator &RHS) const {
+    return !(*this == RHS);
+  }
+};
+
 /// \brief The virtual file system interface.
 class FileSystem : public llvm::ThreadSafeRefCountedBase<FileSystem> {
 public:
@@ -172,8 +203,6 @@ public:
   /// \note The 'end' iterator is directory_iterator()
   virtual directory_iterator dir_begin(const Twine &Dir,
                                        std::error_code &EC) = 0;
-
-  // TODO: recursive directory iterators
 };
 
 /// \brief Gets an \p vfs::FileSystem for the 'real' file system, as seen by
diff --git a/lib/Basic/VirtualFileSystem.cpp b/lib/Basic/VirtualFileSystem.cpp
index 253680e855..3d5e1ad24a 100644
--- a/lib/Basic/VirtualFileSystem.cpp
+++ b/lib/Basic/VirtualFileSystem.cpp
@@ -1168,3 +1168,38 @@ std::error_code VFSFromYamlDirIterImpl::increment() {
   }
   return std::error_code();
 }
+
+vfs::recursive_directory_iterator::recursive_directory_iterator(FileSystem &FS_,
+                                                           const Twine &Path,
+                                                           std::error_code &EC)
+    : FS(&FS_) {
+  directory_iterator I = FS->dir_begin(Path, EC);
+  if (!EC && I != directory_iterator()) {
+    State = std::make_shared<IterState>();
+    State->push(I);
+  }
+}
+
+vfs::recursive_directory_iterator &
+recursive_directory_iterator::increment(std::error_code &EC) {
+  assert(FS && State && !State->empty() && "incrementing past end");
+  assert(State->top()->isStatusKnown() && "non-canonical end iterator");
+  vfs::directory_iterator End;
+  if (State->top()->isDirectory()) {
+    vfs::directory_iterator I = FS->dir_begin(State->top()->getName(), EC);
+    if (EC)
+      return *this;
+    if (I != End) {
+      State->push(I);
+      return *this;
+    }
+  }
+
+  while (!State->empty() && State->top().increment(EC) == End)
+    State->pop();
+
+  if (State->empty())
+    State.reset(); // end iterator
+
+  return *this;
+}
\ No newline at end of file
diff --git a/unittests/Basic/VirtualFileSystemTest.cpp b/unittests/Basic/VirtualFileSystemTest.cpp
index c416aa6c49..9289b22f8e 100644
--- a/unittests/Basic/VirtualFileSystemTest.cpp
+++ b/unittests/Basic/VirtualFileSystemTest.cpp
@@ -54,12 +54,20 @@ public:
     std::map<std::string, vfs::Status> &FilesAndDirs;
     std::map<std::string, vfs::Status>::iterator I;
     std::string Path;
+    bool isInPath(StringRef S) {
+      if (Path.size() < S.size() && S.find(Path) == 0) {
+        auto LastSep = S.find_last_of('/');
+        if (LastSep == Path.size() || LastSep == Path.size()-1)
+          return true;
+      }
+      return false;
+    }
     DirIterImpl(std::map<std::string, vfs::Status> &FilesAndDirs,
                 const Twine &_Path)
         : FilesAndDirs(FilesAndDirs), I(FilesAndDirs.begin()),
           Path(_Path.str()) {
       for ( ; I != FilesAndDirs.end(); ++I) {
-        if (Path.size() < I->first.size() && I->first.find(Path) == 0 && I->first.find_last_of('/') <= Path.size()) {
+        if (isInPath(I->first)) {
           CurrentEntry = I->second;
           break;
         }
@@ -68,7 +76,7 @@ public:
     std::error_code increment() override {
       ++I;
       for ( ; I != FilesAndDirs.end(); ++I) {
-        if (Path.size() < I->first.size() && I->first.find(Path) == 0 && I->first.find_last_of('/') <= Path.size()) {
+        if (isInPath(I->first)) {
           CurrentEntry = I->second;
           break;
         }
@@ -306,6 +314,46 @@ TEST(VirtualFileSystemTest, BasicRealFSIteration) {
   EXPECT_EQ(vfs::directory_iterator(), I);
 }
 
+TEST(VirtualFileSystemTest, BasicRealFSRecursiveIteration) {
+  ScopedDir TestDirectory("virtual-file-system-test", /*Unique*/true);
+  IntrusiveRefCntPtr<vfs::FileSystem> FS = vfs::getRealFileSystem();
+
+  std::error_code EC;
+  auto I = vfs::recursive_directory_iterator(*FS, Twine(TestDirectory), EC);
+  ASSERT_FALSE(EC);
+  EXPECT_EQ(vfs::recursive_directory_iterator(), I); // empty directory is empty
+
+  ScopedDir _a(TestDirectory+"/a");
+  ScopedDir _ab(TestDirectory+"/a/b");
+  ScopedDir _c(TestDirectory+"/c");
+  ScopedDir _cd(TestDirectory+"/c/d");
+
+  I = vfs::recursive_directory_iterator(*FS, Twine(TestDirectory), EC);
+  ASSERT_FALSE(EC);
+  ASSERT_NE(vfs::recursive_directory_iterator(), I);
+
+
+  std::vector<std::string> Contents;
+  for (auto E = vfs::recursive_directory_iterator(); !EC && I != E;
+       I.increment(EC)) {
+    Contents.push_back(I->getName());
+  }
+
+  // Check contents, which may be in any order
+  EXPECT_EQ(4U, Contents.size());
+  int Counts[4] = { 0, 0, 0, 0 };
+  for (const std::string &Name : Contents) {
+    ASSERT_FALSE(Name.empty());
+    int Index = Name[Name.size()-1] - 'a';
+    ASSERT_TRUE(Index >= 0 && Index < 4);
+    Counts[Index]++;
+  }
+  EXPECT_EQ(1, Counts[0]); // a
+  EXPECT_EQ(1, Counts[1]); // b
+  EXPECT_EQ(1, Counts[2]); // c
+  EXPECT_EQ(1, Counts[3]); // d
+}
+
 template <typename T, size_t N>
 std::vector<StringRef> makeStringRefVector(const T (&Arr)[N]) {
   std::vector<StringRef> Vec;
@@ -314,17 +362,17 @@ std::vector<StringRef> makeStringRefVector(const T (&Arr)[N]) {
   return Vec;
 }
 
-static void checkContents(vfs::directory_iterator I,
-                          ArrayRef<StringRef> Expected) {
+template <typename DirIter>
+static void checkContents(DirIter I, ArrayRef<StringRef> Expected) {
   std::error_code EC;
   auto ExpectedIter = Expected.begin(), ExpectedEnd = Expected.end();
-  for (vfs::directory_iterator E;
+  for (DirIter E;
        !EC && I != E && ExpectedIter != ExpectedEnd;
        I.increment(EC), ++ExpectedIter)
     EXPECT_EQ(*ExpectedIter, I->getName());
 
   EXPECT_EQ(ExpectedEnd, ExpectedIter);
-  EXPECT_EQ(vfs::directory_iterator(), I);
+  EXPECT_EQ(DirIter(), I);
 }
 
 TEST(VirtualFileSystemTest, OverlayIteration) {
@@ -357,6 +405,54 @@ TEST(VirtualFileSystemTest, OverlayIteration) {
   }
 }
 
+TEST(VirtualFileSystemTest, OverlayRecursiveIteration) {
+  IntrusiveRefCntPtr<DummyFileSystem> Lower(new DummyFileSystem());
+  IntrusiveRefCntPtr<DummyFileSystem> Middle(new DummyFileSystem());
+  IntrusiveRefCntPtr<DummyFileSystem> Upper(new DummyFileSystem());
+  IntrusiveRefCntPtr<vfs::OverlayFileSystem> O(
+      new vfs::OverlayFileSystem(Lower));
+  O->pushOverlay(Middle);
+  O->pushOverlay(Upper);
+
+  std::error_code EC;
+  checkContents(vfs::recursive_directory_iterator(*O, "/", EC),
+                ArrayRef<StringRef>());
+
+  Lower->addRegularFile("/file1");
+  checkContents(vfs::recursive_directory_iterator(*O, "/", EC),
+                ArrayRef<StringRef>("/file1"));
+
+  Upper->addDirectory("/dir");
+  Upper->addRegularFile("/dir/file2");
+  {
+    const char *Contents[] = {"/dir", "/dir/file2", "/file1"};
+    checkContents(vfs::recursive_directory_iterator(*O, "/", EC),
+                  makeStringRefVector(Contents));
+  }
+
+  Lower->addDirectory("/dir1");
+  Lower->addRegularFile("/dir1/foo");
+  Lower->addDirectory("/dir1/a");
+  Lower->addRegularFile("/dir1/a/b");
+  Middle->addDirectory("/a");
+  Middle->addDirectory("/a/b");
+  Middle->addDirectory("/a/b/c");
+  Middle->addRegularFile("/a/b/c/d");
+  Middle->addRegularFile("/hiddenByUp");
+  Upper->addDirectory("/dir2");
+  Upper->addRegularFile("/dir2/foo");
+  Upper->addRegularFile("/hiddenByUp");
+  checkContents(vfs::recursive_directory_iterator(*O, "/dir2", EC),
+                ArrayRef<StringRef>("/dir2/foo"));
+  {
+    const char *Contents[] = { "/dir", "/dir/file2", "/dir2", "/dir2/foo",
+        "/hiddenByUp", "/a", "/a/b", "/a/b/c", "/a/b/c/d", "/dir1", "/dir1/a",
+        "/dir1/a/b", "/dir1/foo", "/file1" };
+    checkContents(vfs::recursive_directory_iterator(*O, "/", EC),
+                  makeStringRefVector(Contents));
+  }
+}
+
 TEST(VirtualFileSystemTest, ThreeLevelIteration) {
   IntrusiveRefCntPtr<DummyFileSystem> Lower(new DummyFileSystem());
   IntrusiveRefCntPtr<DummyFileSystem> Middle(new DummyFileSystem());