Index: include/clang/Basic/VirtualFileSystem.h
===================================================================
--- include/clang/Basic/VirtualFileSystem.h
+++ include/clang/Basic/VirtualFileSystem.h
@@ -197,6 +197,25 @@
   /// \note The 'end' iterator is directory_iterator().
   virtual directory_iterator dir_begin(const Twine &Dir,
                                        std::error_code &EC) = 0;
+
+  /// Set the working directory. This will affect all following operations on
+  /// this file system and may propagate down for nested file systems.
+  virtual std::error_code setCurrentWorkingDirectory(const Twine &Path) = 0;
+  /// Get the working directory of this file system.
+  virtual llvm::ErrorOr<std::string> getCurrentWorkingDirectory() const = 0;
+
+  /// Make \a Path an absolute path.
+  ///
+  /// Makes \a Path absolute using the current directory if it is not already.
+  /// An empty \a Path will result in the current directory.
+  ///
+  /// /absolute/path   => /absolute/path
+  /// relative/../path => <current-directory>/relative/../path
+  ///
+  /// \param Path A path that is modified to be an absolute path.
+  /// \returns errc::success if \a path has been made absolute, otherwise a
+  ///          platform-specific error_code.
+  std::error_code make_absolute(SmallVectorImpl<char> &Path) const;
 };
 
 /// \brief Gets an \p vfs::FileSystem for the 'real' file system, as seen by
@@ -228,6 +247,8 @@
   llvm::ErrorOr<std::unique_ptr<File>>
   openFileForRead(const Twine &Path) override;
   directory_iterator dir_begin(const Twine &Dir, std::error_code &EC) override;
+  llvm::ErrorOr<std::string> getCurrentWorkingDirectory() const override;
+  std::error_code setCurrentWorkingDirectory(const Twine &Path) override;
 
   typedef FileSystemList::reverse_iterator iterator;
   
@@ -239,6 +260,35 @@
   iterator overlays_end() { return FSList.rend(); }
 };
 
+namespace detail {
+class InMemoryDirectory;
+} // end namespace detail
+
+/// An in-memory file system.
+class InMemoryFileSystem : public FileSystem {
+  std::unique_ptr<detail::InMemoryDirectory> Root;
+  std::string WorkingDirectory;
+
+public:
+  InMemoryFileSystem();
+  ~InMemoryFileSystem() override;
+  void addFile(const Twine &Path, time_t ModificationTime,
+               std::unique_ptr<llvm::MemoryBuffer> Buffer);
+  StringRef toString() const;
+
+  llvm::ErrorOr<Status> status(const Twine &Path) override;
+  llvm::ErrorOr<std::unique_ptr<File>>
+  openFileForRead(const Twine &Path) override;
+  directory_iterator dir_begin(const Twine &Dir, std::error_code &EC) override;
+  llvm::ErrorOr<std::string> getCurrentWorkingDirectory() const override {
+    return WorkingDirectory;
+  }
+  std::error_code setCurrentWorkingDirectory(const Twine &Path) override {
+    WorkingDirectory = Path.str();
+    return std::error_code();
+  }
+};
+
 /// \brief Get a globally unique ID for a virtual file or directory.
 llvm::sys::fs::UniqueID getNextVirtualUniqueID();
 
Index: lib/Basic/VirtualFileSystem.cpp
===================================================================
--- lib/Basic/VirtualFileSystem.cpp
+++ lib/Basic/VirtualFileSystem.cpp
@@ -77,6 +77,63 @@
   return (*F)->getBuffer(Name, FileSize, RequiresNullTerminator, IsVolatile);
 }
 
