3 files changed, 76 insertions, 0 deletions

diff --git a/lldb/include/lldb/Host/FileSystem.h b/lldb/include/lldb/Host/FileSystem.h
index 37ad7fb3432..115d6d4b9cc 100644
--- a/lldb/include/lldb/Host/FileSystem.h
+++ b/lldb/include/lldb/Host/FileSystem.h
@@ -92,6 +92,28 @@ public:
   void Resolve(FileSpec &file_spec);
   /// @}
 
+  enum EnumerateDirectoryResult {
+    /// Enumerate next entry in the current directory.
+    eEnumerateDirectoryResultNext,
+    /// Recurse into the current entry if it is a directory or symlink, or next
+    /// if not.
+    eEnumerateDirectoryResultEnter,
+    /// Stop directory enumerations at any level.
+    eEnumerateDirectoryResultQuit
+  };
+
+  typedef EnumerateDirectoryResult (*EnumerateDirectoryCallbackType)(
+      void *baton, llvm::sys::fs::file_type file_type, llvm::StringRef);
+
+  typedef std::function<EnumerateDirectoryResult(
+      llvm::sys::fs::file_type file_type, llvm::StringRef)>
+      DirectoryCallback;
+
+  void EnumerateDirectory(llvm::Twine path, bool find_directories,
+                          bool find_files, bool find_other,
+                          EnumerateDirectoryCallbackType callback,
+                          void *callback_baton);
+
   std::error_code GetRealPath(const llvm::Twine &path,
                               llvm::SmallVectorImpl<char> &output) const;
 
diff --git a/lldb/source/Host/common/FileSystem.cpp b/lldb/source/Host/common/FileSystem.cpp
index c985ebb8115..aa181ecdd11 100644
--- a/lldb/source/Host/common/FileSystem.cpp
+++ b/lldb/source/Host/common/FileSystem.cpp
@@ -96,6 +96,36 @@ bool FileSystem::Readable(const FileSpec &file_spec) const {
   return Readable(file_spec.GetPath());
 }
 
+void FileSystem::EnumerateDirectory(Twine path, bool find_directories,
+                                    bool find_files, bool find_other,
+                                    EnumerateDirectoryCallbackType callback,
+                                    void *callback_baton) {
+  std::error_code EC;
+  vfs::recursive_directory_iterator Iter(*m_fs, path, EC);
+  vfs::recursive_directory_iterator End;
+  for (; Iter != End && !EC; Iter.increment(EC)) {
+    const auto &Item = *Iter;
+    ErrorOr<vfs::Status> Status = m_fs->status(Item.path());
+    if (!Status)
+      break;
+    if (!find_files && Status->isRegularFile())
+      continue;
+    if (!find_directories && Status->isDirectory())
+      continue;
+    if (!find_other && Status->isOther())
+      continue;
+
+    auto Result = callback(callback_baton, Status->getType(), Item.path());
+    if (Result == eEnumerateDirectoryResultQuit)
+      return;
+    if (Result == eEnumerateDirectoryResultNext) {
+      // Default behavior is to recurse. Opt out if the callback doesn't want
+      // this behavior.
+      Iter.no_push();
+    }
+  }
+}
+
 std::error_code FileSystem::MakeAbsolute(SmallVectorImpl<char> &path) const {
   return m_fs->makeAbsolute(path);
 }
diff --git a/lldb/unittests/Host/FileSystemTest.cpp b/lldb/unittests/Host/FileSystemTest.cpp
index 2a3c626b087..142ad7bb9ce 100644
--- a/lldb/unittests/Host/FileSystemTest.cpp
+++ b/lldb/unittests/Host/FileSystemTest.cpp
@@ -265,3 +265,27 @@ TEST(FileSystemTest, Resolve) {
     EXPECT_EQ("bogus", file_spec.GetPath());
   }
 }
+
+FileSystem::EnumerateDirectoryResult
+VFSCallback(void *baton, llvm::sys::fs::file_type file_type,
+            llvm::StringRef path) {
+  auto visited = static_cast<std::vector<std::string> *>(baton);
+  visited->push_back(path.str());
+  return FileSystem::eEnumerateDirectoryResultNext;
+}
+
+TEST(FileSystemTest, EnumerateDirectory) {
+  FileSystem fs(GetSimpleDummyFS());
+
+  std::vector<std::string> visited;
+
+  constexpr bool find_directories = true;
+  constexpr bool find_files = true;
+  constexpr bool find_other = true;
+
+  fs.EnumerateDirectory("/", find_directories, find_files, find_other,
+                        VFSCallback, &visited);
+
+  EXPECT_THAT(visited,
+              testing::UnorderedElementsAre("/foo", "/bar", "/baz", "/qux"));
+}