[modules] Substantially improve handling of #undef:

 * Track override set across module load and save
 * Track originating module to allow proper re-export of #undef
 * Make override set properly transitive when it picks up a #undef

This fixes nearly all of the remaining macro issues with self-host.

llvm-svn: 213922

4 files changed, 215 insertions, 149 deletions

diff --git a/clang/lib/Lex/PPDirectives.cpp b/clang/lib/Lex/PPDirectives.cpp
index d7ed0b4859e..da50bba598f 100644
--- a/clang/lib/Lex/PPDirectives.cpp
+++ b/clang/lib/Lex/PPDirectives.cpp
@@ -64,25 +64,30 @@ MacroInfo *Preprocessor::AllocateDeserializedMacroInfo(SourceLocation L,
 
 DefMacroDirective *
 Preprocessor::AllocateDefMacroDirective(MacroInfo *MI, SourceLocation Loc,
-                                        bool isImported) {
-  DefMacroDirective *MD = BP.Allocate<DefMacroDirective>();
-  new (MD) DefMacroDirective(MI, Loc, isImported);
-  return MD;
+                                        unsigned ImportedFromModuleID,
+                                        ArrayRef<unsigned> Overrides) {
+  unsigned NumExtra = (ImportedFromModuleID ? 1 : 0) + Overrides.size();
+  return new (BP.Allocate(sizeof(DefMacroDirective) +
+                              sizeof(unsigned) * NumExtra,
+                          llvm::alignOf<DefMacroDirective>()))
+      DefMacroDirective(MI, Loc, ImportedFromModuleID, Overrides);
 }
 
 UndefMacroDirective *
-Preprocessor::AllocateUndefMacroDirective(SourceLocation UndefLoc) {
-  UndefMacroDirective *MD = BP.Allocate<UndefMacroDirective>();
-  new (MD) UndefMacroDirective(UndefLoc);
-  return MD;
+Preprocessor::AllocateUndefMacroDirective(SourceLocation UndefLoc,
+                                          unsigned ImportedFromModuleID,
+                                          ArrayRef<unsigned> Overrides) {
+  unsigned NumExtra = (ImportedFromModuleID ? 1 : 0) + Overrides.size();
+  return new (BP.Allocate(sizeof(UndefMacroDirective) +
+                              sizeof(unsigned) * NumExtra,
+                          llvm::alignOf<UndefMacroDirective>()))
+      UndefMacroDirective(UndefLoc, ImportedFromModuleID, Overrides);
 }
 
 VisibilityMacroDirective *
 Preprocessor::AllocateVisibilityMacroDirective(SourceLocation Loc,
                                                bool isPublic) {
-  VisibilityMacroDirective *MD = BP.Allocate<VisibilityMacroDirective>();
-  new (MD) VisibilityMacroDirective(Loc, isPublic);
-  return MD;
+  return new (BP) VisibilityMacroDirective(Loc, isPublic);
 }
 
 /// \brief Clean up a MacroInfo that was allocated but not used due to an
diff --git a/clang/lib/Lex/Pragma.cpp b/clang/lib/Lex/Pragma.cpp
index a5f1a62a708..6a8b0a6bd31 100644
--- a/clang/lib/Lex/Pragma.cpp
+++ b/clang/lib/Lex/Pragma.cpp
@@ -605,7 +605,7 @@ void Preprocessor::HandlePragmaPopMacro(Token &PopMacroTok) {
     if (MacroToReInstall) {
       // Reinstall the previously pushed macro.
       appendDefMacroDirective(IdentInfo, MacroToReInstall, MessageLoc,
-                              /*isImported=*/false);
+                              /*isImported=*/false, /*Overrides*/None);
     }
 
     // Pop PragmaPushMacroInfo stack.
diff --git a/clang/lib/Serialization/ASTReader.cpp b/clang/lib/Serialization/ASTReader.cpp
index e3a16c3c5f2..6ed41c39702 100644
--- a/clang/lib/Serialization/ASTReader.cpp
+++ b/clang/lib/Serialization/ASTReader.cpp
@@ -774,12 +774,13 @@ IdentifierInfo *ASTIdentifierLookupTrait::ReadData(const internal_key_type& k,
     DataLen -= 4;
     SmallVector<uint32_t, 8> LocalMacroIDs;
     if (hasSubmoduleMacros) {
-      while (uint32_t LocalMacroID =
-                 endian::readNext<uint32_t, little, unaligned>(d)) {
+      while (true) {
+        uint32_t LocalMacroID =
+            endian::readNext<uint32_t, little, unaligned>(d);
         DataLen -= 4;
+        if (LocalMacroID == 0xdeadbeef) break;
         LocalMacroIDs.push_back(LocalMacroID);
       }
-      DataLen -= 4;
     }
 
     if (F.Kind == MK_Module) {
@@ -1732,10 +1733,12 @@ struct ASTReader::ModuleMacroInfo {
     return llvm::makeArrayRef(Overrides + 1, *Overrides);
   }
 
-  DefMacroDirective *import(Preprocessor &PP, SourceLocation ImportLoc) const {
+  MacroDirective *import(Preprocessor &PP, SourceLocation ImportLoc) const {
     if (!MI)
-      return nullptr;
-    return PP.AllocateDefMacroDirective(MI, ImportLoc, /*isImported=*/true);
+      return PP.AllocateUndefMacroDirective(ImportLoc, SubModID,
+                                            getOverriddenSubmodules());
+    return PP.AllocateDefMacroDirective(MI, ImportLoc, SubModID,
+                                        getOverriddenSubmodules());
   }
 };
 
@@ -1790,7 +1793,7 @@ void ASTReader::resolvePendingMacro(IdentifierInfo *II,
     // install if we make this module visible.
     HiddenNamesMap[Owner].HiddenMacros.insert(std::make_pair(II, MMI));
   } else {
-    installImportedMacro(II, MMI, Owner, /*FromFinalization*/false);
+    installImportedMacro(II, MMI, Owner);
   }
 }
 
@@ -1828,23 +1831,36 @@ void ASTReader::installPCHMacroDirectives(IdentifierInfo *II,
     case MacroDirective::MD_Define: {
       GlobalMacroID GMacID = getGlobalMacroID(M, Record[Idx++]);
       MacroInfo *MI = getMacro(GMacID);
-      bool isImported = Record[Idx++];
-      bool isAmbiguous = Record[Idx++];
+      SubmoduleID ImportedFrom = Record[Idx++];
+      bool IsAmbiguous = Record[Idx++];
+      llvm::SmallVector<unsigned, 4> Overrides;
+      if (ImportedFrom) {
+        Overrides.insert(Overrides.end(),
+                         &Record[Idx] + 1, &Record[Idx] + 1 + Record[Idx]);
+        Idx += Overrides.size() + 1;
+      }
       DefMacroDirective *DefMD =
-          PP.AllocateDefMacroDirective(MI, Loc, isImported);
-      DefMD->setAmbiguous(isAmbiguous);
+          PP.AllocateDefMacroDirective(MI, Loc, ImportedFrom, Overrides);
+      DefMD->setAmbiguous(IsAmbiguous);
       MD = DefMD;
       break;
     }
-    case MacroDirective::MD_Undefine:
-      MD = PP.AllocateUndefMacroDirective(Loc);
+    case MacroDirective::MD_Undefine: {
+      SubmoduleID ImportedFrom = Record[Idx++];
+      llvm::SmallVector<unsigned, 4> Overrides;
+      if (ImportedFrom) {
+        Overrides.insert(Overrides.end(),
+                         &Record[Idx] + 1, &Record[Idx] + 1 + Record[Idx]);
+        Idx += Overrides.size() + 1;
+      }
+      MD = PP.AllocateUndefMacroDirective(Loc, ImportedFrom, Overrides);
       break;
-    case MacroDirective::MD_Visibility: {
+    }
+    case MacroDirective::MD_Visibility:
       bool isPublic = Record[Idx++];
       MD = PP.AllocateVisibilityMacroDirective(Loc, isPublic);
       break;
     }
-    }
 
     if (!Latest)
       Latest = MD;
@@ -1877,6 +1893,7 @@ static bool areDefinedInSystemModules(MacroInfo *PrevMI, MacroInfo *NewMI,
 }
 
 void ASTReader::removeOverriddenMacros(IdentifierInfo *II,
+                                       SourceLocation ImportLoc,
                                        AmbiguousMacros &Ambig,
                                        ArrayRef<SubmoduleID> Overrides) {
   for (unsigned OI = 0, ON = Overrides.size(); OI != ON; ++OI) {
@@ -1887,9 +1904,12 @@ void ASTReader::removeOverriddenMacros(IdentifierInfo *II,
     HiddenNames &Hidden = HiddenNamesMap[Owner];
     HiddenMacrosMap::iterator HI = Hidden.HiddenMacros.find(II);
     if (HI != Hidden.HiddenMacros.end()) {
+      // Register the macro now so we don't lose it when we re-export.
+      PP.appendMacroDirective(II, HI->second->import(PP, ImportLoc));
+
       auto SubOverrides = HI->second->getOverriddenSubmodules();
       Hidden.HiddenMacros.erase(HI);
-      removeOverriddenMacros(II, Ambig, SubOverrides);
+      removeOverriddenMacros(II, ImportLoc, Ambig, SubOverrides);
     }
 
     // If this macro is already in our list of conflicts, remove it from there.
@@ -1903,6 +1923,7 @@ void ASTReader::removeOverriddenMacros(IdentifierInfo *II,
 
 ASTReader::AmbiguousMacros *
 ASTReader::removeOverriddenMacros(IdentifierInfo *II,
+                                  SourceLocation ImportLoc,
                                   ArrayRef<SubmoduleID> Overrides) {
   MacroDirective *Prev = PP.getMacroDirective(II);
   if (!Prev && Overrides.empty())
@@ -1915,7 +1936,7 @@ ASTReader::removeOverriddenMacros(IdentifierInfo *II,
     AmbiguousMacros &Ambig = AmbiguousMacroDefs[II];
     Ambig.push_back(PrevDef);
 
-    removeOverriddenMacros(II, Ambig, Overrides);
+    removeOverriddenMacros(II, ImportLoc, Ambig, Overrides);
 
     if (!Ambig.empty())
       return &Ambig;
@@ -1927,7 +1948,7 @@ ASTReader::removeOverriddenMacros(IdentifierInfo *II,
     if (PrevDef)
       Ambig.push_back(PrevDef);
 
-    removeOverriddenMacros(II, Ambig, Overrides);
+    removeOverriddenMacros(II, ImportLoc, Ambig, Overrides);
 
     if (!Ambig.empty()) {
       AmbiguousMacros &Result = AmbiguousMacroDefs[II];
@@ -1941,11 +1962,11 @@ ASTReader::removeOverriddenMacros(IdentifierInfo *II,
 }
 
 void ASTReader::installImportedMacro(IdentifierInfo *II, ModuleMacroInfo *MMI,
-                                     Module *Owner, bool FromFinalization) {
+                                     Module *Owner) {
   assert(II && Owner);
 
   SourceLocation ImportLoc = Owner->MacroVisibilityLoc;
-  if (ImportLoc.isInvalid() && !FromFinalization) {
+  if (ImportLoc.isInvalid()) {
     // FIXME: If we made macros from this module visible but didn't provide a
     // source location for the import, we don't have a location for the macro.
     // Use the location at which the containing module file was first imported
@@ -1955,18 +1976,16 @@ void ASTReader::installImportedMacro(IdentifierInfo *II, ModuleMacroInfo *MMI,
   }
 
   AmbiguousMacros *Prev =
-      removeOverriddenMacros(II, MMI->getOverriddenSubmodules());
+      removeOverriddenMacros(II, ImportLoc, MMI->getOverriddenSubmodules());
 
   // Create a synthetic macro definition corresponding to the import (or null
   // if this was an undefinition of the macro).
-  DefMacroDirective *MD = MMI->import(PP, ImportLoc);
+  MacroDirective *Imported = MMI->import(PP, ImportLoc);
+  DefMacroDirective *MD = dyn_cast<DefMacroDirective>(Imported);
 
   // If there's no ambiguity, just install the macro.
   if (!Prev) {
-    if (MD)
-      PP.appendMacroDirective(II, MD);
-    else
-      PP.appendMacroDirective(II, PP.AllocateUndefMacroDirective(ImportLoc));
+    PP.appendMacroDirective(II, Imported);
     return;
   }
   assert(!Prev->empty());
@@ -1974,10 +1993,14 @@ void ASTReader::installImportedMacro(IdentifierInfo *II, ModuleMacroInfo *MMI,
   if (!MD) {
     // We imported a #undef that didn't remove all prior definitions. The most
     // recent prior definition remains, and we install it in the place of the
-    // imported directive.
+    // imported directive, as if by a local #pragma pop_macro.
     MacroInfo *NewMI = Prev->back()->getInfo();
     Prev->pop_back();
-    MD = PP.AllocateDefMacroDirective(NewMI, ImportLoc, /*Imported*/true);
+    MD = PP.AllocateDefMacroDirective(NewMI, ImportLoc);
+
+    // Install our #undef first so that we don't lose track of it. We'll replace
+    // this with whichever macro definition ends up winning.
+    PP.appendMacroDirective(II, Imported);
   }
 
   // We're introducing a macro definition that creates or adds to an ambiguity.
@@ -3336,8 +3359,13 @@ void ASTReader::makeNamesVisible(const HiddenNames &Names, Module *Owner,
 
   assert((FromFinalization || Owner->NameVisibility >= Module::MacrosVisible) &&
          "nothing to make visible?");
-  for (const auto &Macro : Names.HiddenMacros)
-    installImportedMacro(Macro.first, Macro.second, Owner, FromFinalization);
+  for (const auto &Macro : Names.HiddenMacros) {
+    if (FromFinalization)
+      PP.appendMacroDirective(Macro.first,
+                              Macro.second->import(PP, SourceLocation()));
+    else
+      installImportedMacro(Macro.first, Macro.second, Owner);
+  }
 }
 
 void ASTReader::makeModuleVisible(Module *Mod,
diff --git a/clang/lib/Serialization/ASTWriter.cpp b/clang/lib/Serialization/ASTWriter.cpp
index 82932426032..ac38f303800 100644
--- a/clang/lib/Serialization/ASTWriter.cpp
+++ b/clang/lib/Serialization/ASTWriter.cpp
@@ -1901,10 +1901,8 @@ static bool shouldIgnoreMacro(MacroDirective *MD, bool IsModule,
 
   if (IsModule) {
     // Re-export any imported directives.
-    // FIXME: Also ensure we re-export imported #undef directives.
-    if (auto *DMD = dyn_cast<DefMacroDirective>(MD))
-      if (DMD->isImported())
-        return false;
+    if (MD->isImported())
+      return false;
 
     SourceLocation Loc = MD->getLocation();
     if (Loc.isInvalid())
@@ -1983,16 +1981,24 @@ void ASTWriter::WritePreprocessor(const Preprocessor &PP, bool IsModule) {
 
       AddSourceLocation(MD->getLocation(), Record);
       Record.push_back(MD->getKind());
-      if (DefMacroDirective *DefMD = dyn_cast<DefMacroDirective>(MD)) {
+      if (auto *DefMD = dyn_cast<DefMacroDirective>(MD)) {
         MacroID InfoID = getMacroRef(DefMD->getInfo(), Name);
         Record.push_back(InfoID);
-        Record.push_back(DefMD->isImported());
+        Record.push_back(DefMD->getOwningModuleID());
         Record.push_back(DefMD->isAmbiguous());
-
-      } else if (VisibilityMacroDirective *
-                   VisMD = dyn_cast<VisibilityMacroDirective>(MD)) {
+      } else if (auto *UndefMD = dyn_cast<UndefMacroDirective>(MD)) {
+        Record.push_back(UndefMD->getOwningModuleID());
+      } else {
+        auto *VisMD = cast<VisibilityMacroDirective>(MD);
         Record.push_back(VisMD->isPublic());
       }
+
+      if (MD->isImported()) {
+        auto Overrides = MD->getOverriddenModules();
+        Record.push_back(Overrides.size());
+        for (auto Override : Overrides)
+          Record.push_back(Override);
+      }
     }
     if (Record.empty())
       continue;
@@ -2992,115 +2998,140 @@ class ASTIdentifierTableTrait {
     if (Macro || (Macro = PP.getMacroDirectiveHistory(II))) {
       if (!IsModule)
         return !shouldIgnoreMacro(Macro, IsModule, PP);
-      SubmoduleID ModID;
-      if (getFirstPublicSubmoduleMacro(Macro, ModID))
+
+      MacroState State;
+      if (getFirstPublicSubmoduleMacro(Macro, State))
         return true;
     }
 
     return false;
   }
 
-  typedef llvm::SmallVectorImpl<SubmoduleID> OverriddenList;
+  enum class SubmoduleMacroState {
+    /// We've seen nothing about this macro.
+    None,
+    /// We've seen a public visibility directive.
+    Public,
+    /// We've either exported a macro for this module or found that the
+    /// module's definition of this macro is private.
+    Done
+  };
+  typedef llvm::DenseMap<SubmoduleID, SubmoduleMacroState> MacroState;
 
   MacroDirective *
-  getFirstPublicSubmoduleMacro(MacroDirective *MD, SubmoduleID &ModID) {
-    ModID = 0;
-    llvm::SmallVector<SubmoduleID, 1> Overridden;
-    if (MacroDirective *NextMD = getPublicSubmoduleMacro(MD, ModID, Overridden))
-      if (!shouldIgnoreMacro(NextMD, IsModule, PP))
-        return NextMD;
+  getFirstPublicSubmoduleMacro(MacroDirective *MD, MacroState &State) {
+    if (MacroDirective *NextMD = getPublicSubmoduleMacro(MD, State))
+      return NextMD;
     return nullptr;
   }
 
   MacroDirective *
-  getNextPublicSubmoduleMacro(MacroDirective *MD, SubmoduleID &ModID,
-                              OverriddenList &Overridden) {
+  getNextPublicSubmoduleMacro(MacroDirective *MD, MacroState &State) {
     if (MacroDirective *NextMD =
-            getPublicSubmoduleMacro(MD->getPrevious(), ModID, Overridden))
-      if (!shouldIgnoreMacro(NextMD, IsModule, PP))
-        return NextMD;
+            getPublicSubmoduleMacro(MD->getPrevious(), State))
+      return NextMD;
     return nullptr;
   }
 
-  /// \brief Traverses the macro directives history and returns the latest
-  /// public macro definition or undefinition that is not in ModID.
+  /// \brief Traverses the macro directives history and returns the next
+  /// public macro definition or undefinition that has not been found so far.
+  ///
   /// A macro that is defined in submodule A and undefined in submodule B
   /// will still be considered as defined/exported from submodule A.
-  /// ModID is updated to the module containing the returned directive.
-  ///
-  /// FIXME: This process breaks down if a module defines a macro, imports
-  ///        another submodule that changes the macro, then changes the
-  ///        macro again itself.
   MacroDirective *getPublicSubmoduleMacro(MacroDirective *MD,
-                                          SubmoduleID &ModID,
-                                          OverriddenList &Overridden) {
-    Overridden.clear();
+                                          MacroState &State) {
     if (!MD)
       return nullptr;
 
-    SubmoduleID OrigModID = ModID;
     Optional<bool> IsPublic;
     for (; MD; MD = MD->getPrevious()) {
-      SubmoduleID ThisModID = getSubmoduleID(MD);
-      if (ThisModID == 0) {
-        IsPublic = Optional<bool>();
-
-        // If we have no directive location, this macro was installed when
-        // finalizing the ASTReader.
-        if (DefMacroDirective *DefMD = dyn_cast<DefMacroDirective>(MD))
-          if (DefMD->getInfo()->getOwningModuleID())
-            return MD;
-        // Skip imports that only produce #undefs for now.
-        // FIXME: We should still re-export them!
+      // Once we hit an ignored macro, we're done: the rest of the chain
+      // will all be ignored macros.
+      if (shouldIgnoreMacro(MD, IsModule, PP))
+        break;
+
+      // If this macro was imported, re-export it.
+      if (MD->isImported())
+        return MD;
+
+      SubmoduleID ModID = getSubmoduleID(MD);
+      auto &S = State[ModID];
+      assert(ModID && "found macro in no submodule");
 
+      if (S == SubmoduleMacroState::Done)
         continue;
+
+      if (auto *VisMD = dyn_cast<VisibilityMacroDirective>(MD)) {
+        // The latest visibility directive for a name in a submodule affects all
+        // the directives that come before it.
+        if (S == SubmoduleMacroState::None)
+          S = VisMD->isPublic() ? SubmoduleMacroState::Public
+                                : SubmoduleMacroState::Done;
+      } else {
+        S = SubmoduleMacroState::Done;
+        return MD;
       }
-      if (ThisModID != ModID) {
-        ModID = ThisModID;
-        IsPublic = Optional<bool>();
-      }
+    }
+
+    return nullptr;
+  }
+
+  ArrayRef<SubmoduleID>
+  getOverriddenSubmodules(MacroDirective *MD,
+                          SmallVectorImpl<SubmoduleID> &ScratchSpace) {
+    assert(!isa<VisibilityMacroDirective>(MD) &&
+           "only #define and #undef can override");
+    if (MD->isImported())
+      return MD->getOverriddenModules();
+
+    ScratchSpace.clear();
+    SubmoduleID ModID = getSubmoduleID(MD);
+    for (MD = MD->getPrevious(); MD; MD = MD->getPrevious()) {
+      if (shouldIgnoreMacro(MD, IsModule, PP))
+        break;
 
       // If this is a definition from a submodule import, that submodule's
       // definition is overridden by the definition or undefinition that we
       // started with.
-      // FIXME: This should only apply to macros defined in OrigModID.
-      // We can't do that currently, because a #include of a different submodule
-      // of the same module just leaks through macros instead of providing new
-      // DefMacroDirectives for them.
-      if (DefMacroDirective *DefMD = dyn_cast<DefMacroDirective>(MD)) {
-        // Figure out which submodule the macro was originally defined within.
-        SubmoduleID SourceID = DefMD->getInfo()->getOwningModuleID();
-        if (!SourceID) {
-          SourceLocation DefLoc = DefMD->getInfo()->getDefinitionLoc();
-          if (DefLoc == MD->getLocation())
-            SourceID = ThisModID;
-          else
-            SourceID = Writer.inferSubmoduleIDFromLocation(DefLoc);
+      if (MD->isImported()) {
+        if (auto *DefMD = dyn_cast<DefMacroDirective>(MD)) {
+          SubmoduleID DefModuleID = DefMD->getInfo()->getOwningModuleID();
+          assert(DefModuleID && "imported macro has no owning module");
+          ScratchSpace.push_back(DefModuleID);
+        } else if (auto *UndefMD = dyn_cast<UndefMacroDirective>(MD)) {
+          // If we override a #undef, we override anything that #undef overrides.
+          // We don't need to override it, since an active #undef doesn't affect
+          // the meaning of a macro.
+          auto Overrides = UndefMD->getOverriddenModules();
+          ScratchSpace.insert(ScratchSpace.end(),
+                              Overrides.begin(), Overrides.end());
         }
-        if (OrigModID && SourceID != OrigModID)
-          Overridden.push_back(SourceID);
       }
 
-      // We are looking for a definition in a different submodule than the one
-      // that we started with. If a submodule has re-definitions of the same
-      // macro, only the last definition will be used as the "exported" one.
-      if (ModID == OrigModID)
-        continue;
-
-      // The latest visibility directive for a name in a submodule affects all
-      // the directives that come before it.
-      if (VisibilityMacroDirective *VisMD =
-              dyn_cast<VisibilityMacroDirective>(MD)) {
-        if (!IsPublic.hasValue())
-          IsPublic = VisMD->isPublic();
-      } else if (!IsPublic.hasValue() || IsPublic.getValue()) {
-        // FIXME: If we find an imported macro, we should include its list of
-        // overrides in our export.
-        return MD;
+      // Stop once we leave the original macro's submodule.
+      //
+      // Either this submodule #included another submodule of the same
+      // module or it just happened to be built after the other module.
+      // In the former case, we override the submodule's macro.
+      //
+      // FIXME: In the latter case, we shouldn't do so, but we can't tell
+      // these cases apart.
+      //
+      // FIXME: We can leave this submodule and re-enter it if it #includes a
+      // header within a different submodule of the same module. In such cases
+      // the overrides list will be incomplete.
+      SubmoduleID DirectiveModuleID = getSubmoduleID(MD);
+      if (DirectiveModuleID != ModID) {
+        if (DirectiveModuleID && !MD->isImported())
+          ScratchSpace.push_back(DirectiveModuleID);
+        break;
       }
     }
 
-    return nullptr;
+    std::sort(ScratchSpace.begin(), ScratchSpace.end());
+    ScratchSpace.erase(std::unique(ScratchSpace.begin(), ScratchSpace.end()),
+                       ScratchSpace.end());
+    return ScratchSpace;
   }
 
   SubmoduleID getSubmoduleID(MacroDirective *MD) {
@@ -3136,27 +3167,23 @@ public:
       if (hadMacroDefinition(II, Macro)) {
         DataLen += 4; // MacroDirectives offset.
         if (IsModule) {
-          SubmoduleID ModID;
-          llvm::SmallVector<SubmoduleID, 4> Overridden;
-          for (MacroDirective *
-                 MD = getFirstPublicSubmoduleMacro(Macro, ModID);
-                 MD; MD = getNextPublicSubmoduleMacro(MD, ModID, Overridden)) {
-            // Previous macro's overrides.
-            if (!Overridden.empty())
-              DataLen += 4 * (1 + Overridden.size());
+          MacroState State;
+          SmallVector<SubmoduleID, 16> Scratch;
+          for (MacroDirective *MD = getFirstPublicSubmoduleMacro(Macro, State);
+               MD; MD = getNextPublicSubmoduleMacro(MD, State)) {
             DataLen += 4; // MacroInfo ID or ModuleID.
+            if (unsigned NumOverrides =
+                    getOverriddenSubmodules(MD, Scratch).size())
+              DataLen += 4 * (1 + NumOverrides);
           }
-          // Previous macro's overrides.
-          if (!Overridden.empty())
-            DataLen += 4 * (1 + Overridden.size());
-          DataLen += 4;
+          DataLen += 4; // 0 terminator.
         }
       }
 
       for (IdentifierResolver::iterator D = IdResolver.begin(II),
                                      DEnd = IdResolver.end();
            D != DEnd; ++D)
-        DataLen += sizeof(DeclID);
+        DataLen += 4;
     }
     using namespace llvm::support;
     endian::Writer<little> LE(Out);
@@ -3183,8 +3210,10 @@ public:
       using namespace llvm::support;
       endian::Writer<little> LE(Out);
       LE.write<uint32_t>(Overridden.size() | 0x80000000U);
-      for (unsigned I = 0, N = Overridden.size(); I != N; ++I)
+      for (unsigned I = 0, N = Overridden.size(); I != N; ++I) {
+        assert(Overridden[I] && "zero module ID for override");
         LE.write<uint32_t>(Overridden[I]);
+      }
     }
   }
 
@@ -3216,24 +3245,28 @@ public:
       LE.write<uint32_t>(Writer.getMacroDirectivesOffset(II));
       if (IsModule) {
         // Write the IDs of macros coming from different submodules.
-        SubmoduleID ModID;
-        llvm::SmallVector<SubmoduleID, 4> Overridden;
-        for (MacroDirective *
-               MD = getFirstPublicSubmoduleMacro(Macro, ModID);
-               MD; MD = getNextPublicSubmoduleMacro(MD, ModID, Overridden)) {
-          MacroID InfoID = 0;
-          emitMacroOverrides(Out, Overridden);
+        MacroState State;
+        SmallVector<SubmoduleID, 16> Scratch;
+        for (MacroDirective *MD = getFirstPublicSubmoduleMacro(Macro, State);
+             MD; MD = getNextPublicSubmoduleMacro(MD, State)) {
           if (DefMacroDirective *DefMD = dyn_cast<DefMacroDirective>(MD)) {
-            InfoID = Writer.getMacroID(DefMD->getInfo());
+            // FIXME: If this macro directive was created by #pragma pop_macros,
+            // or if it was created implicitly by resolving conflicting macros,
+            // it may be for a different submodule from the one in the MacroInfo
+            // object. If so, we should write out its owning ModuleID.
+            MacroID InfoID = Writer.getMacroID(DefMD->getInfo());
             assert(InfoID);
             LE.write<uint32_t>(InfoID << 1);
           } else {
-            assert(isa<UndefMacroDirective>(MD));
-            LE.write<uint32_t>((ModID << 1) | 1);
+            auto *UndefMD = cast<UndefMacroDirective>(MD);
+            SubmoduleID Mod = UndefMD->isImported()
+                                  ? UndefMD->getOwningModuleID()
+                                  : getSubmoduleID(UndefMD);
+            LE.write<uint32_t>((Mod << 1) | 1);
           }
+          emitMacroOverrides(Out, getOverriddenSubmodules(MD, Scratch));
         }
-        emitMacroOverrides(Out, Overridden);
-        LE.write<uint32_t>(0);
+        LE.write<uint32_t>(0xdeadbeef);
       }
     }