[clang][Preprocessor] Replace the slow translateFile call by a new, f…

clang/include/clang/Basic/SourceManager.h

@@ -811,6 +811,11 @@ class SourceManager : public RefCountedBase<SourceManager> {
     MainFileID = FID;
   }
 
+  /// Returns true when the given FileEntry corresponds to the main file.
+  ///
+  /// The main file should be set prior to calling this function.
+  bool isMainFile(FileEntryRef SourceFile);
+
   /// Set the file ID for the precompiled preamble.
   void setPreambleFileID(FileID Preamble) {
     assert(PreambleFileID.isInvalid() && "PreambleFileID already set!");

clang/lib/Basic/SourceManager.cpp

@@ -389,6 +389,14 @@ void SourceManager::clearIDTables() {
   createExpansionLoc(SourceLocation(), SourceLocation(), SourceLocation(), 1);
 }
 
+bool SourceManager::isMainFile(FileEntryRef SourceFile) {
+  assert(MainFileID.isValid() && "expected initialized SourceManager");
+  auto FE = getFileEntryRefForID(MainFileID);
+  if (!FE)
+    return false;
+  return FE->getUID() == SourceFile.getUID();
+}
+
 void SourceManager::initializeForReplay(const SourceManager &Old) {
   assert(MainFileID.isInvalid() && "expected uninitialized SourceManager");
 

clang/lib/Lex/PPDirectives.cpp

@@ -1937,8 +1937,7 @@ Preprocessor::ImportAction Preprocessor::HandleHeaderIncludeOrImport(
   // some directives (e.g. #endif of a header guard) will never be seen.
   // Since this will lead to confusing errors, avoid the inclusion.
   if (File && PreambleConditionalStack.isRecording() &&
-      SourceMgr.translateFile(&File->getFileEntry()) ==
-          SourceMgr.getMainFileID()) {
+      SourceMgr.isMainFile(*File)) {
     Diag(FilenameTok.getLocation(),
          diag::err_pp_including_mainfile_in_preamble);
     return {ImportAction::None};

clang/unittests/Basic/SourceManagerTest.cpp

@@ -485,6 +485,30 @@ TEST_F(SourceManagerTest, isBeforeInTranslationUnitWithMacroInInclude) {
   EXPECT_TRUE(SourceMgr.isBeforeInTranslationUnit(Macros[10].Loc, Macros[11].Loc));
 }
 
+TEST_F(SourceManagerTest, isMainFile) {
+  const char *Source = "int x;";
+
+  std::unique_ptr<llvm::MemoryBuffer> Buf =
+      llvm::MemoryBuffer::getMemBuffer(Source);
+  const FileEntry *SourceFile =
+      FileMgr.getVirtualFile("mainFile.cpp", Buf->getBufferSize(), 0);
+  SourceMgr.overrideFileContents(SourceFile, std::move(Buf));
+
+  std::unique_ptr<llvm::MemoryBuffer> Buf2 =
+      llvm::MemoryBuffer::getMemBuffer(Source);
+  const FileEntry *SecondFile =
+      FileMgr.getVirtualFile("file2.cpp", Buf2->getBufferSize(), 0);
+  SourceMgr.overrideFileContents(SecondFile, std::move(Buf2));
+
+  FileID MainFileID = SourceMgr.getOrCreateFileID(SourceFile, SrcMgr::C_User);
+  SourceMgr.setMainFileID(MainFileID);
+
+  EXPECT_TRUE(SourceMgr.isMainFile(FileEntryRef("mainFile.cpp", *SourceFile)));
+  EXPECT_TRUE(
+      SourceMgr.isMainFile(FileEntryRef("anotherName.cpp", *SourceFile)));
+  EXPECT_FALSE(SourceMgr.isMainFile(FileEntryRef("mainFile.cpp", *SecondFile)));
+}
+
 #endif
 
 } // anonymous namespace