+// FIXME: This is copypasted from LLVM's file system implementation.
+std::error_code FileSystem::make_absolute(SmallVectorImpl<char> &path) const {
+  StringRef p(path.data(), path.size());
+
+  bool rootDirectory = llvm::sys::path::has_root_directory(p),
+#ifdef LLVM_ON_WIN32
+       rootName = llvm::sys::path::has_root_name(p);
+#else
+       rootName = true;
+#endif
+
+  // Already absolute.
+  if (rootName && rootDirectory)
+    return std::error_code();
+
+  // All of the following conditions will need the current directory.
+  auto working_dir = getCurrentWorkingDirectory();
+  if (!working_dir)
+    return working_dir.getError();
+
+  SmallString<128> current_dir(working_dir->begin(), working_dir->end());
+
+  // Relative path. Prepend the current directory.
+  if (!rootName && !rootDirectory) {
+    // Append path to the current directory.
+    llvm::sys::path::append(current_dir, p);
+    // Set path to the result.
+    path.swap(current_dir);
+    return std::error_code();
+  }
+
+  if (!rootName && rootDirectory) {
+    StringRef cdrn = llvm::sys::path::root_name(current_dir);
+    SmallString<128> curDirRootName(cdrn.begin(), cdrn.end());
+    llvm::sys::path::append(curDirRootName, p);
+    // Set path to the result.
+    path.swap(curDirRootName);
+    return std::error_code();
+  }
+
+  if (rootName && !rootDirectory) {
+    StringRef pRootName = llvm::sys::path::root_name(p);
+    StringRef bRootDirectory = llvm::sys::path::root_directory(current_dir);
+    StringRef bRelativePath = llvm::sys::path::relative_path(current_dir);
+    StringRef pRelativePath = llvm::sys::path::relative_path(p);
+
+    SmallString<128> res;
+    llvm::sys::path::append(res, pRootName, bRootDirectory, bRelativePath,
+                            pRelativePath);
+    path.swap(res);
+    return std::error_code();
+  }
+
+  llvm_unreachable("All rootName and rootDirectory combinations should have "
+                   "occurred above!");
+}
+
 //===-----------------------------------------------------------------------===/
 // RealFileSystem implementation
 //===-----------------------------------------------------------------------===/
@@ -153,6 +210,9 @@
   ErrorOr<Status> status(const Twine &Path) override;
   ErrorOr<std::unique_ptr<File>> openFileForRead(const Twine &Path) override;
   directory_iterator dir_begin(const Twine &Dir, std::error_code &EC) override;
+
+  llvm::ErrorOr<std::string> getCurrentWorkingDirectory() const override;
+  std::error_code setCurrentWorkingDirectory(const Twine &Path) override;
 };
 } // end anonymous namespace
 
@@ -175,6 +235,28 @@
   return std::move(Result);
 }
 
+llvm::ErrorOr<std::string> RealFileSystem::getCurrentWorkingDirectory() const {
+  SmallString<256> Dir;
+  if (std::error_code EC = llvm::sys::fs::current_path(Dir))
+    return EC;
+  return Dir.str().str();
+}
+
+std::error_code RealFileSystem::setCurrentWorkingDirectory(const Twine &Path) {
+  // FIXME: chdir is thread hostile; on the other hand, creating the same
+  // behavior as chdir is complex: chdir resolves the path once, thus
+  // guaranteeing that all subsequent relative path operations work
+  // on the same path the original chdir resulted in. This makes a
+  // difference for example on network filesystems, where symlinks might be
+  // switched during runtime of the tool. Fixing this depends on having a
+  // file system abstraction that allows openat() style interactions.
+  SmallString<256> Storage;
+  StringRef Dir = Path.toNullTerminatedStringRef(Storage);
+  if (int Err = ::chdir(Dir.data()))
+    return std::error_code(Err, std::generic_category());
+  return std::error_code();
+}
+
 IntrusiveRefCntPtr<FileSystem> vfs::getRealFileSystem() {
   static IntrusiveRefCntPtr<FileSystem> FS = new RealFileSystem();
   return FS;
@@ -229,6 +311,9 @@
 
 void OverlayFileSystem::pushOverlay(IntrusiveRefCntPtr<FileSystem> FS) {
   FSList.push_back(FS);
+  // Synchronize added file systems by duplicating the working directory from
+  // the first one in the list.
+  FS->setCurrentWorkingDirectory(getCurrentWorkingDirectory().get());
 }
 
 ErrorOr<Status> OverlayFileSystem::status(const Twine &Path) {
@@ -252,6 +337,19 @@
   return make_error_code(llvm::errc::no_such_file_or_directory);
 }
 
+llvm::ErrorOr<std::string>
+OverlayFileSystem::getCurrentWorkingDirectory() const {
+  // All file systems are synchronized, just take the first working directory.
+  return FSList.front()->getCurrentWorkingDirectory();
+}
+std::error_code
+OverlayFileSystem::setCurrentWorkingDirectory(const Twine &Path) {
+  for (auto &FS : FSList)
+    if (std::error_code EC = FS->setCurrentWorkingDirectory(Path))
+      return EC;
+  return std::error_code();
+}
+
 clang::vfs::detail::DirIterImpl::~DirIterImpl() { }
 
 namespace {
@@ -320,6 +418,252 @@
       std::make_shared<OverlayFSDirIterImpl>(Dir, *this, EC));
 }
 
