--- a/clang/lib/Frontend/Rewrite/InclusionRewriter.cpp	Tue Jun 15 19:13:43 2021 +0900
+++ b/clang/lib/Frontend/Rewrite/InclusionRewriter.cpp	Tue Jun 15 19:15:29 2021 +0900
@@ -40,17 +40,17 @@
   SourceManager &SM; ///< Used to read and manage source files.
   raw_ostream &OS; ///< The destination stream for rewritten contents.
   StringRef MainEOL; ///< The line ending marker to use.
-  const llvm::MemoryBuffer *PredefinesBuffer; ///< The preprocessor predefines.
+  llvm::MemoryBufferRef PredefinesBuffer; ///< The preprocessor predefines.
   bool ShowLineMarkers; ///< Show #line markers.
   bool UseLineDirectives; ///< Use of line directives or line markers.
   /// Tracks where inclusions that change the file are found.
-  std::map<unsigned, IncludedFile> FileIncludes;
+  std::map<SourceLocation, IncludedFile> FileIncludes;
   /// Tracks where inclusions that import modules are found.
-  std::map<unsigned, const Module *> ModuleIncludes;
+  std::map<SourceLocation, const Module *> ModuleIncludes;
   /// Tracks where inclusions that enter modules (in a module build) are found.
-  std::map<unsigned, const Module *> ModuleEntryIncludes;
+  std::map<SourceLocation, const Module *> ModuleEntryIncludes;
   /// Tracks where #if and #elif directives get evaluated and whether to true.
-  std::map<unsigned, bool> IfConditions;
+  std::map<SourceLocation, bool> IfConditions;
   /// Used transitively for building up the FileIncludes mapping over the
   /// various \c PPCallbacks callbacks.
   SourceLocation LastInclusionLocation;
@@ -59,14 +59,14 @@
                     bool UseLineDirectives);
   void Process(FileID FileId, SrcMgr::CharacteristicKind FileType,
                const DirectoryLookup *DirLookup);
