[VFS] Add support for "no_push" to VFS recursive iterators.

[JDevlieghere]

Summary

The "regular" file system has a useful feature that makes it possible to stop recursing when using the recursive directory iterators. This functionality was missing for the VFS recursive iterator and this patch adds that.

[JDevlieghere]

ping

[vsapsai]

I am gravitating towards having a separate unit test(s) for no_push functionality. Maybe I'm wrong but I expect it to be smaller and easier to understand, though there'll be some boilerplate.

Implementation looks reasonable but I'm not sure about a few edge cases:

• call no_push before starting iteration;
• call no_push after we finished iteration (here I mean when we are pointing to the last directory, before the iterator becomes end iterator);
• call no_push for a directory that has no children and we wouldn't descend into anyway.

I think in these cases VFS iterator should behave as llvm::sys::fs::recursive_directory_iterator unless it is doing something strange. Not sure all these cases deserve a test, depends on the fact if they execute different paths in implementation.

llvm/unittests/Support/VirtualFileSystemTest.cpp

479–481

Is Counts[5] supposed to be // f?

[JDevlieghere]

Thanks for the review, Volodymyr! I've extended the test which I believe covers all the cases you mentioned. The implementation is almost identical to how no_push is handled in the non-VFS filesystem implementation, so the behavior should indeed be identical.

[vsapsai]

Looks good to me.

llvm/include/llvm/Support/VirtualFileSystem.h

public:
   }
   bool operator!=(const directory_iterator &RHS) const {
     return !(*this == RHS);
   }
 };
 
 class FileSystem;
 
+namespace detail {
+
+/// Keeps state for the recursive_directory_iterator.
+struct RecDirIterState {
+  std::stack<directory_iterator, std::vector<directory_iterator>> Stack;
+  bool HasNoPushRequest = false;
+};
+
+} // end namespace detail
+
 /// An input iterator over the recursive contents of a virtual path,
 /// similar to llvm::sys::fs::recursive_directory_iterator.
 class recursive_directory_iterator {
-  using IterState =
-      std::stack<directory_iterator, std::vector<directory_iterator>>;
-
   FileSystem *FS;
-  std::shared_ptr<IterState> State; // Input iterator semantics on copy.
+  std::shared_ptr<detail::RecDirIterState>
+      State; // Input iterator semantics on copy.
 
 public:
   recursive_directory_iterator(FileSystem &FS, const Twine &Path,
                                std::error_code &EC);
 
   /// Construct an 'end' iterator.
   recursive_directory_iterator() = default;
 
   /// Equivalent to operator++, with an error code.
   recursive_directory_iterator &increment(std::error_code &EC);
 
-  const directory_entry &operator*() const { return *State->top(); }
-  const directory_entry *operator->() const { return &*State->top(); }
+  const directory_entry &operator*() const { return *State->Stack.top(); }
+  const directory_entry *operator->() const { return &*State->Stack.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);
   }
 
   /// Gets the current level. Starting path is at level 0.
   int level() const {
-    assert(!State->empty() && "Cannot get level without any iteration state");
-    return State->size() - 1;
+    assert(!State->Stack.empty() &&
+           "Cannot get level without any iteration state");
+    return State->Stack.size() - 1;
   }
