[clangd] Fix toHalfOpenFileRange where start/end endpoints are in different files due to #include

[sammccall]

Summary

https://github.com/clangd/clangd/issues/129

[sammccall]

kick phabricator to try to get it to mail

[sammccall]

clang-tools-extra/clangd/SourceCode.cpp

318

I dropped this function because it's just SM.isSpelledInSameFile

338

this is an unrelated bug fix: E1 < E2, std::min() etc don't do anything sensible for SourceLocation

358

This is half of the real fix: once we've hit the top of the macro tree, we should walk up the #include stack in case we find a common ancestor there.

443

this is one part of the main fix: the start and endpoint may be in different real files in the presence of #include

[SureYeaah]

LGTM

clang-tools-extra/clangd/SourceCode.cpp

284

nit: use this as a while condition?

while(Offset-- && Buf[Offset] != '\n')

[sammccall]

clang-tools-extra/clangd/SourceCode.cpp

284

Ooh, clever... a little too subtle for my taste though :-)
(In particular, allowing the variable to overflow but then ensuring we never read from it...)

clang-tools-extra/clangd/SourceCode.h

 /// 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);
 
+/// Returns the #include location through which IncludedFIle was loaded.
+/// Where SM.getIncludeLoc() returns the location of the *filename*, which may
+/// be in a macro, includeHashLoc() returns the location of the #.
+SourceLocation includeHashLoc(FileID IncludedFile, const SourceManager &SM);
+
 /// Returns true if the token at Loc is spelled in the source code.
 /// This is not the case for:
 ///   * symbols formed via macro concatenation, the spelling location will
 ///     be "<scratch space>"
 ///   * symbols controlled and defined by a compile command-line option
 ///     `-DName=foo`, the spelling location will be "<command line>".
 bool isSpelledInSource(SourceLocation Loc, const SourceManager &SM);
 

clang-tools-extra/clangd/SourceCode.cpp

bool halfOpenRangeContains(const SourceManager &Mgr, SourceRange R,
   return BeginFID == LFid && BeginOffset <= LOffset && LOffset < EndOffset;
 }
 
 bool halfOpenRangeTouches(const SourceManager &Mgr, SourceRange R,
                           SourceLocation L) {
   return L == R.getEnd() || halfOpenRangeContains(Mgr, R, L);
 }
 
