[clangd] Fix an assertion crash in ReferenceFinder.

[hokein]

Summary

The assertion is almost correct, but it fails on refs from non-preamble

clang-tools-extra/clangd/XRefs.cpp

 //===--- XRefs.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 "XRefs.h"
public:
   }
 
   bool
   handleDeclOccurrence(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");
-    if (!CanonicalTargets.count(D))
+    const SourceManager &SM = AST.getSourceManager();
+    if (!CanonicalTargets.count(D) || !isInsideMainFile(Loc, SM))
       return true;
     const auto &TB = AST.getTokens();
-    const SourceManager &SM = AST.getSourceManager();
     Loc = SM.getFileLoc(Loc);
-    // We are only traversing decls *inside* the main file, so this should hold.
-    assert(isInsideMainFile(Loc, SM));
     if (const auto *Tok = TB.spelledTokenAt(Loc))
       References.push_back({*Tok, Roles});
     return true;
   }
 
 private:
   llvm::SmallSet<const Decl *, 4> CanonicalTargets;
   std::vector<Reference> References;

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

 //===-- XRefsTests.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"
for (const char *Test : Tests) {
     for (const auto &R : T.ranges())
       ExpectedLocations.push_back(RangeIs(R));
     EXPECT_THAT(findReferences(AST, T.point(), 0).References,
                 ElementsAreArray(ExpectedLocations))
         << Test;
   }
 }
 
+TEST(FindReferences, MainFileReferencesOnly) {
+  llvm::StringRef Test =
+      R"cpp(
+        void test() {
+          int [[fo^o]] = 1;
+          // refs not from main file should not be included.
+          #include "foo.inc"
+        })cpp";
+
+  Annotations Code(Test);
+  auto TU = TestTU::withCode(Code.code());
+  TU.AdditionalFiles["foo.inc"] = R"cpp(
+      foo = 3;
+    )cpp";
+  auto AST = TU.build();
+
+  std::vector<Matcher<Location>> ExpectedLocations;
+  for (const auto &R : Code.ranges())
+    ExpectedLocations.push_back(RangeIs(R));
+  EXPECT_THAT(findReferences(AST, Code.point(), 0).References,
+              ElementsAreArray(ExpectedLocations))
+      << Test;
+}
+
 TEST(FindReferences, ExplicitSymbols) {
   const char *Tests[] = {
       R"cpp(
       struct Foo { Foo* [[self]]() const; };
       void f() {
         Foo foo;
         if (Foo* T = foo.[[^self]]()) {} // Foo member call expr.
       }