+namespace clang {
+namespace vfs {
+namespace detail {
+
+enum InMemoryNodeKind { IME_File, IME_Directory };
+
+/// The in memory file system is a tree of Nodes. Every node can either be a
+/// file or a directory.
+class InMemoryNode {
+  Status Stat;
+  InMemoryNodeKind Kind;
+
+public:
+  InMemoryNode(Status Stat, InMemoryNodeKind Kind)
+      : Stat(std::move(Stat)), Kind(Kind) {}
+  virtual ~InMemoryNode() {}
+  Status &getStatus() { return Stat; }
+  InMemoryNodeKind getKind() const { return Kind; }
+  virtual std::string toString(unsigned Indent) = 0;
+};
+
+namespace {
+class InMemoryFile : public InMemoryNode {
+  std::unique_ptr<llvm::MemoryBuffer> Buffer;
+
+public:
+  InMemoryFile(Status Stat, std::unique_ptr<llvm::MemoryBuffer> Buffer)
+      : InMemoryNode(std::move(Stat), IME_File), Buffer(std::move(Buffer)) {}
+
+  llvm::MemoryBuffer *getBuffer() { return Buffer.get(); }
+  std::string toString(unsigned Indent) override {
+    return (std::string(Indent, ' ') + getStatus().getName() + "\n").str();
+  }
+  static bool classof(const InMemoryNode *N) {
+    return N->getKind() == IME_File;
+  }
+};
+
+/// Adapt a InMemoryFile for VFS' File interface.
+class InMemoryFileAdaptor : public File {
+  InMemoryFile &Parent;
+
+public:
+  explicit InMemoryFileAdaptor(InMemoryFile &Parent) : Parent(Parent) {}
+
+  llvm::ErrorOr<Status> status() override { return Parent.getStatus(); }
+  llvm::ErrorOr<std::unique_ptr<llvm::MemoryBuffer>>
+  getBuffer(const Twine &Name, int64_t FileSize = -1,
+            bool RequiresNullTerminator = true,
+            bool IsVolatile = false) override {
+    llvm::MemoryBuffer *Buf = Parent.getBuffer();
+    return llvm::MemoryBuffer::getMemBuffer(
+        Buf->getBuffer(), Buf->getBufferIdentifier(), RequiresNullTerminator);
+  }
+  std::error_code close() override { return std::error_code(); }
+  void setName(StringRef Name) override { Parent.getStatus().setName(Name); }
+};
+} // end anonymous namespace
+
+class InMemoryDirectory : public InMemoryNode {
+  std::map<std::string, std::unique_ptr<InMemoryNode>> Entries;
+
+public:
+  InMemoryDirectory(Status Stat)
+      : InMemoryNode(std::move(Stat), IME_Directory) {}
+  InMemoryNode *getChild(StringRef Name) {
+    auto I = Entries.find(Name);
+    if (I != Entries.end())
+      return I->second.get();
+    return nullptr;
+  }
+  InMemoryNode *addChild(StringRef Name, std::unique_ptr<InMemoryNode> Child) {
+    return Entries.insert(make_pair(Name, std::move(Child)))
+        .first->second.get();
+  }
+
+  typedef decltype(Entries)::const_iterator const_iterator;
+  const_iterator begin() const { return Entries.begin(); }
+  const_iterator end() const { return Entries.end(); }
+
+  std::string toString(unsigned Indent) override {
+    std::string Result =
+        (std::string(Indent, ' ') + getStatus().getName() + "\n").str();
+    for (const auto &Entry : Entries) {
+      Result += Entry.second->toString(Indent + 2);
+    }
+    return Result;
+  }
+  static bool classof(const InMemoryNode *N) {
+    return N->getKind() == IME_Directory;
+  }
+};
+}
+
+InMemoryFileSystem::InMemoryFileSystem()
+    : Root(new detail::InMemoryDirectory(Status(
+          "", "", getNextVirtualUniqueID(), llvm::sys::TimeValue::MinTime(), 0,
+          0, 0, llvm::sys::fs::file_type::directory_file,
+          llvm::sys::fs::perms::all_all))) {}
+
+InMemoryFileSystem::~InMemoryFileSystem() {}
+
+StringRef InMemoryFileSystem::toString() const {
+  return Root->toString(/*Indent=*/0);
+}
+
+void InMemoryFileSystem::addFile(const Twine &P, time_t ModificationTime,
+                                 std::unique_ptr<llvm::MemoryBuffer> Buffer) {
+  SmallString<128> Path;
+  P.toVector(Path);
+
+  // Fix up relative paths. This just prepends the current working directory.
+  std::error_code EC = make_absolute(Path);
+  assert(!EC);
+  (void)EC;
+
+  detail::InMemoryDirectory *Dir = Root.get();
+  auto I = llvm::sys::path::begin(Path), E = llvm::sys::path::end(Path);
+  while (true) {
+    StringRef Name = *I;
+    detail::InMemoryNode *Node = Dir->getChild(Name);
+    ++I;
+    if (!Node) {
+      if (I == E) {
+        // End of the path, create a new file.
+        Status Stat(Path, "", getNextVirtualUniqueID(),
+                    llvm::sys::TimeValue(ModificationTime), 0, 0,
+                    Buffer->getBufferSize(),
+                    llvm::sys::fs::file_type::regular_file,
+                    llvm::sys::fs::all_all);
+        Dir->addChild(Name, llvm::make_unique<detail::InMemoryFile>(
+                                std::move(Stat), std::move(Buffer)));
+        return;
+      }
+
+      // Create a new directory. Use the path up to here.
+      Status Stat(
+          StringRef(Path.str().begin(), Name.end() - Path.str().begin()), "",
+          getNextVirtualUniqueID(), llvm::sys::TimeValue(ModificationTime), 0,
+          0, Buffer->getBufferSize(), llvm::sys::fs::file_type::directory_file,
+          llvm::sys::fs::all_all);
+      Dir = cast<detail::InMemoryDirectory>(Dir->addChild(
+          Name, llvm::make_unique<detail::InMemoryDirectory>(std::move(Stat))));
+      continue;
+    }
+
+    if (auto *NewDir = dyn_cast<detail::InMemoryDirectory>(Node))
+      Dir = NewDir;
+  }
+}
+
+static ErrorOr<detail::InMemoryNode *>
+lookupInMemoryNode(const InMemoryFileSystem &FS, detail::InMemoryDirectory *Dir,
+                   const Twine &P) {
+  SmallString<128> Path;
+  P.toVector(Path);
+
+  // Fix up relative paths. This just prepends the current working directory.
+  std::error_code EC = FS.make_absolute(Path);
+  assert(!EC);
+  (void)EC;
+
+  auto I = llvm::sys::path::begin(Path), E = llvm::sys::path::end(Path);
+  while (true) {
+    detail::InMemoryNode *Node = Dir->getChild(*I);
+    ++I;
+    if (!Node)
+      return errc::no_such_file_or_directory;
+
+    // Return the file if it's at the end of the path.
+    if (auto File = dyn_cast<detail::InMemoryFile>(Node)) {
+      if (I == E)
+        return File;
+      return errc::no_such_file_or_directory;
+    }
+
+    // Traverse directories.
+    Dir = cast<detail::InMemoryDirectory>(Node);
+    if (I == E)
+      return Dir;
+  }
+}
+
+llvm::ErrorOr<Status> InMemoryFileSystem::status(const Twine &Path) {
+  auto Node = lookupInMemoryNode(*this, Root.get(), Path);
+  if (Node)
+    return (*Node)->getStatus();
+  return Node.getError();
+}
+
+llvm::ErrorOr<std::unique_ptr<File>>
+InMemoryFileSystem::openFileForRead(const Twine &Path) {
+  auto Node = lookupInMemoryNode(*this, Root.get(), Path);
+  if (!Node)
+    return Node.getError();
+
+  // When we have a file provide a heap-allocated wrapper for the memory buffer
+  // to match the ownership semantics for File.
+  if (auto *F = dyn_cast<detail::InMemoryFile>(*Node))
+    return std::unique_ptr<File>(new detail::InMemoryFileAdaptor(*F));
+
+  // FIXME: errc::not_a_file?
+  return make_error_code(llvm::errc::invalid_argument);
+}
+
+namespace {
+/// Adaptor from InMemoryDir::iterator to directory_iterator.
+class InMemoryDirIterator : public clang::vfs::detail::DirIterImpl {
+  detail::InMemoryDirectory::const_iterator I;
+  detail::InMemoryDirectory::const_iterator E;
+
+public:
+  InMemoryDirIterator() {}
+  explicit InMemoryDirIterator(detail::InMemoryDirectory &Dir)
+      : I(Dir.begin()), E(Dir.end()) {
+    if (I != E)
+      CurrentEntry = I->second->getStatus();
+  }
+
+  std::error_code increment() override {
+    ++I;
+    // When we're at the end, make CurrentEntry invalid and DirIterImpl will do
+    // the rest.
+    CurrentEntry = I != E ? I->second->getStatus() : Status();
+    return std::error_code();
+  }
+};
+} // end anonymous namespace
+
+directory_iterator InMemoryFileSystem::dir_begin(const Twine &Dir,
+                                                 std::error_code &EC) {
+  auto Node = lookupInMemoryNode(*this, Root.get(), Dir);
+  if (!Node) {
+    EC = Node.getError();
+    return directory_iterator(std::make_shared<InMemoryDirIterator>());
+  }
+
+  if (auto *DirNode = dyn_cast<detail::InMemoryDirectory>(*Node))
+    return directory_iterator(std::make_shared<InMemoryDirIterator>(*DirNode));
+
+  EC = make_error_code(llvm::errc::not_a_directory);
+  return directory_iterator(std::make_shared<InMemoryDirIterator>());
+}
+}
+}
+
 //===-----------------------------------------------------------------------===/
 // VFSFromYAML implementation
 //===-----------------------------------------------------------------------===/