+SourceLocation includeHashLoc(FileID IncludedFile, const SourceManager &SM) {
+  assert(SM.getLocForEndOfFile(IncludedFile).isFileID());
+  FileID IncludingFile;
+  unsigned Offset;
+  std::tie(IncludingFile, Offset) =
+      SM.getDecomposedExpansionLoc(SM.getIncludeLoc(IncludedFile));
+  bool Invalid = false;
+  llvm::StringRef Buf = SM.getBufferData(IncludingFile, &Invalid);
+  if (Invalid)
+    return SourceLocation();
+  // Now buf is "...\n#include <foo>\n..."
+  // and Offset points here:   ^
+  // Rewind to the preceding # on the line.
+  assert(Offset < Buf.size());
+  for (;; --Offset) {
+    if (Buf[Offset] == '#')
+      return SM.getComposedLoc(IncludingFile, Offset);
+    if (Buf[Offset] == '\n' || Offset == 0) // no hash, what's going on?
+      return SourceLocation();
+  }
+}
+
+
 static unsigned getTokenLengthAtLoc(SourceLocation Loc, const SourceManager &SM,
                                     const LangOptions &LangOpts) {
   Token TheTok;
   if (Lexer::getRawToken(Loc, TheTok, SM, LangOpts))
     return 0;
   // FIXME: Here we check whether the token at the location is a greatergreater
   // (>>) token and consider it as a single greater (>). This is to get it
   // working for templates but it isn't correct for the right shift operator. We
static SourceRange toTokenRange(CharSourceRange Range, const SourceManager &SM,
   return Range.getAsRange();
 }
 // Returns the union of two token ranges.
 // To find the maximum of the Ends of the ranges, we compare the location of the
 // last character of the token.
 static SourceRange unionTokenRange(SourceRange R1, SourceRange R2,
                                    const SourceManager &SM,
                                    const LangOptions &LangOpts) {
-  SourceLocation E1 = getLocForTokenEnd(R1.getEnd(), SM, LangOpts);
-  SourceLocation E2 = getLocForTokenEnd(R2.getEnd(), SM, LangOpts);
-  return SourceRange(std::min(R1.getBegin(), R2.getBegin()),
-                     E1 < E2 ? R2.getEnd() : R1.getEnd());
-}
-
-// Check if two locations have the same file id.
-static bool inSameFile(SourceLocation Loc1, SourceLocation Loc2,
-                       const SourceManager &SM) {
-  return SM.getFileID(Loc1) == SM.getFileID(Loc2);
-}
-
-// Find an expansion range (not necessarily immediate) the ends of which are in
-// the same file id.
-static SourceRange
-getExpansionTokenRangeInSameFile(SourceLocation Loc, const SourceManager &SM,
+  SourceLocation Begin =
+      SM.isBeforeInTranslationUnit(R1.getBegin(), R2.getBegin())
+          ? R1.getBegin()
+          : R2.getBegin();
+  SourceLocation End =
+      SM.isBeforeInTranslationUnit(getLocForTokenEnd(R1.getEnd(), SM, LangOpts),
+                                   getLocForTokenEnd(R2.getEnd(), SM, LangOpts))
+          ? R2.getEnd()
+          : R1.getEnd();
+  return SourceRange(Begin, End);
+}
+
+// Given a range whose endpoints may be in different expansions or files,
+// tries to find a range within a common file by following up the expansion and
+// include location in each.
+static SourceRange rangeInCommonFile(SourceRange R, const SourceManager &SM,
                                      const LangOptions &LangOpts) {
-  SourceRange ExpansionRange =
-      toTokenRange(SM.getImmediateExpansionRange(Loc), SM, LangOpts);
   // Fast path for most common cases.
-  if (inSameFile(ExpansionRange.getBegin(), ExpansionRange.getEnd(), SM))
-    return ExpansionRange;
+  if (SM.isWrittenInSameFile(R.getBegin(), R.getEnd()))
+    return R;
   // Record the stack of expansion locations for the beginning, keyed by FileID.
   llvm::DenseMap<FileID, SourceLocation> BeginExpansions;
-  for (SourceLocation Begin = ExpansionRange.getBegin(); Begin.isValid();
+  for (SourceLocation Begin = R.getBegin(); Begin.isValid();
        Begin = Begin.isFileID()
-                   ? SourceLocation()
+                   ? includeHashLoc(SM.getFileID(Begin), SM)
                    : SM.getImmediateExpansionRange(Begin).getBegin()) {
     BeginExpansions[SM.getFileID(Begin)] = Begin;
   }
   // Move up the stack of expansion locations for the end until we find the
   // location in BeginExpansions with that has the same file id.
-  for (SourceLocation End = ExpansionRange.getEnd(); End.isValid();
-       End = End.isFileID() ? SourceLocation()
+  for (SourceLocation End = R.getEnd(); End.isValid();
+       End = End.isFileID() ? includeHashLoc(SM.getFileID(End), SM)
                             : toTokenRange(SM.getImmediateExpansionRange(End),
                                            SM, LangOpts)
                                   .getEnd()) {
     auto It = BeginExpansions.find(SM.getFileID(End));
-    if (It != BeginExpansions.end())
+    if (It != BeginExpansions.end()) {
+      if (SM.getFileOffset(It->second) > SM.getFileOffset(End))
+        return SourceLocation();
       return {It->second, End};
     }
-  llvm_unreachable(
-      "We should able to find a common ancestor in the expansion tree.");
   }
+  return SourceRange();
+}
+
+// Find an expansion range (not necessarily immediate) the ends of which are in
+// the same file id.
+static SourceRange
+getExpansionTokenRangeInSameFile(SourceLocation Loc, const SourceManager &SM,
+                                 const LangOptions &LangOpts) {
+  return rangeInCommonFile(
+      toTokenRange(SM.getImmediateExpansionRange(Loc), SM, LangOpts), SM,
+      LangOpts);
+}
+
 // Returns the file range for a given Location as a Token Range
 // This is quite similar to getFileLoc in SourceManager as both use
 // getImmediateExpansionRange and getImmediateSpellingLoc (for macro IDs).
 // However:
 // - We want to maintain the full range information as we move from one file to
 //   the next. getFileLoc only uses the BeginLoc of getImmediateExpansionRange.
 // - We want to split '>>' tokens as the lexer parses the '>>' in nested
 //   template instantiations as a '>>' instead of two '>'s.
 // There is also getExpansionRange but it simply calls
 // getImmediateExpansionRange on the begin and ends separately which is wrong.
 static SourceRange getTokenFileRange(SourceLocation Loc,
                                      const SourceManager &SM,
                                      const LangOptions &LangOpts) {
   SourceRange FileRange = Loc;
   while (!FileRange.getBegin().isFileID()) {
     if (SM.isMacroArgExpansion(FileRange.getBegin())) {
       FileRange = unionTokenRange(
           SM.getImmediateSpellingLoc(FileRange.getBegin()),
           SM.getImmediateSpellingLoc(FileRange.getEnd()), SM, LangOpts);
-      assert(inSameFile(FileRange.getBegin(), FileRange.getEnd(), SM));
+      assert(SM.isWrittenInSameFile(FileRange.getBegin(), FileRange.getEnd()));
     } else {
       SourceRange ExpansionRangeForBegin =
           getExpansionTokenRangeInSameFile(FileRange.getBegin(), SM, LangOpts);
       SourceRange ExpansionRangeForEnd =
           getExpansionTokenRangeInSameFile(FileRange.getEnd(), SM, LangOpts);
-      assert(inSameFile(ExpansionRangeForBegin.getBegin(),
-                        ExpansionRangeForEnd.getBegin(), SM) &&
+      if (ExpansionRangeForBegin.isInvalid() ||
+          ExpansionRangeForEnd.isInvalid())
+        return SourceRange();
+      assert(SM.isWrittenInSameFile(ExpansionRangeForBegin.getBegin(),
+                                    ExpansionRangeForEnd.getBegin()) &&
              "Both Expansion ranges should be in same file.");
       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);
   if (!isValidFileRange(SM, R2))
     return llvm::None;
 
-  SourceRange Result = unionTokenRange(R1, R2, SM, LangOpts);
+  SourceRange Result =
+      rangeInCommonFile(unionTokenRange(R1, R2, SM, LangOpts), SM, LangOpts);
   unsigned TokLen = getTokenLengthAtLoc(Result.getEnd(), SM, LangOpts);
   // Convert from closed token range to half-open (char) range
   Result.setEnd(Result.getEnd().getLocWithOffset(TokLen));
   if (!isValidFileRange(SM, Result))
     return llvm::None;
 
   return Result;
 }

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

for (const SelectionTree::Node *N : allNodes(T))
         Complete.push_back(nodeRange(N, AST));
       else if (N->Selected == SelectionTree::Partial)
         Partial.push_back(nodeRange(N, AST));
     EXPECT_THAT(Complete, UnorderedElementsAreArray(Test.ranges("C"))) << C;
     EXPECT_THAT(Partial, UnorderedElementsAreArray(Test.ranges())) << C;
   }
 }
 
+TEST(SelectionTest, PathologicalPreprocessor) {
+  const char *Case = R"cpp(
+#define MACRO while(1)
+    void test() {
+#include "Expand.inc"
+        br^eak;
+    }
+  )cpp";
+  Annotations Test(Case);
+  auto TU = TestTU::withCode(Test.code());
+  TU.AdditionalFiles["Expand.inc"] = "MACRO\n";
+  auto AST = TU.build();
+  EXPECT_THAT(AST.getDiagnostics(), ::testing::IsEmpty());
+  auto T = makeSelectionTree(Case, AST);
+
+  EXPECT_EQ("BreakStmt", T.commonAncestor()->kind());
+  EXPECT_EQ("WhileStmt", T.commonAncestor()->Parent->kind());
+}
+
 TEST(SelectionTest, Implicit) {
   const char* Test = R"cpp(
     struct S { S(const char*); };
     int f(S);
     int x = f("^");
   )cpp";
   auto AST = TestTU::withCode(Annotations(Test).code()).build();
   auto T = makeSelectionTree(Test, AST);

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

 //===-- SourceCodeTests.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 "Annotations.h"
 #include "Context.h"
 #include "Protocol.h"
 #include "SourceCode.h"
 #include "TestTU.h"
 #include "clang/Basic/LangOptions.h"
+#include "clang/Basic/SourceLocation.h"
 #include "clang/Format/Format.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/raw_os_ostream.h"
+#include "llvm/Testing/Support/Annotations.h"
 #include "llvm/Testing/Support/Error.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 namespace clang {
 namespace clangd {
 namespace {
 
TEST(SourceCodeTests, HalfOpenFileRange) {
   CheckRange("a");
   CheckRange("b");
   CheckRange("c");
   CheckRange("d");
   CheckRange("e");
   CheckRange("f");
 }
 
+TEST(SourceCodeTests, HalfOpenFileRangePathologicalPreprocessor) {
+  const char *Case = R"cpp(
+#define MACRO while(1)
+    void test() {
+[[#include "Expand.inc"
+        br^eak]];
+    }
+  )cpp";
+  Annotations Test(Case);
+  auto TU = TestTU::withCode(Test.code());
+  TU.AdditionalFiles["Expand.inc"] = "MACRO\n";
+  auto AST = TU.build();
+
+  const auto &Func = cast<FunctionDecl>(findDecl(AST, "test"));
+  const auto &Body = cast<CompoundStmt>(Func.getBody());
+  const auto &Loop = cast<WhileStmt>(*Body->child_begin());
+  llvm::Optional<SourceRange> Range = toHalfOpenFileRange(
+      AST.getSourceManager(), AST.getASTContext().getLangOpts(),
+      Loop->getSourceRange());
+  ASSERT_TRUE(Range) << "Failed to get file range";
+  EXPECT_EQ(AST.getSourceManager().getFileOffset(Range->getBegin()),
+            Test.llvm::Annotations::range().Begin);
+  EXPECT_EQ(AST.getSourceManager().getFileOffset(Range->getEnd()),
+            Test.llvm::Annotations::range().End);
+}
+
+TEST(SourceCodeTests, IncludeHashLoc) {
+  const char *Case = R"cpp(
+$foo^#include "foo.inc"
+#define HEADER "bar.inc"
+  $bar^#  include HEADER
+  )cpp";
+  Annotations Test(Case);
+  auto TU = TestTU::withCode(Test.code());
+  TU.AdditionalFiles["foo.inc"] = "int foo;\n";
+  TU.AdditionalFiles["bar.inc"] = "int bar;\n";
+  auto AST = TU.build();
+  const auto& SM = AST.getSourceManager();
+
+  FileID Foo = SM.getFileID(findDecl(AST, "foo").getLocation());
+  EXPECT_EQ(SM.getFileOffset(includeHashLoc(Foo, SM)),
+            Test.llvm::Annotations::point("foo"));
+  FileID Bar = SM.getFileID(findDecl(AST, "bar").getLocation());
+  EXPECT_EQ(SM.getFileOffset(includeHashLoc(Bar, SM)),
+            Test.llvm::Annotations::point("foo"));
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang