3 files changed, 42 insertions, 20 deletions

diff --git a/include/clang/Serialization/ASTReader.h b/include/clang/Serialization/ASTReader.h
index e167985e30..5d663105cc 100644
--- a/include/clang/Serialization/ASTReader.h
+++ b/include/clang/Serialization/ASTReader.h
@@ -436,6 +436,9 @@ class ModuleManager {
   /// \brief FileManager that handles translating between filenames and
   /// FileEntry *.
   FileManager FileMgr;
+  
+  /// \brief A lookup of in-memory (virtual file) buffers
+  llvm::DenseMap<const FileEntry *, llvm::MemoryBuffer *> InMemoryBuffers;
 
 public:
   typedef SmallVector<Module*, 2>::iterator ModuleIterator;
@@ -481,12 +484,18 @@ public:
 
   /// \brief Returns the module associated with the given name
   Module *lookup(StringRef Name);
+  
+  /// \brief Returns the in-memory (virtual file) buffer with the given name
+  llvm::MemoryBuffer *lookupBuffer(StringRef Name);
 
   /// \brief Number of modules loaded
   unsigned size() const { return Chain.size(); }
 
   /// \brief Creates a new module and adds it to the list of known modules
   Module &addModule(StringRef FileName, ModuleKind Type);
+  
+  /// \brief Add an in-memory buffer the list of known buffers
+  void addInMemoryBuffer(StringRef FileName, llvm::MemoryBuffer *Buffer);
 
   /// \brief Exports the list of loaded modules with their corresponding names
   void exportLookup(SmallVector<ModuleOffset, 16> &Target);
@@ -559,10 +568,6 @@ private:
   /// \brief The AST consumer.
   ASTConsumer *Consumer;
 
-  /// \brief AST buffers for chained PCHs created and stored in memory.
-  /// First (not depending on another) PCH in chain is in front.
-  std::vector<llvm::MemoryBuffer *> ASTBuffers;
-
   /// \brief The module manager which manages modules and their dependencies
   ModuleManager ModuleMgr;
 
@@ -1138,11 +1143,9 @@ public:
   /// \brief Sets and initializes the given Context.
   void InitializeContext(ASTContext &Context);
 
-  /// \brief Set AST buffers for chained PCHs created and stored in memory.
-  /// First (not depending on another) PCH in chain is first in array.
-  void setASTMemoryBuffers(llvm::MemoryBuffer **bufs, unsigned numBufs) {
-    ASTBuffers.clear();
-    ASTBuffers.insert(ASTBuffers.begin(), bufs, bufs + numBufs);
+  /// \brief Add in-memory (virtual file) buffer.
+  void addInMemoryBuffer(StringRef &FileName, llvm::MemoryBuffer *Buffer) {
+    ModuleMgr.addInMemoryBuffer(FileName, Buffer);
   }
 
   /// \brief Retrieve the name of the named (primary) AST file
diff --git a/lib/Serialization/ASTReader.cpp b/lib/Serialization/ASTReader.cpp
index 0241b22f9e..1b351a30dc 100644
--- a/lib/Serialization/ASTReader.cpp
+++ b/lib/Serialization/ASTReader.cpp
@@ -2764,9 +2764,8 @@ ASTReader::ASTReadResult ASTReader::ReadASTCore(StringRef FileName,
     if (CurrentDir.empty()) CurrentDir = ".";
   }
 
-  if (!ASTBuffers.empty()) {
-    F.Buffer.reset(ASTBuffers.back());
-    ASTBuffers.pop_back();
+  if (llvm::MemoryBuffer *Buffer = ModuleMgr.lookupBuffer(FileName)) {
+    F.Buffer.reset(Buffer);
     assert(F.Buffer && "Passed null buffer");
   } else {
     // Open the AST file.
@@ -5586,6 +5585,11 @@ Module *ModuleManager::lookup(StringRef Name) {
   return Modules[Entry];
 }
 
+llvm::MemoryBuffer *ModuleManager::lookupBuffer(StringRef Name) {
+  const FileEntry *Entry = FileMgr.getFile(Name);
+  return InMemoryBuffers[Entry];
+}
+
 /// \brief Creates a new module and adds it to the list of known modules
 Module &ModuleManager::addModule(StringRef FileName, ModuleKind Type) {
   Module *Prev = !size() ? 0 : &getLastModule();
@@ -5605,6 +5609,13 @@ Module &ModuleManager::addModule(StringRef FileName, ModuleKind Type) {
   return *Current;
 }
 
+void ModuleManager::addInMemoryBuffer(StringRef FileName, 
+  llvm::MemoryBuffer *Buffer) {
+  
+  const FileEntry *Entry = FileMgr.getVirtualFile(FileName, 
+    Buffer->getBufferSize(), 0);
+  InMemoryBuffers[Entry] = Buffer;
+}
 /// \brief Exports the list of loaded modules with their corresponding names
 void ModuleManager::exportLookup(SmallVector<ModuleOffset, 16> &Target) {
   Target.reserve(size());
diff --git a/lib/Serialization/ChainedIncludesSource.cpp b/lib/Serialization/ChainedIncludesSource.cpp
index a4a9f08f99..000e76846e 100644
--- a/lib/Serialization/ChainedIncludesSource.cpp
+++ b/lib/Serialization/ChainedIncludesSource.cpp
@@ -26,15 +26,18 @@
 using namespace clang;
 
 static ASTReader *createASTReader(CompilerInstance &CI,
-                                  StringRef pchFile,
-                                  llvm::MemoryBuffer **memBufs,
-                                  unsigned numBufs,
+                                  StringRef pchFile,  
+                                  SmallVector<llvm::MemoryBuffer *, 4> &memBufs,
+                                  SmallVector<std::string, 4> &bufNames,
                              ASTDeserializationListener *deserialListener = 0) {
   Preprocessor &PP = CI.getPreprocessor();
   llvm::OwningPtr<ASTReader> Reader;
   Reader.reset(new ASTReader(PP, &CI.getASTContext(), /*isysroot=*/"",
                              /*DisableValidation=*/true));
-  Reader->setASTMemoryBuffers(memBufs, numBufs);
+  for (unsigned ti = 0; ti < bufNames.size(); ++ti) {
+    StringRef sr(bufNames[ti]);
+    Reader->addInMemoryBuffer(sr, memBufs[ti]);
+  }
   Reader->setDeserializationListener(deserialListener);
   switch (Reader->ReadAST(pchFile, serialization::MK_PCH)) {
   case ASTReader::Success:
@@ -63,6 +66,7 @@ ChainedIncludesSource *ChainedIncludesSource::create(CompilerInstance &CI) {
   InputKind IK = CI.getFrontendOpts().Inputs[0].first;
 
   SmallVector<llvm::MemoryBuffer *, 4> serialBufs;
+  SmallVector<std::string, 4> serialBufNames;
 
   for (unsigned i = 0, e = includes.size(); i != e; ++i) {
     bool firstInclude = (i == 0);
@@ -125,9 +129,13 @@ ChainedIncludesSource *ChainedIncludesSource::create(CompilerInstance &CI) {
       llvm::raw_string_ostream os(pchName);
       os << ".pch" << i-1;
       os.flush();
+      
+      serialBufNames.push_back(pchName);
+
       llvm::OwningPtr<ExternalASTSource> Reader;
-      Reader.reset(createASTReader(*Clang, pchName, bufs.data(), bufs.size(),
-                      Clang->getASTConsumer().GetASTDeserializationListener()));
+
+      Reader.reset(createASTReader(*Clang, pchName, bufs, serialBufNames, 
+        Clang->getASTConsumer().GetASTDeserializationListener()));
       if (!Reader)
         return 0;
       Clang->getASTContext().setExternalSource(Reader);
@@ -147,9 +155,9 @@ ChainedIncludesSource *ChainedIncludesSource::create(CompilerInstance &CI) {
 
   assert(!serialBufs.empty());
   std::string pchName = includes.back() + ".pch-final";
+  serialBufNames.push_back(pchName);
   llvm::OwningPtr<ASTReader> Reader;
-  Reader.reset(createASTReader(CI, pchName,
-                               serialBufs.data(), serialBufs.size()));
+  Reader.reset(createASTReader(CI, pchName, serialBufs, serialBufNames));
   if (!Reader)
     return 0;