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

[hokein]

Summary

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

[ilya-biryukov]

Neat! Many thanks, that's a very useful cleanup.
A few suggestions from my side.

clang-tools-extra/clangd/SourceCode.cpp

332 ↗

(On Diff #210542)

NIT: why not getExpansionLoc? Its implementation and description seems simpler and (I believe) it produces locations inside the same files as getFileLoc (as both the macro arg and the macro name of the expansion are always coming from the same file)

clang-tools-extra/clangd/SourceCode.h

81 ↗

(On Diff #210542)

Could we add a description without mentioning other SourceManager functions here?
I believe this should be a fair description of what this function is doing:

For macro locations, returns iff the macro is being expanded inside the main file.

[hokein]

address review comments, and add unittest.

[hokein]

clang-tools-extra/clangd/SourceCode.cpp

332 ↗

(On Diff #210542)

yes, since we are checking the file (not the range), there is no difference between these two methods. Changed to getExpansionLoc

[ilya-biryukov]

LGTM

clang-tools-extra/clangd/unittests/SourceCodeTests.cpp

448 ↗

(On Diff #210569)

NIT: not sure if it's easy to get a hold of those, but ...
maybe add tests for locations inside the <builtin>, <command-line> and <scratch> buffers?

clang-tools-extra/trunk/clangd/ClangdUnit.cpp

 class DeclTrackingASTConsumer : public ASTConsumer {
 public:
   DeclTrackingASTConsumer(std::vector<Decl *> &TopLevelDecls)
       : TopLevelDecls(TopLevelDecls) {}
 
   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.
       if (isa<ObjCMethodDecl>(D))
         continue;
 
       TopLevelDecls.push_back(D);
     }
ASTDiags.setLevelAdjuster([&CTContext](DiagnosticsEngine::Level DiagLevel,
 
           // Check for suppression comment. Skip the check for diagnostics not
           // in the main file, because we don't want that function to query the
           // source buffer for preamble files. For the same reason, we ask
           // ShouldSuppressDiagnostic not to follow macro expansions, since
           // 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)) {
             return DiagnosticsEngine::Ignored;
           }
         }
       }
       return DiagLevel;

clang-tools-extra/trunk/clangd/Diagnostics.cpp

void adjustDiagFromHeader(Diag &D, const clang::Diagnostic &Info,
   N.File = FE->getName();
   N.Message = "error occurred here";
   N.Range = diagnosticRange(Info, LangOpts);
 
   // Update message to mention original file.
   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) {
   return L == DiagnosticsEngine::Note || L == DiagnosticsEngine::Remark;
 }
 
 llvm::StringRef diagLeveltoString(DiagnosticsEngine::Level Lvl) {
   switch (Lvl) {

clang-tools-extra/trunk/clangd/Headers.cpp

public:
   // in the main file are collected.
   void InclusionDirective(SourceLocation HashLoc, const Token & /*IncludeTok*/,
                           llvm::StringRef FileName, bool IsAngled,
                           CharSourceRange FilenameRange, const FileEntry *File,
                           llvm::StringRef /*SearchPath*/,
                           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);
       Inc.Written =
           (IsAngled ? "<" + FileName + ">" : "\"" + FileName + "\"").str();
       Inc.Resolved = File ? File->tryGetRealPathName() : "";
       Inc.HashOffset = SM.getFileOffset(HashLoc);
       Inc.FileKind = FileKind;

clang-tools-extra/trunk/clangd/IncludeFixer.cpp

public:
   TypoCorrection CorrectTypo(const DeclarationNameInfo &Typo, int LookupKind,
                              Scope *S, CXXScopeSpec *SS,
                              CorrectionCandidateCallback &CCC,
                              DeclContext *MemberContext, bool EnteringContext,
                              const ObjCObjectPointerType *OPT) override {
     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
     // relatively cheap.
     auto Extracted = extractUnresolvedNameCheaply(
         SemaPtr->SourceMgr, Typo, SS, SemaPtr->LangOpts,
         static_cast<Sema::LookupNameKind>(LookupKind) ==
             Sema::LookupNameKind::LookupNestedNameSpecifierName);

clang-tools-extra/trunk/clangd/Quality.cpp

 //===--- Quality.cpp ---------------------------------------------*- C++-*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "Quality.h"
 #include "AST.h"
 #include "FileDistance.h"
+#include "SourceCode.h"
 #include "URI.h"
 #include "index/Symbol.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/Decl.h"
 #include "clang/AST/DeclCXX.h"
 #include "clang/AST/DeclTemplate.h"
 #include "clang/AST/DeclVisitor.h"
 #include "clang/Basic/CharInfo.h"
static bool isReserved(llvm::StringRef Name) {
   // FIXME: Should we exclude _Bool and others recognized by the standard?
   return Name.size() >= 2 && Name[0] == '_' &&
          (isUppercase(Name[1]) || Name[1] == '_');
 }
 
 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;
 }
 
 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;
 }
 
 static SymbolQualitySignals::SymbolCategory categorize(const NamedDecl &ND) {
   if (const auto *FD = dyn_cast<FunctionDecl>(&ND)) {
     if (FD->isOverloadedOperator())

clang-tools-extra/trunk/clangd/SourceCode.h

llvm::Optional<Range> getTokenRange(const SourceManager &SM,
                                     const LangOptions &LangOpts,
                                     SourceLocation TokLoc);
 
 /// Return the file location, corresponding to \p P. Note that one should take
 /// care to avoid comparing the result with expansion locations.
 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.
 ///
 /// File locations always point to a particular offset in a file, i.e. they
 /// never refer to a location inside a macro expansion. Turning locations from
 /// macro expansions into file locations is ambiguous - one can use
 /// SourceManager::{getExpansion|getFile|getSpelling}Loc. This function

clang-tools-extra/trunk/clangd/SourceCode.cpp

if (SM.isMacroArgExpansion(FileRange.getBegin())) {
           SM.getImmediateExpansionRange(FileRange.getEnd()), SM, LangOpts);
       FileRange = unionTokenRange(ExpansionRangeForBegin, ExpansionRangeForEnd,
                                   SM, LangOpts);
     }
   }
   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) {
   SourceRange R1 = getTokenFileRange(R.getBegin(), SM, LangOpts);
   if (!isValidFileRange(SM, R1))
     return llvm::None;
 
   SourceRange R2 = getTokenFileRange(R.getEnd(), SM, LangOpts);

clang-tools-extra/trunk/clangd/XRefs.cpp

public:
   bool
   handleDeclOccurence(const Decl *D, index::SymbolRoleSet Roles,
                       llvm::ArrayRef<index::SymbolRelation> Relations,
                       SourceLocation Loc,
                       index::IndexDataConsumer::ASTNodeInfo ASTNode) override {
     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;
   }
 
 private:
   llvm::SmallSet<const Decl *, 4> CanonicalTargets;
   std::vector<Reference> References;
   const ASTContext &AST;

clang-tools-extra/trunk/clangd/index/SymbolCollector.cpp

 // in a header file, in which case clangd would prefer to use ND as a canonical
 // declaration.
 // FIXME: handle symbol types that are not TagDecl (e.g. functions), if using
 // the first seen declaration as canonical declaration is not a good enough
 // heuristic.
 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) {
   return static_cast<RefKind>(static_cast<unsigned>(RefKind::All) & Roles);
 }
 
 bool shouldIndexRelation(const index::SymbolRelation &R) {
   // We currently only index BaseOf relations, for type hierarchy subtypes.

clang-tools-extra/trunk/clangd/refactor/Rename.cpp

 llvm::Optional<ReasonToReject> renamableWithinFile(const Decl &RenameDecl,
                                                    StringRef MainFile,
                                                    const SymbolIndex *Index) {
   if (llvm::isa<NamespaceDecl>(&RenameDecl))
     return ReasonToReject::UnsupportedSymbol;
   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.
   if (DeclaredInMainFile && !MainFileIsHeader)
     return None;
 
   // Below are cases where the symbol is declared in the header.
   // If the symbol is function-local, we rename it.

clang-tools-extra/trunk/clangd/unittests/SourceCodeTests.cpp

TEST(SourceCodeTests, GetMacros) {
   auto AST = TU.build();
   auto Loc = getBeginningOfIdentifier(AST, Code.point(),
                                       AST.getSourceManager().getMainFileID());
   auto Result = locateMacroAt(Loc, AST.getPreprocessor());
   ASSERT_TRUE(Result);
   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
   // and each name should be unique.
   Annotations Test(R"cpp(
     #define FOO(X, Y) int Y = ++X
     #define BAR(X) X + 1
     #define ECHO(X) X

clang-tools-extra/trunk/clangd/unittests/SymbolCollectorTests.cpp

public:
   }
 
   // build() must have been called.
   bool shouldCollect(llvm::StringRef Name, bool Qualified = true) {
     assert(AST.hasValue());
     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);
   }
 
 protected:
   std::string HeaderName = "f.h";
   std::string FileName = "f.cpp";
   TestTU File;