-  void setPredefinesBuffer(const llvm::MemoryBuffer *Buf) {
+  void setPredefinesBuffer(const llvm::MemoryBufferRef &Buf) {
     PredefinesBuffer = Buf;
   }
   void detectMainFileEOL();
   void handleModuleBegin(Token &Tok) {
     assert(Tok.getKind() == tok::annot_module_begin);
-    ModuleEntryIncludes.insert({Tok.getLocation().getRawEncoding(),
-                                (Module *)Tok.getAnnotationValue()});
+    ModuleEntryIncludes.insert(
+        {Tok.getLocation(), (Module *)Tok.getAnnotationValue()});
   }
 private:
   void FileChanged(SourceLocation Loc, FileChangeReason Reason,
@@ -88,12 +88,11 @@
                      SrcMgr::CharacteristicKind FileType,
                      StringRef Extra = StringRef());
   void WriteImplicitModuleImport(const Module *Mod);
-  void OutputContentUpTo(const MemoryBuffer &FromFile,
-                         unsigned &WriteFrom, unsigned WriteTo,
-                         StringRef EOL, int &lines,
+  void OutputContentUpTo(const MemoryBufferRef &FromFile, unsigned &WriteFrom,
+                         unsigned WriteTo, StringRef EOL, int &lines,
                          bool EnsureNewline);
   void CommentOutDirective(Lexer &DirectivesLex, const Token &StartToken,
-                           const MemoryBuffer &FromFile, StringRef EOL,
+                           const MemoryBufferRef &FromFile, StringRef EOL,
                            unsigned &NextToWrite, int &Lines);
   const IncludedFile *FindIncludeAtLocation(SourceLocation Loc) const;
   const Module *FindModuleAtLocation(SourceLocation Loc) const;
@@ -109,8 +108,7 @@
                                      bool ShowLineMarkers,
                                      bool UseLineDirectives)
     : PP(PP), SM(PP.getSourceManager()), OS(OS), MainEOL("\n"),
-      PredefinesBuffer(nullptr), ShowLineMarkers(ShowLineMarkers),
-      UseLineDirectives(UseLineDirectives),
+      ShowLineMarkers(ShowLineMarkers), UseLineDirectives(UseLineDirectives),
       LastInclusionLocation(SourceLocation()) {}
 
 /// Write appropriate line information as either #line directives or GNU line
@@ -164,7 +162,7 @@
     return;
   FileID Id = FullSourceLoc(Loc, SM).getFileID();
   auto P = FileIncludes.insert(
-      std::make_pair(LastInclusionLocation.getRawEncoding(),
+      std::make_pair(LastInclusionLocation,
                      IncludedFile(Id, NewFileType, PP.GetCurDirLookup())));
   (void)P;
   assert(P.second && "Unexpected revisitation of the same include directive");
@@ -199,8 +197,7 @@
                                            const Module *Imported,
                                            SrcMgr::CharacteristicKind FileType){
   if (Imported) {
-    auto P = ModuleIncludes.insert(
-        std::make_pair(HashLoc.getRawEncoding(), Imported));
+    auto P = ModuleIncludes.insert(std::make_pair(HashLoc, Imported));
     (void)P;
     assert(P.second && "Unexpected revisitation of the same include directive");
   } else
@@ -209,8 +206,7 @@
 
 void InclusionRewriter::If(SourceLocation Loc, SourceRange ConditionRange,
                            ConditionValueKind ConditionValue) {
-  auto P = IfConditions.insert(
-      std::make_pair(Loc.getRawEncoding(), ConditionValue == CVK_True));
+  auto P = IfConditions.insert(std::make_pair(Loc, ConditionValue == CVK_True));
   (void)P;
   assert(P.second && "Unexpected revisitation of the same if directive");
 }
@@ -218,8 +214,7 @@
 void InclusionRewriter::Elif(SourceLocation Loc, SourceRange ConditionRange,
                              ConditionValueKind ConditionValue,
                              SourceLocation IfLoc) {
-  auto P = IfConditions.insert(
-      std::make_pair(Loc.getRawEncoding(), ConditionValue == CVK_True));
+  auto P = IfConditions.insert(std::make_pair(Loc, ConditionValue == CVK_True));
   (void)P;
   assert(P.second && "Unexpected revisitation of the same elif directive");
 }
@@ -228,7 +223,7 @@
 /// an inclusion directive) in the map of inclusion information, FileChanges.
 const InclusionRewriter::IncludedFile *
 InclusionRewriter::FindIncludeAtLocation(SourceLocation Loc) const {
-  const auto I = FileIncludes.find(Loc.getRawEncoding());
+  const auto I = FileIncludes.find(Loc);
   if (I != FileIncludes.end())
     return &I->second;
   return nullptr;
@@ -238,7 +233,7 @@
 /// an inclusion directive) in the map of module inclusion information.
 const Module *
 InclusionRewriter::FindModuleAtLocation(SourceLocation Loc) const {
-  const auto I = ModuleIncludes.find(Loc.getRawEncoding());
+  const auto I = ModuleIncludes.find(Loc);
   if (I != ModuleIncludes.end())
     return I->second;
   return nullptr;
@@ -248,14 +243,14 @@
 /// an inclusion directive) in the map of module entry information.
 const Module *
 InclusionRewriter::FindEnteredModule(SourceLocation Loc) const {
-  const auto I = ModuleEntryIncludes.find(Loc.getRawEncoding());
+  const auto I = ModuleEntryIncludes.find(Loc);
   if (I != ModuleEntryIncludes.end())
     return I->second;
   return nullptr;
 }
 
 bool InclusionRewriter::IsIfAtLocationTrue(SourceLocation Loc) const {
-  const auto I = IfConditions.find(Loc.getRawEncoding());
+  const auto I = IfConditions.find(Loc);
   if (I != IfConditions.end())
     return I->second;
   return false;
@@ -263,7 +258,7 @@
 
 /// Detect the likely line ending style of \p FromFile by examining the first
 /// newline found within it.
-static StringRef DetectEOL(const MemoryBuffer &FromFile) {
+static StringRef DetectEOL(const MemoryBufferRef &FromFile) {
   // Detect what line endings the file uses, so that added content does not mix
   // the style. We need to check for "\r\n" first because "\n\r" will match
   // "\r\n\r\n".
@@ -278,23 +273,22 @@
 }
 
 void InclusionRewriter::detectMainFileEOL() {
-  bool Invalid;
-  const MemoryBuffer &FromFile = *SM.getBuffer(SM.getMainFileID(), &Invalid);
-  assert(!Invalid);
-  if (Invalid)
+  Optional<MemoryBufferRef> FromFile = *SM.getBufferOrNone(SM.getMainFileID());
+  assert(FromFile);
+  if (!FromFile)
     return; // Should never happen, but whatever.
-  MainEOL = DetectEOL(FromFile);
+  MainEOL = DetectEOL(*FromFile);
 }
 
 /// Writes out bytes from \p FromFile, starting at \p NextToWrite and ending at
 /// \p WriteTo - 1.
-void InclusionRewriter::OutputContentUpTo(const MemoryBuffer &FromFile,
+void InclusionRewriter::OutputContentUpTo(const MemoryBufferRef &FromFile,
                                           unsigned &WriteFrom, unsigned WriteTo,
                                           StringRef LocalEOL, int &Line,
                                           bool EnsureNewline) {
   if (WriteTo <= WriteFrom)
     return;
-  if (&FromFile == PredefinesBuffer) {
+  if (FromFile == PredefinesBuffer) {
     // Ignore the #defines of the predefines buffer.
     WriteFrom = WriteTo;
     return;
@@ -341,7 +335,7 @@
 /// through the \p FromFile buffer.
 void InclusionRewriter::CommentOutDirective(Lexer &DirectiveLex,
                                             const Token &StartToken,
-                                            const MemoryBuffer &FromFile,
+                                            const MemoryBufferRef &FromFile,
                                             StringRef LocalEOL,
                                             unsigned &NextToWrite, int &Line) {
   OutputContentUpTo(FromFile, NextToWrite,
@@ -351,7 +345,7 @@
   do {
     DirectiveLex.LexFromRawLexer(DirectiveToken);
   } while (!DirectiveToken.is(tok::eod) && DirectiveToken.isNot(tok::eof));
-  if (&FromFile == PredefinesBuffer) {
+  if (FromFile == PredefinesBuffer) {
     // OutputContentUpTo() would not output anything anyway.
     return;
   }
@@ -379,11 +373,15 @@
 void InclusionRewriter::Process(FileID FileId,
                                 SrcMgr::CharacteristicKind FileType,
                                 const DirectoryLookup *DirLookup) {
-  bool Invalid;
-  const MemoryBuffer &FromFile = *SM.getBuffer(FileId, &Invalid);
-  assert(!Invalid && "Attempting to process invalid inclusion");
+  MemoryBufferRef FromFile;
+  {
+    auto B = SM.getBufferOrNone(FileId);
+    assert(B && "Attempting to process invalid inclusion");
+    if (B)
+      FromFile = *B;
+  }
   StringRef FileName = FromFile.getBufferIdentifier();
-  Lexer RawLex(FileId, &FromFile, PP.getSourceManager(), PP.getLangOpts());
+  Lexer RawLex(FileId, FromFile, PP.getSourceManager(), PP.getLangOpts());
   RawLex.SetCommentRetentionState(false);
 
   StringRef LocalEOL = DetectEOL(FromFile);
@@ -560,7 +558,7 @@
     if (Tok.is(tok::annot_module_begin))
       Rewrite->handleModuleBegin(Tok);
   } while (Tok.isNot(tok::eof));
-  Rewrite->setPredefinesBuffer(SM.getBuffer(PP.getPredefinesFileID()));
+  Rewrite->setPredefinesBuffer(SM.getBufferOrFake(PP.getPredefinesFileID()));
   Rewrite->Process(PP.getPredefinesFileID(), SrcMgr::C_User, nullptr);
   Rewrite->Process(SM.getMainFileID(), SrcMgr::C_User, nullptr);
   OS->flush();