@@ -453,6 +797,8 @@
   /// \brief The file system to use for external references.
   IntrusiveRefCntPtr<FileSystem> ExternalFS;
 
+  std::string WorkingDirectory;
+
   /// @name Configuration
   /// @{
 
@@ -496,6 +842,14 @@
   ErrorOr<Status> status(const Twine &Path) override;
   ErrorOr<std::unique_ptr<File>> openFileForRead(const Twine &Path) override;
 
+  llvm::ErrorOr<std::string> getCurrentWorkingDirectory() const override {
+    return WorkingDirectory;
+  }
+  std::error_code setCurrentWorkingDirectory(const Twine &Path) override {
+    WorkingDirectory = Path.str();
+    return std::error_code();
+  }
+
   directory_iterator dir_begin(const Twine &Dir, std::error_code &EC) override{
     ErrorOr<Entry *> E = lookupPath(Dir);
     if (!E) {
@@ -866,7 +1220,7 @@
   Path_.toVector(Path);
 
   // Handle relative paths
-  if (std::error_code EC = sys::fs::make_absolute(Path))
+  if (std::error_code EC = make_absolute(Path))
     return EC;
 
   if (Path.empty())
Index: unittests/Basic/VirtualFileSystemTest.cpp
===================================================================
--- unittests/Basic/VirtualFileSystemTest.cpp
+++ unittests/Basic/VirtualFileSystemTest.cpp
@@ -43,6 +43,12 @@
   openFileForRead(const Twine &Path) override {
     llvm_unreachable("unimplemented");
   }
+  llvm::ErrorOr<std::string> getCurrentWorkingDirectory() const override {
+    return std::string();
+  }
+  std::error_code setCurrentWorkingDirectory(const Twine &Path) override {
+    return std::error_code();
+  }
 
   struct DirIterImpl : public clang::vfs::detail::DirIterImpl {
     std::map<std::string, vfs::Status> &FilesAndDirs;
@@ -515,6 +521,73 @@
   }
 }
 
+class InMemoryFileSystemTest : public ::testing::Test {
+protected:
+  clang::vfs::InMemoryFileSystem FS;
+};
+
+TEST_F(InMemoryFileSystemTest, IsEmpty) {
+  auto Stat = FS.status("/a");
+  ASSERT_EQ(errc::no_such_file_or_directory, Stat.getError()) << FS.toString();
+  Stat = FS.status("/");
+  ASSERT_EQ(errc::no_such_file_or_directory, Stat.getError()) << FS.toString();
+}
+
+TEST_F(InMemoryFileSystemTest, WindowsPath) {
+  FS.addFile("c:/windows/system128/foo.cpp", 0, MemoryBuffer::getMemBuffer(""));
+  auto Stat = FS.status("c:");
+  ASSERT_FALSE(Stat.getError()) << Stat.getError() << FS.toString();
+  Stat = FS.status("c:/windows/system128/foo.cpp");
+  ASSERT_FALSE(Stat.getError()) << Stat.getError() << FS.toString();
+  FS.addFile("d:/windows/foo.cpp", 0, MemoryBuffer::getMemBuffer(""));
+  Stat = FS.status("d:/windows/foo.cpp");
+  ASSERT_FALSE(Stat.getError()) << Stat.getError() << FS.toString();
+}
+
+TEST_F(InMemoryFileSystemTest, OverlayFile) {
+  FS.addFile("/a", 0, MemoryBuffer::getMemBuffer("a"));
+  auto Stat = FS.status("/");
+  ASSERT_FALSE(Stat.getError()) << Stat.getError() << FS.toString();
+  Stat = FS.status("/a");
+  ASSERT_FALSE(Stat.getError()) << Stat.getError() << "\n" << FS.toString();
+  ASSERT_EQ("/a", Stat->getName());
+}
+
+TEST_F(InMemoryFileSystemTest, OpenFileForRead) {
+  FS.addFile("/a", 0, MemoryBuffer::getMemBuffer("a"));
+  auto File = FS.openFileForRead("/a");
+  ASSERT_EQ("a", (*(*File)->getBuffer("ignored"))->getBuffer());
+  File = FS.openFileForRead("/a"); // Open again.
+  ASSERT_EQ("a", (*(*File)->getBuffer("ignored"))->getBuffer());
+  File = FS.openFileForRead("/");
+  ASSERT_EQ(errc::invalid_argument, File.getError()) << FS.toString();
+  File = FS.openFileForRead("/b");
+  ASSERT_EQ(errc::no_such_file_or_directory, File.getError()) << FS.toString();
+}
+
+TEST_F(InMemoryFileSystemTest, DirectoryIteration) {
+  FS.addFile("/a", 0, MemoryBuffer::getMemBuffer(""));
+  FS.addFile("/b/c", 0, MemoryBuffer::getMemBuffer(""));
+
+  std::error_code EC;
+  vfs::directory_iterator I = FS.dir_begin("/", EC);
+  ASSERT_FALSE(EC);
+  ASSERT_EQ("/a", I->getName());
+  I.increment(EC);
+  ASSERT_FALSE(EC);
+  ASSERT_EQ("/b", I->getName());
+  I.increment(EC);
+  ASSERT_FALSE(EC);
+  ASSERT_EQ(vfs::directory_iterator(), I);
+
+  I = FS.dir_begin("/b", EC);
+  ASSERT_FALSE(EC);
+  ASSERT_EQ("/b/c", I->getName());
+  I.increment(EC);
+  ASSERT_FALSE(EC);
+  ASSERT_EQ(vfs::directory_iterator(), I);
+}
+
 // NOTE: in the tests below, we use '//root/' as our root directory, since it is
 // a legal *absolute* path on Windows as well as *nix.
 class VFSFromYAMLTest : public ::testing::Test {
Index: unittests/Tooling/RewriterTestContext.h
===================================================================
--- unittests/Tooling/RewriterTestContext.h
+++ unittests/Tooling/RewriterTestContext.h
@@ -34,15 +34,18 @@
 /// methods, which help with writing tests that change files.
 class RewriterTestContext {
  public:
-  RewriterTestContext()
-      : DiagOpts(new DiagnosticOptions()),
-        Diagnostics(IntrusiveRefCntPtr<DiagnosticIDs>(new DiagnosticIDs),
-                    &*DiagOpts),
-        DiagnosticPrinter(llvm::outs(), &*DiagOpts),
-        Files((FileSystemOptions())),
-        Sources(Diagnostics, Files),
-        Rewrite(Sources, Options) {
+   RewriterTestContext()
+       : DiagOpts(new DiagnosticOptions()),
+         Diagnostics(IntrusiveRefCntPtr<DiagnosticIDs>(new DiagnosticIDs),
+                     &*DiagOpts),
+         DiagnosticPrinter(llvm::outs(), &*DiagOpts),
+         InMemoryFileSystem(new vfs::InMemoryFileSystem),
+         OverlayFileSystem(
+             new vfs::OverlayFileSystem(vfs::getRealFileSystem())),
+         Files(FileSystemOptions(), OverlayFileSystem),
+         Sources(Diagnostics, Files), Rewrite(Sources, Options) {
     Diagnostics.setClient(&DiagnosticPrinter, false);
+    OverlayFileSystem->pushOverlay(InMemoryFileSystem);
   }
 
   ~RewriterTestContext() {}
@@ -50,9 +53,9 @@
   FileID createInMemoryFile(StringRef Name, StringRef Content) {
     std::unique_ptr<llvm::MemoryBuffer> Source =
         llvm::MemoryBuffer::getMemBuffer(Content);
-    const FileEntry *Entry =
-      Files.getVirtualFile(Name, Source->getBufferSize(), 0);
-    Sources.overrideFileContents(Entry, std::move(Source));
+    InMemoryFileSystem->addFile(Name, 0, std::move(Source));
+
+    const FileEntry *Entry = Files.getFile(Name);
     assert(Entry != nullptr);
     return Sources.createFileID(Entry, SourceLocation(), SrcMgr::C_User);
   }
@@ -109,6 +112,8 @@
   IntrusiveRefCntPtr<DiagnosticOptions> DiagOpts;
   DiagnosticsEngine Diagnostics;
   TextDiagnosticPrinter DiagnosticPrinter;
+  IntrusiveRefCntPtr<vfs::InMemoryFileSystem> InMemoryFileSystem;
+  IntrusiveRefCntPtr<vfs::OverlayFileSystem> OverlayFileSystem;
   FileManager Files;
   SourceManager Sources;
   LangOptions Options;