diff --git a/clang/include/clang/Basic/SourceManager.h b/clang/include/clang/Basic/SourceManager.h --- a/clang/include/clang/Basic/SourceManager.h +++ b/clang/include/clang/Basic/SourceManager.h @@ -813,6 +813,11 @@ 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!"); diff --git a/clang/lib/Basic/SourceManager.cpp b/clang/lib/Basic/SourceManager.cpp --- a/clang/lib/Basic/SourceManager.cpp +++ b/clang/lib/Basic/SourceManager.cpp @@ -389,6 +389,14 @@ 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"); diff --git a/clang/lib/Lex/PPDirectives.cpp b/clang/lib/Lex/PPDirectives.cpp --- a/clang/lib/Lex/PPDirectives.cpp +++ b/clang/lib/Lex/PPDirectives.cpp @@ -2054,8 +2054,7 @@ // some directives (e.g. #endif of a header guard) will never be seen. // Since this will lead to confusing errors, avoid the inclusion. if (Action == Enter && 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}; diff --git a/clang/unittests/Basic/SourceManagerTest.cpp b/clang/unittests/Basic/SourceManagerTest.cpp --- a/clang/unittests/Basic/SourceManagerTest.cpp +++ b/clang/unittests/Basic/SourceManagerTest.cpp @@ -491,6 +491,30 @@ EXPECT_TRUE(SourceMgr.isBeforeInTranslationUnit(Macros[10].Loc, Macros[11].Loc)); } +TEST_F(SourceManagerTest, isMainFile) { + const char *Source = "int x;"; + + std::unique_ptr Buf = + llvm::MemoryBuffer::getMemBuffer(Source); + const FileEntry *SourceFile = + FileMgr.getVirtualFile("mainFile.cpp", Buf->getBufferSize(), 0); + SourceMgr.overrideFileContents(SourceFile, std::move(Buf)); + + std::unique_ptr 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