[clangd] cleanup: unify the implemenation of checking a location is inside main file.

Summary: We have variant implementations in the codebase, this patch unifies them.

Reviewers: ilya-biryukov, kadircet

Subscribers: MaskRay, jkorous, arphaman, cfe-commits

Tags: #clang

Differential Revision: https://reviews.llvm.org/D64915

llvm-svn: 366541

Author: hokein

Diff for: clang-tools-extra/clangd/ClangdUnit.cpp

@@ -68,7 +68,7 @@ class DeclTrackingASTConsumer : public ASTConsumer {
   bool HandleTopLevelDecl(DeclGroupRef DG) override {
     for (Decl *D : DG) {
       auto &SM = D->getASTContext().getSourceManager();
-      if (!SM.isWrittenInMainFile(SM.getExpansionLoc(D->getLocation())))
+      if (!isInsideMainFile(D->getLocation(), SM))
         continue;
 
       // ObjCMethodDecl are not actually top-level decls.
@@ -355,8 +355,7 @@ ParsedAST::build(std::unique_ptr<CompilerInvocation> CI,
           // those might take us into a preamble file as well.
           bool IsInsideMainFile =
               Info.hasSourceManager() &&
-              Info.getSourceManager().isWrittenInMainFile(
-                  Info.getSourceManager().getFileLoc(Info.getLocation()));
+              isInsideMainFile(Info.getLocation(), Info.getSourceManager());
           if (IsInsideMainFile && tidy::ShouldSuppressDiagnostic(
                                       DiagLevel, Info, *CTContext,
                                       /* CheckMacroExpansion = */ false)) {

Diff for: clang-tools-extra/clangd/Diagnostics.cpp

@@ -140,15 +140,11 @@ void adjustDiagFromHeader(Diag &D, const clang::Diagnostic &Info,
   D.Message = llvm::Twine("in included file: ", D.Message).str();
 }
 
-bool isInsideMainFile(const SourceLocation Loc, const SourceManager &M) {
-  return Loc.isValid() && M.isWrittenInMainFile(M.getFileLoc(Loc));
-}
-
 bool isInsideMainFile(const clang::Diagnostic &D) {
   if (!D.hasSourceManager())
     return false;
 
-  return isInsideMainFile(D.getLocation(), D.getSourceManager());
+  return clangd::isInsideMainFile(D.getLocation(), D.getSourceManager());
 }
 
 bool isNote(DiagnosticsEngine::Level L) {

Diff for: clang-tools-extra/clangd/Headers.cpp

@@ -35,7 +35,7 @@ class RecordHeaders : public PPCallbacks {
                           llvm::StringRef /*RelativePath*/,
                           const Module * /*Imported*/,
                           SrcMgr::CharacteristicKind FileKind) override {
-    if (SM.isWrittenInMainFile(HashLoc)) {
+    if (isInsideMainFile(HashLoc, SM)) {
       Out->MainFileIncludes.emplace_back();
       auto &Inc = Out->MainFileIncludes.back();
       Inc.R = halfOpenToRange(SM, FilenameRange);

Diff for: clang-tools-extra/clangd/IncludeFixer.cpp

@@ -318,7 +318,7 @@ class IncludeFixer::UnresolvedNameRecorder : public ExternalSemaSource {
     assert(SemaPtr && "Sema must have been set.");
     if (SemaPtr->isSFINAEContext())
       return TypoCorrection();
-    if (!SemaPtr->SourceMgr.isWrittenInMainFile(Typo.getLoc()))
+    if (!isInsideMainFile(Typo.getLoc(), SemaPtr->SourceMgr))
       return clang::TypoCorrection();
 
     // This is not done lazily because `SS` can get out of scope and it's

Diff for: clang-tools-extra/clangd/Quality.cpp

@@ -9,6 +9,7 @@
 #include "Quality.h"
 #include "AST.h"
 #include "FileDistance.h"
+#include "SourceCode.h"
 #include "URI.h"
 #include "index/Symbol.h"
 #include "clang/AST/ASTContext.h"
@@ -42,8 +43,7 @@ static bool isReserved(llvm::StringRef Name) {
 static bool hasDeclInMainFile(const Decl &D) {
   auto &SourceMgr = D.getASTContext().getSourceManager();
   for (auto *Redecl : D.redecls()) {
-    auto Loc = SourceMgr.getSpellingLoc(Redecl->getLocation());
-    if (SourceMgr.isWrittenInMainFile(Loc))
+    if (isInsideMainFile(Redecl->getLocation(), SourceMgr))
       return true;
   }
   return false;
@@ -53,8 +53,7 @@ static bool hasUsingDeclInMainFile(const CodeCompletionResult &R) {
   const auto &Context = R.Declaration->getASTContext();
   const auto &SourceMgr = Context.getSourceManager();
   if (R.ShadowDecl) {
-    const auto Loc = SourceMgr.getExpansionLoc(R.ShadowDecl->getLocation());
-    if (SourceMgr.isWrittenInMainFile(Loc))
+    if (isInsideMainFile(R.ShadowDecl->getLocation(), SourceMgr))
       return true;
   }
   return false;

Diff for: clang-tools-extra/clangd/SourceCode.cpp

@@ -328,6 +328,10 @@ static SourceRange getTokenFileRange(SourceLocation Loc,
   return FileRange;
 }
 
+bool isInsideMainFile(SourceLocation Loc, const SourceManager &SM) {
+  return Loc.isValid() && SM.isWrittenInMainFile(SM.getExpansionLoc(Loc));
+}
+
 llvm::Optional<SourceRange> toHalfOpenFileRange(const SourceManager &SM,
                                                 const LangOptions &LangOpts,
                                                 SourceRange R) {

Diff for: clang-tools-extra/clangd/SourceCode.h

@@ -75,6 +75,14 @@ llvm::Optional<Range> getTokenRange(const SourceManager &SM,
 llvm::Expected<SourceLocation> sourceLocationInMainFile(const SourceManager &SM,
                                                         Position P);
 
+/// Returns true iff \p Loc is inside the main file. This function handles
+/// file & macro locations. For macro locations, returns iff the macro is being
+/// expanded inside the main file.
+///
+/// The function is usually used to check whether a declaration is inside the
+/// the main file.
+bool isInsideMainFile(SourceLocation Loc, const SourceManager &SM);
+
 /// Turns a token range into a half-open range and checks its correctness.
 /// The resulting range will have only valid source location on both sides, both
 /// of which are file locations.

Diff for: clang-tools-extra/clangd/XRefs.cpp

@@ -406,7 +406,7 @@ class ReferenceFinder : public index::IndexDataConsumer {
     assert(D->isCanonicalDecl() && "expect D to be a canonical declaration");
     const SourceManager &SM = AST.getSourceManager();
     Loc = SM.getFileLoc(Loc);
-    if (SM.isWrittenInMainFile(Loc) && CanonicalTargets.count(D))
+    if (isInsideMainFile(Loc, SM) && CanonicalTargets.count(D))
       References.push_back({D, Loc, Roles});
     return true;
   }

Diff for: clang-tools-extra/clangd/index/SymbolCollector.cpp

@@ -185,8 +185,7 @@ getTokenLocation(SourceLocation TokLoc, const SourceManager &SM,
 bool isPreferredDeclaration(const NamedDecl &ND, index::SymbolRoleSet Roles) {
   const auto &SM = ND.getASTContext().getSourceManager();
   return (Roles & static_cast<unsigned>(index::SymbolRole::Definition)) &&
-         isa<TagDecl>(&ND) &&
-         !SM.isWrittenInMainFile(SM.getExpansionLoc(ND.getLocation()));
+         isa<TagDecl>(&ND) && !isInsideMainFile(ND.getLocation(), SM);
 }
 
 RefKind toRefKind(index::SymbolRoleSet Roles) {

Diff for: clang-tools-extra/clangd/refactor/Rename.cpp

@@ -104,8 +104,7 @@ llvm::Optional<ReasonToReject> renamableWithinFile(const Decl &RenameDecl,
   auto &ASTCtx = RenameDecl.getASTContext();
   const auto &SM = ASTCtx.getSourceManager();
   bool MainFileIsHeader = ASTCtx.getLangOpts().IsHeaderFile;
-  bool DeclaredInMainFile =
-      SM.isWrittenInMainFile(SM.getExpansionLoc(RenameDecl.getLocation()));
+  bool DeclaredInMainFile = isInsideMainFile(RenameDecl.getBeginLoc(), SM);
 
   // If the symbol is declared in the main file (which is not a header), we
   // rename it.

Diff for: clang-tools-extra/clangd/unittests/SourceCodeTests.cpp

@@ -422,6 +422,36 @@ TEST(SourceCodeTests, GetMacros) {
   EXPECT_THAT(*Result, MacroName("MACRO"));
 }
 
+TEST(SourceCodeTests, IsInsideMainFile){
+  TestTU TU;
+  TU.HeaderCode = R"cpp(
+    #define DEFINE_CLASS(X) class X {};
+    #define DEFINE_YY DEFINE_CLASS(YY)
+
+    class Header1 {};
+    DEFINE_CLASS(Header2)
+    class Header {};
+  )cpp";
+  TU.Code = R"cpp(
+    class Main1 {};
+    DEFINE_CLASS(Main2)
+    DEFINE_YY
+    class Main {};
+  )cpp";
+  TU.ExtraArgs.push_back("-DHeader=Header3");
+  TU.ExtraArgs.push_back("-DMain=Main3");
+  auto AST = TU.build();
+  const auto& SM = AST.getSourceManager();
+  auto DeclLoc = [&AST](llvm::StringRef Name) {
+    return findDecl(AST, Name).getLocation();
+  };
+  for (const auto *HeaderDecl : {"Header1", "Header2", "Header3"})
+    EXPECT_FALSE(isInsideMainFile(DeclLoc(HeaderDecl), SM));
+
+  for (const auto *MainDecl : {"Main1", "Main2", "Main3", "YY"})
+    EXPECT_TRUE(isInsideMainFile(DeclLoc(MainDecl), SM));
+}
+
 // Test for functions toHalfOpenFileRange and getHalfOpenFileRange
 TEST(SourceCodeTests, HalfOpenFileRange) {
   // Each marked range should be the file range of the decl with the same name

Diff for: clang-tools-extra/clangd/unittests/SymbolCollectorTests.cpp

@@ -124,8 +124,7 @@ class ShouldCollectSymbolTest : public ::testing::Test {
     const NamedDecl &ND =
         Qualified ? findDecl(*AST, Name) : findUnqualifiedDecl(*AST, Name);
     const SourceManager &SM = AST->getSourceManager();
-    bool MainFile =
-        SM.isWrittenInMainFile(SM.getExpansionLoc(ND.getBeginLoc()));
+    bool MainFile = isInsideMainFile(ND.getBeginLoc(), SM);
     return SymbolCollector::shouldCollectSymbol(
         ND, AST->getASTContext(), SymbolCollector::Options(), MainFile);
   }