Index: llvm/include/llvm/Support/FileCollector.h
===================================================================
--- llvm/include/llvm/Support/FileCollector.h
+++ 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);
@@ -55,7 +57,10 @@
   bool getRealPath(StringRef SrcPath, SmallVectorImpl<char> &Result);
 
   void addFileToMapping(StringRef VirtualPath, StringRef RealPath) {
-    VFSWriter.addFileMapping(VirtualPath, RealPath);
+    if (sys::fs::is_directory(VirtualPath))
+      VFSWriter.addDirectoryMapping(VirtualPath, RealPath);
+    else
+      VFSWriter.addFileMapping(VirtualPath, RealPath);
   }
 
 protected:
Index: llvm/lib/Support/FileCollector.cpp
===================================================================
--- llvm/lib/Support/FileCollector.cpp
+++ llvm/lib/Support/FileCollector.cpp
@@ -68,6 +68,21 @@
     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 (Iter->type() == sys::fs::file_type::regular_file ||
+        Iter->type() == sys::fs::file_type::directory_file ||
+        Iter->type() == sys::fs::file_type::symlink_file) {
+      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 +228,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);
   }
Index: llvm/unittests/Support/FileCollectorTest.cpp
===================================================================
--- llvm/unittests/Support/FileCollectorTest.cpp
+++ 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");