[clang-tools-extra] 2be4569 - [clangd] Fix isInsideMainFile to be aware of preamble.

[Sam McCall via cfe-commits]

Author: Sam McCall
Date: 2020-03-04T23:57:21+01:00
New Revision: 2be4569719357c55fa9680265bbdcfb798468e9a

URL: https://github.com/llvm/llvm-project/commit/2be4569719357c55fa9680265bbdcfb798468e9a
DIFF: https://github.com/llvm/llvm-project/commit/2be4569719357c55fa9680265bbdcfb798468e9a.diff

LOG: [clangd] Fix isInsideMainFile to be aware of preamble.

Reviewers: kadircet

Subscribers: ilya-biryukov, MaskRay, jkorous, arphaman, usaxena95, cfe-commits

Tags: #clang

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

Added: 
    

Modified: 
    clang-tools-extra/clangd/SourceCode.cpp
    clang-tools-extra/clangd/unittests/SourceCodeTests.cpp

Removed: 
    


################################################################################
diff  --git a/clang-tools-extra/clangd/SourceCode.cpp b/clang-tools-extra/clangd/SourceCode.cpp
index d18daa910d18..75154723bb72 100644
--- a/clang-tools-extra/clangd/SourceCode.cpp
+++ b/clang-tools-extra/clangd/SourceCode.cpp
@@ -420,7 +420,10 @@ static SourceRange getTokenFileRange(SourceLocation Loc,
 }
 
 bool isInsideMainFile(SourceLocation Loc, const SourceManager &SM) {
-  return Loc.isValid() && SM.isWrittenInMainFile(SM.getExpansionLoc(Loc));
+  if (!Loc.isValid())
+    return false;
+  FileID FID = SM.getFileID(SM.getExpansionLoc(Loc));
+  return FID == SM.getMainFileID() || FID == SM.getPreambleFileID();
 }
 
 llvm::Optional<SourceRange> toHalfOpenFileRange(const SourceManager &SM,

diff  --git a/clang-tools-extra/clangd/unittests/SourceCodeTests.cpp b/clang-tools-extra/clangd/unittests/SourceCodeTests.cpp
index 5c45675ab74d..49f6a73cec48 100644
--- a/clang-tools-extra/clangd/unittests/SourceCodeTests.cpp
+++ b/clang-tools-extra/clangd/unittests/SourceCodeTests.cpp
@@ -461,10 +461,12 @@ TEST(SourceCodeTests, IsInsideMainFile){
     class Header {};
   )cpp";
   TU.Code = R"cpp(
+    #define DEFINE_MAIN4 class Main4{};
     class Main1 {};
     DEFINE_CLASS(Main2)
     DEFINE_YY
     class Main {};
+    DEFINE_MAIN4
   )cpp";
   TU.ExtraArgs.push_back("-DHeader=Header3");
   TU.ExtraArgs.push_back("-DMain=Main3");
@@ -474,10 +476,13 @@ TEST(SourceCodeTests, IsInsideMainFile){
     return findDecl(AST, Name).getLocation();
   };
   for (const auto *HeaderDecl : {"Header1", "Header2", "Header3"})
-    EXPECT_FALSE(isInsideMainFile(DeclLoc(HeaderDecl), SM));
+    EXPECT_FALSE(isInsideMainFile(DeclLoc(HeaderDecl), SM)) << HeaderDecl;
 
-  for (const auto *MainDecl : {"Main1", "Main2", "Main3", "YY"})
-    EXPECT_TRUE(isInsideMainFile(DeclLoc(MainDecl), SM));
+  for (const auto *MainDecl : {"Main1", "Main2", "Main3", "Main4", "YY"})
+    EXPECT_TRUE(isInsideMainFile(DeclLoc(MainDecl), SM)) << MainDecl;
+
+  // Main4 is *spelled* in the preamble, but in the main-file part of it.
+  EXPECT_TRUE(isInsideMainFile(SM.getSpellingLoc(DeclLoc("Main4")), SM));
 }
 
 // Test for functions toHalfOpenFileRange and getHalfOpenFileRange