-frewrite-includes: Rework how includes and modules are differentiated

The map of FileChange structs here was storing two disjoint types of
information:

1. A pointer to the Module that an #include directive implicitly
   imported

2. A FileID and FileType for an included file. These would be left
   uninitialized in the Module case.

This change splits these two kinds of information into their own maps,
which both simplifies how we access either and avoids the undefined
behaviour we were hitting due to the uninitialized fields in the
included file case.

Mostly NFC, but fixes some errors found by self-host with ubsan.

git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@241140 91177308-0d34-0410-b5e6-96231b3b80d8

lib/Frontend/Rewrite/InclusionRewriter.cpp

@@ -29,13 +29,11 @@ namespace {
 class InclusionRewriter : public PPCallbacks {
   /// Information about which #includes were actually performed,
   /// created by preprocessor callbacks.
-  struct FileChange {
-    const Module *Mod;
-    SourceLocation From;
+  struct IncludedFile {
     FileID Id;
     SrcMgr::CharacteristicKind FileType;
-    FileChange(SourceLocation From, const Module *Mod) : Mod(Mod), From(From) {
-    }
+    IncludedFile(FileID Id, SrcMgr::CharacteristicKind FileType)
+        : Id(Id), FileType(FileType) {}
   };
   Preprocessor &PP; ///< Used to find inclusion directives.
   SourceManager &SM; ///< Used to read and manage source files.
@@ -44,11 +42,13 @@ class InclusionRewriter : public PPCallbacks {
   const llvm::MemoryBuffer *PredefinesBuffer; ///< The preprocessor predefines.
   bool ShowLineMarkers; ///< Show #line markers.
   bool UseLineDirectives; ///< Use of line directives or line markers.
-  typedef std::map<unsigned, FileChange> FileChangeMap;
-  FileChangeMap FileChanges; ///< Tracks which files were included where.
-  /// Used transitively for building up the FileChanges mapping over the
+  /// Tracks where inclusions that change the file are found.
+  std::map<unsigned, IncludedFile> FileIncludes;
+  /// Tracks where inclusions that import modules are found.
+  std::map<unsigned, const Module *> ModuleIncludes;
+  /// Used transitively for building up the FileIncludes mapping over the
   /// various \c PPCallbacks callbacks.
-  FileChangeMap::iterator LastInsertedFileChange;
+  SourceLocation LastInclusionLocation;
 public:
   InclusionRewriter(Preprocessor &PP, raw_ostream &OS, bool ShowLineMarkers,
                     bool UseLineDirectives);
@@ -82,7 +82,8 @@ private:
   bool HandleHasInclude(FileID FileId, Lexer &RawLex,
                         const DirectoryLookup *Lookup, Token &Tok,
                         bool &FileExists);
-  const FileChange *FindFileChangeLocation(SourceLocation Loc) const;
+  const IncludedFile *FindIncludeAtLocation(SourceLocation Loc) const;
+  const Module *FindModuleAtLocation(SourceLocation Loc) const;
   StringRef NextIdentifierName(Lexer &RawLex, Token &RawToken);
 };
 
@@ -95,7 +96,7 @@ InclusionRewriter::InclusionRewriter(Preprocessor &PP, raw_ostream &OS,
     : PP(PP), SM(PP.getSourceManager()), OS(OS), MainEOL("\n"),
       PredefinesBuffer(nullptr), ShowLineMarkers(ShowLineMarkers),
       UseLineDirectives(UseLineDirectives),
-      LastInsertedFileChange(FileChanges.end()) {}
+      LastInclusionLocation(SourceLocation()) {}
 
 /// Write appropriate line information as either #line directives or GNU line
 /// markers depending on what mode we're in, including the \p Filename and
@@ -143,12 +144,14 @@ void InclusionRewriter::FileChanged(SourceLocation Loc,
                                     FileID) {
   if (Reason != EnterFile)
     return;
-  if (LastInsertedFileChange == FileChanges.end())
+  if (LastInclusionLocation.isInvalid())
     // we didn't reach this file (eg: the main file) via an inclusion directive
     return;
-  LastInsertedFileChange->second.Id = FullSourceLoc(Loc, SM).getFileID();
-  LastInsertedFileChange->second.FileType = NewFileType;
-  LastInsertedFileChange = FileChanges.end();
+  FileID Id = FullSourceLoc(Loc, SM).getFileID();
+  auto P = FileIncludes.emplace(LastInclusionLocation.getRawEncoding(),
+                                IncludedFile(Id, NewFileType));
+  assert(P.second && "Unexpected revisitation of the same include directive");
+  LastInclusionLocation = SourceLocation();
 }
 
 /// Called whenever an inclusion is skipped due to canonical header protection
@@ -156,10 +159,9 @@ void InclusionRewriter::FileChanged(SourceLocation Loc,
 void InclusionRewriter::FileSkipped(const FileEntry &/*SkippedFile*/,
                                     const Token &/*FilenameTok*/,
                                     SrcMgr::CharacteristicKind /*FileType*/) {
-  assert(LastInsertedFileChange != FileChanges.end() && "A file, that wasn't "
-    "found via an inclusion directive, was skipped");
-  FileChanges.erase(LastInsertedFileChange);
-  LastInsertedFileChange = FileChanges.end();
+  assert(!LastInclusionLocation.isInvalid() &&
+         "A file, that wasn't found via an inclusion directive, was skipped");
+  LastInclusionLocation = SourceLocation();
 }
 
 /// This should be called whenever the preprocessor encounters include
@@ -176,25 +178,36 @@ void InclusionRewriter::InclusionDirective(SourceLocation HashLoc,
                                            StringRef /*SearchPath*/,
                                            StringRef /*RelativePath*/,
                                            const Module *Imported) {
-  assert(LastInsertedFileChange == FileChanges.end() && "Another inclusion "
-    "directive was found before the previous one was processed");
-  std::pair<FileChangeMap::iterator, bool> p = FileChanges.insert(
-    std::make_pair(HashLoc.getRawEncoding(), FileChange(HashLoc, Imported)));
-  assert(p.second && "Unexpected revisitation of the same include directive");
-  if (!Imported)
-    LastInsertedFileChange = p.first;
+  assert(LastInclusionLocation.isInvalid() &&
+         "Another inclusion directive was found before the previous one "
+         "was processed");
+  if (Imported) {
+    auto P = ModuleIncludes.emplace(HashLoc.getRawEncoding(), Imported);
+    assert(P.second && "Unexpected revisitation of the same include directive");
+  } else
+    LastInclusionLocation = HashLoc;
 }
 
 /// Simple lookup for a SourceLocation (specifically one denoting the hash in
 /// an inclusion directive) in the map of inclusion information, FileChanges.
-const InclusionRewriter::FileChange *
-InclusionRewriter::FindFileChangeLocation(SourceLocation Loc) const {
-  FileChangeMap::const_iterator I = FileChanges.find(Loc.getRawEncoding());
-  if (I != FileChanges.end())
+const InclusionRewriter::IncludedFile *
+InclusionRewriter::FindIncludeAtLocation(SourceLocation Loc) const {
+  const auto I = FileIncludes.find(Loc.getRawEncoding());
+  if (I != FileIncludes.end())
     return &I->second;
   return nullptr;
 }
 
+/// Simple lookup for a SourceLocation (specifically one denoting the hash in
+/// an inclusion directive) in the map of module inclusion information.
+const Module *
+InclusionRewriter::FindModuleAtLocation(SourceLocation Loc) const {
+  const auto I = ModuleIncludes.find(Loc.getRawEncoding());
+  if (I != ModuleIncludes.end())
+    return I->second;
+  return nullptr;
+}
+
 /// Detect the likely line ending style of \p FromFile by examining the first
 /// newline found within it.
 static StringRef DetectEOL(const MemoryBuffer &FromFile) {
@@ -388,8 +401,7 @@ bool InclusionRewriter::Process(FileID FileId,
 {
   bool Invalid;
   const MemoryBuffer &FromFile = *SM.getBuffer(FileId, &Invalid);
-  if (Invalid) // invalid inclusion
-    return false;
+  assert(!Invalid && "Attempting to process invalid inclusion");
   const char *FileName = FromFile.getBufferIdentifier();
   Lexer RawLex(FileId, &FromFile, PP.getSourceManager(), PP.getLangOpts());
   RawLex.SetCommentRetentionState(false);
@@ -433,13 +445,12 @@ bool InclusionRewriter::Process(FileID FileId,
             if (FileId != PP.getPredefinesFileID())
               WriteLineInfo(FileName, Line - 1, FileType, "");
             StringRef LineInfoExtra;
-            if (const FileChange *Change = FindFileChangeLocation(
-                HashToken.getLocation())) {
-              if (Change->Mod) {
-                WriteImplicitModuleImport(Change->Mod);
-
-              // else now include and recursively process the file
-              } else if (Process(Change->Id, Change->FileType)) {
+            SourceLocation Loc = HashToken.getLocation();
+            if (const Module *Mod = FindModuleAtLocation(Loc))
+              WriteImplicitModuleImport(Mod);
+            else if (const IncludedFile *Inc = FindIncludeAtLocation(Loc)) {
+              // include and recursively process the file
+              if (Process(Inc->Id, Inc->FileType)) {
                 // and set lineinfo back to this file, if the nested one was
                 // actually included
                 // `2' indicates returning to a file (after having included