+
+  void no_push() { State->HasNoPushRequest = true; }
 };
 
 /// The virtual file system interface.
 class FileSystem : public llvm::ThreadSafeRefCountedBase<FileSystem> {
 public:
   virtual ~FileSystem();
 
   /// Get the status of the entry at \p Path, if one exists.

llvm/lib/Support/VirtualFileSystem.cpp

std::error_code VFSFromYamlDirIterImpl::incrementImpl(bool IsFirstTime) {
   llvm_unreachable("returned above");
 }
 
 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 (I != directory_iterator()) {
-    State = std::make_shared<IterState>();
-    State->push(I);
+    State = std::make_shared<detail::RecDirIterState>();
+    State->Stack.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()->path().empty() && "non-canonical end iterator");
+  assert(FS && State && !State->Stack.empty() && "incrementing past end");
+  assert(!State->Stack.top()->path().empty() && "non-canonical end iterator");
   vfs::directory_iterator End;
-  if (State->top()->type() == sys::fs::file_type::directory_file) {
-    vfs::directory_iterator I = FS->dir_begin(State->top()->path(), EC);
+
+  if (State->HasNoPushRequest)
+    State->HasNoPushRequest = false;
+  else {
+    if (State->Stack.top()->type() == sys::fs::file_type::directory_file) {
+      vfs::directory_iterator I = FS->dir_begin(State->Stack.top()->path(), EC);
       if (I != End) {
-      State->push(I);
+        State->Stack.push(I);
         return *this;
       }
     }
+  }
 
-  while (!State->empty() && State->top().increment(EC) == End)
-    State->pop();
+  while (!State->Stack.empty() && State->Stack.top().increment(EC) == End)
+    State->Stack.pop();
 
-  if (State->empty())
+  if (State->Stack.empty())
     State.reset(); // end iterator
 
   return *this;
 }

llvm/unittests/Support/VirtualFileSystemTest.cpp

for (const std::string &Name : Contents) {
     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
 }
 
+TEST(VirtualFileSystemTest, BasicRealFSRecursiveIterationNoPush) {
+  ScopedDir TestDirectory("virtual-file-system-test", /*Unique*/ true);
+
+  ScopedDir _a(TestDirectory + "/a");
+  ScopedDir _ab(TestDirectory + "/a/b");
+  ScopedDir _c(TestDirectory + "/c");
+  ScopedDir _cd(TestDirectory + "/c/d");
+  ScopedDir _e(TestDirectory + "/e");
+  ScopedDir _ef(TestDirectory + "/e/f");
+  ScopedDir _g(TestDirectory + "/g");
+
+  IntrusiveRefCntPtr<vfs::FileSystem> FS = vfs::getRealFileSystem();
+
+  // Test that calling no_push on entries without subdirectories has no effect.
+  {
+    std::error_code EC;
+    auto I = vfs::recursive_directory_iterator(*FS, Twine(TestDirectory), EC);
+    ASSERT_FALSE(EC);
+
+    std::vector<std::string> Contents;
+    for (auto E = vfs::recursive_directory_iterator(); !EC && I != E;
+         I.increment(EC)) {
+      Contents.push_back(I->path());
+      char last = I->path().back();
+      switch (last) {
+      case 'b':
+      case 'd':
+      case 'f':
+      case 'g':
+        I.no_push();
+        break;
+      default:
+        break;
+      }
+    }
+    EXPECT_EQ(7U, Contents.size());
+  }
+
+  // Test that calling no_push skips subdirectories.
+  {
+    std::error_code EC;
+    auto I = vfs::recursive_directory_iterator(*FS, Twine(TestDirectory), EC);
+    ASSERT_FALSE(EC);
+
+    std::vector<std::string> Contents;
+    for (auto E = vfs::recursive_directory_iterator(); !EC && I != E;
+         I.increment(EC)) {
+      Contents.push_back(I->path());
+      char last = I->path().back();
+      switch (last) {
+      case 'a':
+      case 'c':
+      case 'e':
+        I.no_push();
+        break;
+      default:
+        break;
+      }
+    }
+
+    // Check contents, which may be in any order
+    EXPECT_EQ(4U, Contents.size());
+    int Counts[7] = {0, 0, 0, 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 < 7);
+      Counts[Index]++;
+    }
+    EXPECT_EQ(1, Counts[0]); // a
+    EXPECT_EQ(0, Counts[1]); // b
+    EXPECT_EQ(1, Counts[2]); // c
+    EXPECT_EQ(0, Counts[3]); // d
+    EXPECT_EQ(1, Counts[4]); // e
+    EXPECT_EQ(0, Counts[5]); // f
+    EXPECT_EQ(1, Counts[6]); // g
+  }
+}
+
 #ifdef LLVM_ON_UNIX
 TEST(VirtualFileSystemTest, BrokenSymlinkRealFSRecursiveIteration) {
   ScopedDir TestDirectory("virtual-file-system-test", /*Unique*/ true);
   IntrusiveRefCntPtr<vfs::FileSystem> FS = vfs::getRealFileSystem();
 
   ScopedLink _a("no_such_file", TestDirectory + "/a");
   ScopedDir _b(TestDirectory + "/b");
   ScopedLink _ba("no_such_file", TestDirectory + "/b/a");