2 files changed, 36 insertions, 27 deletions

diff --git a/lib/Serialization/ASTReader.cpp b/lib/Serialization/ASTReader.cpp
index 1d7ea1346f..d81e9d8a36 100644
--- a/lib/Serialization/ASTReader.cpp
+++ b/lib/Serialization/ASTReader.cpp
@@ -2045,7 +2045,7 @@ ASTReader::ReadASTBlock(Module &F) {
         Idx += Length;
 
         // Load the AST file.
-        switch(ReadASTCore(ImportedFile, ImportedKind)) {
+        switch(ReadASTCore(ImportedFile, ImportedKind, &F)) {
         case Failure: return Failure;
           // If we have to ignore the dependency, we'll have to ignore this too.
         case IgnorePCH: return IgnorePCH;
@@ -2724,7 +2724,7 @@ ASTReader::ASTReadResult ASTReader::validateFileEntries() {
 
 ASTReader::ASTReadResult ASTReader::ReadAST(const std::string &FileName,
                                             ModuleKind Type) {
-  switch(ReadASTCore(FileName, Type)) {
+  switch(ReadASTCore(FileName, Type, /*ImportedBy=*/0)) {
   case Failure: return Failure;
   case IgnorePCH: return IgnorePCH;
   case Success: break;
@@ -2829,8 +2829,9 @@ ASTReader::ASTReadResult ASTReader::ReadAST(const std::string &FileName,
 }
 
 ASTReader::ASTReadResult ASTReader::ReadASTCore(StringRef FileName,
-                                                ModuleKind Type) {
-  Module &F = ModuleMgr.addModule(FileName, Type);
+                                                ModuleKind Type,
+                                                Module *ImportedBy) {
+  Module &F = ModuleMgr.addModule(FileName, Type, ImportedBy);
 
   if (FileName != "-") {
     CurrentDir = llvm::sys::path::parent_path(FileName);
@@ -5610,7 +5611,8 @@ ASTReader::~ASTReader() {
 }
 
 Module::Module(ModuleKind Kind)
-  : Kind(Kind), SizeInBits(0), LocalNumSLocEntries(0), SLocEntryBaseID(0),
+  : Kind(Kind), DirectlyImported(false), SizeInBits(0), 
+    LocalNumSLocEntries(0), SLocEntryBaseID(0),
     SLocEntryBaseOffset(0), SLocEntryOffsets(0),
     SLocFileOffsets(0), LocalNumIdentifiers(0), 
     IdentifierOffsets(0), BaseIdentifierID(0), IdentifierTableData(0),
@@ -5708,19 +5710,21 @@ llvm::MemoryBuffer *ModuleManager::lookupBuffer(StringRef Name) {
 }
 
 /// \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();
+Module &ModuleManager::addModule(StringRef FileName, ModuleKind Type,
+                                 Module *ImportedBy) {
   Module *Current = new Module(Type);
-
   Current->FileName = FileName.str();
-
   Chain.push_back(Current);
+
   const FileEntry *Entry = FileMgr.getFile(FileName);
+  // FIXME: Check whether we already loaded this module, before 
   Modules[Entry] = Current;
 
-  if (Prev) {
-    Current->ImportedBy.insert(Prev);
-    Prev->Imports.insert(Current);
+  if (ImportedBy) {
+    Current->ImportedBy.insert(ImportedBy);
+    ImportedBy->Imports.insert(Current);
+  } else {
+    Current->DirectlyImported = true;
   }
   
   return *Current;
diff --git a/lib/Serialization/ASTWriter.cpp b/lib/Serialization/ASTWriter.cpp
index 91c68df477..15a7cbc628 100644
--- a/lib/Serialization/ASTWriter.cpp
+++ b/lib/Serialization/ASTWriter.cpp
@@ -972,18 +972,23 @@ void ASTWriter::WriteMetadata(ASTContext &Context, StringRef isysroot,
   Stream.EmitRecordWithBlob(MetaAbbrevCode, Record, Triple);
 
   if (Chain) {
-    // FIXME: Add all of the "directly imported" modules, not just
-    // "the one we're chaining to".
     serialization::ModuleManager &Mgr = Chain->getModuleManager();
     llvm::SmallVector<char, 128> ModulePaths;
     Record.clear();
-    Module &PrimaryModule = Mgr.getPrimaryModule();
-    Record.push_back((unsigned)PrimaryModule.Kind); // FIXME: Stable encoding
-    // FIXME: Write import location, once it matters.
-    // FIXME: This writes the absolute path for AST files we depend on.
-    const std::string &MainFileName = PrimaryModule.FileName;
-    Record.push_back(MainFileName.size());
-    Record.append(MainFileName.begin(), MainFileName.end());
+
+    for (ModuleManager::ModuleIterator M = Mgr.begin(), MEnd = Mgr.end();
+         M != MEnd; ++M) {
+      // Skip modules that weren't directly imported.
+      if (!(*M)->isDirectlyImported())
+        continue;
+
+      Record.push_back((unsigned)(*M)->Kind); // FIXME: Stable encoding
+      // FIXME: Write import location, once it matters.
+      // FIXME: This writes the absolute path for AST files we depend on.
+      const std::string &FileName = (*M)->FileName;
+      Record.push_back(FileName.size());
+      Record.append(FileName.begin(), FileName.end());
+    }
     Stream.EmitRecord(IMPORTS, Record);
   }
 
@@ -3855,7 +3860,7 @@ void ASTWriter::AddCXXDefinitionData(const CXXRecordDecl *D, RecordDataImpl &Rec
 
 void ASTWriter::ReaderInitialized(ASTReader *Reader) {
   assert(Reader && "Cannot remove chain");
-  assert(!Chain && "Cannot replace chain");
+  assert((!Chain || Chain == Reader) && "Cannot replace chain");
   assert(FirstDeclID == NextDeclID &&
          FirstTypeID == NextTypeID &&
          FirstIdentID == NextIdentID &&
@@ -3865,11 +3870,11 @@ void ASTWriter::ReaderInitialized(ASTReader *Reader) {
 
   Chain = Reader;
 
-  FirstDeclID += Chain->getTotalNumDecls();
-  FirstTypeID += Chain->getTotalNumTypes();
-  FirstIdentID += Chain->getTotalNumIdentifiers();
-  FirstSelectorID += Chain->getTotalNumSelectors();
-  FirstMacroID += Chain->getTotalNumMacroDefinitions();
+  FirstDeclID = NUM_PREDEF_DECL_IDS + Chain->getTotalNumDecls();
+  FirstTypeID = NUM_PREDEF_TYPE_IDS + Chain->getTotalNumTypes();
+  FirstIdentID = NUM_PREDEF_IDENT_IDS + Chain->getTotalNumIdentifiers();
+  FirstSelectorID = NUM_PREDEF_SELECTOR_IDS + Chain->getTotalNumSelectors();
+  FirstMacroID = NUM_PREDEF_MACRO_IDS + Chain->getTotalNumMacroDefinitions();
   NextDeclID = FirstDeclID;
   NextTypeID = FirstTypeID;
   NextIdentID = FirstIdentID;