[clangd] Store the containing symbol for refs

This is needed to implement call hierarchy.

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

5 files changed, 116 insertions, 42 deletions

diff --git a/clang-tools-extra/clangd/index/Ref.h b/clang-tools-extra/clangd/index/Ref.h
index 577db7525442..060d7c971c38 100644
--- a/clang-tools-extra/clangd/index/Ref.h
+++ b/clang-tools-extra/clangd/index/Ref.h
@@ -88,13 +88,19 @@ struct Ref {
   /// The source location where the symbol is named.
   SymbolLocation Location;
   RefKind Kind = RefKind::Unknown;
+  /// The ID of the symbol whose definition contains this reference.
+  /// For example, for a reference inside a function body, this would
+  /// be that function. For top-level definitions this isNull().
+  SymbolID Container;
 };
 
 inline bool operator<(const Ref &L, const Ref &R) {
-  return std::tie(L.Location, L.Kind) < std::tie(R.Location, R.Kind);
+  return std::tie(L.Location, L.Kind, L.Container) <
+         std::tie(R.Location, R.Kind, R.Container);
 }
 inline bool operator==(const Ref &L, const Ref &R) {
-  return std::tie(L.Location, L.Kind) == std::tie(R.Location, R.Kind);
+  return std::tie(L.Location, L.Kind, L.Container) ==
+         std::tie(R.Location, R.Kind, R.Container);
 }
 
 llvm::raw_ostream &operator<<(llvm::raw_ostream &, const Ref &);
diff --git a/clang-tools-extra/clangd/index/Serialization.cpp b/clang-tools-extra/clangd/index/Serialization.cpp
index e7f65f087b1c..6f65bf024c4e 100644
--- a/clang-tools-extra/clangd/index/Serialization.cpp
+++ b/clang-tools-extra/clangd/index/Serialization.cpp
@@ -345,6 +345,7 @@ void writeRefs(const SymbolID &ID, llvm::ArrayRef<Ref> Refs,
   for (const auto &Ref : Refs) {
     OS.write(static_cast<unsigned char>(Ref.Kind));
     writeLocation(Ref.Location, Strings, OS);
+    OS << Ref.Container.raw();
   }
 }
 
@@ -356,6 +357,7 @@ readRefs(Reader &Data, llvm::ArrayRef<llvm::StringRef> Strings) {
   for (auto &Ref : Result.second) {
     Ref.Kind = static_cast<RefKind>(Data.consume8());
     Ref.Location = readLocation(Data, Strings);
+    Ref.Container = Data.consumeID();
   }
   return Result;
 }
diff --git a/clang-tools-extra/clangd/index/SymbolCollector.cpp b/clang-tools-extra/clangd/index/SymbolCollector.cpp
index 511b5a2c6047..a7a6bd992f8b 100644
--- a/clang-tools-extra/clangd/index/SymbolCollector.cpp
+++ b/clang-tools-extra/clangd/index/SymbolCollector.cpp
@@ -344,7 +344,8 @@ bool SymbolCollector::handleDeclOccurrence(
       !isa<NamespaceDecl>(ND) &&
       (Opts.RefsInHeaders ||
        SM.getFileID(SM.getFileLoc(Loc)) == SM.getMainFileID()))
-    DeclRefs[ND].emplace_back(SM.getFileLoc(Loc), Roles);
+    DeclRefs[ND].push_back(
+        SymbolRef{SM.getFileLoc(Loc), Roles, ASTNode.Parent});
   // Don't continue indexing if this is a mere reference.
   if (IsOnlyRef)
     return true;
@@ -422,7 +423,8 @@ bool SymbolCollector::handleMacroOccurrence(const IdentifierInfo *Name,
   // Do not store references to main-file macros.
   if ((static_cast<unsigned>(Opts.RefFilter) & Roles) && !IsMainFileOnly &&
       (Opts.RefsInHeaders || SM.getFileID(SpellingLoc) == SM.getMainFileID()))
-    MacroRefs[ID].push_back({Loc, Roles});
+    // FIXME: Populate container information for macro references.
+    MacroRefs[ID].push_back({Loc, Roles, /*Container=*/nullptr});
 
   // Collect symbols.
   if (!Opts.CollectMacro)
@@ -579,24 +581,22 @@ void SymbolCollector::finish() {
     }
     return Found->second;
   };
-  auto CollectRef =
-      [&](SymbolID ID,
-          const std::pair<SourceLocation, index::SymbolRoleSet> &LocAndRole,
-          bool Spelled = false) {
-        auto FileID = SM.getFileID(LocAndRole.first);
-        // FIXME: use the result to filter out references.
-        shouldIndexFile(FileID);
-        if (auto FileURI = GetURI(FileID)) {
-          auto Range =
-              getTokenRange(LocAndRole.first, SM, ASTCtx->getLangOpts());
-          Ref R;
-          R.Location.Start = Range.first;
-          R.Location.End = Range.second;
-          R.Location.FileURI = FileURI->c_str();
-          R.Kind = toRefKind(LocAndRole.second, Spelled);
-          Refs.insert(ID, R);
-        }
-      };
+  auto CollectRef = [&](SymbolID ID, const SymbolRef &LocAndRole,
+                        bool Spelled = false) {
+    auto FileID = SM.getFileID(LocAndRole.Loc);
+    // FIXME: use the result to filter out references.
+    shouldIndexFile(FileID);
+    if (auto FileURI = GetURI(FileID)) {
+      auto Range = getTokenRange(LocAndRole.Loc, SM, ASTCtx->getLangOpts());
+      Ref R;
+      R.Location.Start = Range.first;
+      R.Location.End = Range.second;
+      R.Location.FileURI = FileURI->c_str();
+      R.Kind = toRefKind(LocAndRole.Roles, Spelled);
+      R.Container = getSymbolID(LocAndRole.Container);
+      Refs.insert(ID, R);
+    }
+  };
   // Populate Refs slab from MacroRefs.
   // FIXME: All MacroRefs are marked as Spelled now, but this should be checked.
   for (const auto &IDAndRefs : MacroRefs)
@@ -607,7 +607,7 @@ void SymbolCollector::finish() {
   for (auto &DeclAndRef : DeclRefs) {
     if (auto ID = getSymbolID(DeclAndRef.first)) {
       for (auto &LocAndRole : DeclAndRef.second) {
-        const auto FileID = SM.getFileID(LocAndRole.first);
+        const auto FileID = SM.getFileID(LocAndRole.Loc);
         // FIXME: It's better to use TokenBuffer by passing spelled tokens from
         // the caller of SymbolCollector.
         if (!FilesToTokensCache.count(FileID))
@@ -617,7 +617,7 @@ void SymbolCollector::finish() {
         // Check if the referenced symbol is spelled exactly the same way the
         // corresponding NamedDecl is. If it is, mark this reference as spelled.
         const auto *IdentifierToken =
-            spelledIdentifierTouching(LocAndRole.first, Tokens);
+            spelledIdentifierTouching(LocAndRole.Loc, Tokens);
         DeclarationName Name = DeclAndRef.first->getDeclName();
         const auto NameKind = Name.getNameKind();
         bool IsTargetKind = NameKind == DeclarationName::Identifier ||
diff --git a/clang-tools-extra/clangd/index/SymbolCollector.h b/clang-tools-extra/clangd/index/SymbolCollector.h
index a1b40a0dba79..1648a98b6220 100644
--- a/clang-tools-extra/clangd/index/SymbolCollector.h
+++ b/clang-tools-extra/clangd/index/SymbolCollector.h
@@ -156,7 +156,11 @@ private:
   std::shared_ptr<GlobalCodeCompletionAllocator> CompletionAllocator;
   std::unique_ptr<CodeCompletionTUInfo> CompletionTUInfo;
   Options Opts;
-  using SymbolRef = std::pair<SourceLocation, index::SymbolRoleSet>;
+  struct SymbolRef {
+    SourceLocation Loc;
+    index::SymbolRoleSet Roles;
+    const Decl *Container;
+  };
   // Symbols referenced from the current TU, flushed on finish().
   llvm::DenseSet<const NamedDecl *> ReferencedDecls;
   llvm::DenseSet<const IdentifierInfo *> ReferencedMacros;
diff --git a/clang-tools-extra/clangd/unittests/SymbolCollectorTests.cpp b/clang-tools-extra/clangd/unittests/SymbolCollectorTests.cpp
index 47071ac2da38..d0c4dd4497fa 100644
--- a/clang-tools-extra/clangd/unittests/SymbolCollectorTests.cpp
+++ b/clang-tools-extra/clangd/unittests/SymbolCollectorTests.cpp
@@ -70,21 +70,16 @@ MATCHER_P(IncludeHeader, P, "") {
 MATCHER_P2(IncludeHeaderWithRef, IncludeHeader, References, "") {
   return (arg.IncludeHeader == IncludeHeader) && (arg.References == References);
 }
+bool rangesMatch(const SymbolLocation &Loc, const Range &R) {
+  return std::make_tuple(Loc.Start.line(), Loc.Start.column(), Loc.End.line(),
+                         Loc.End.column()) ==
+         std::make_tuple(R.start.line, R.start.character, R.end.line,
+                         R.end.character);
+}
 MATCHER_P(DeclRange, Pos, "") {
-  return std::make_tuple(arg.CanonicalDeclaration.Start.line(),
-                         arg.CanonicalDeclaration.Start.column(),
-                         arg.CanonicalDeclaration.End.line(),
-                         arg.CanonicalDeclaration.End.column()) ==
-         std::make_tuple(Pos.start.line, Pos.start.character, Pos.end.line,
-                         Pos.end.character);
-}
-MATCHER_P(DefRange, Pos, "") {
-  return std::make_tuple(
-             arg.Definition.Start.line(), arg.Definition.Start.column(),
-             arg.Definition.End.line(), arg.Definition.End.column()) ==
-         std::make_tuple(Pos.start.line, Pos.start.character, Pos.end.line,
-                         Pos.end.character);
+  return rangesMatch(arg.CanonicalDeclaration, Pos);
 }
+MATCHER_P(DefRange, Pos, "") { return rangesMatch(arg.Definition, Pos); }
 MATCHER_P(RefCount, R, "") { return int(arg.References) == R; }
 MATCHER_P(ForCodeCompletion, IsIndexedForCodeCompletion, "") {
   return static_cast<bool>(arg.Flags & Symbol::IndexedForCodeCompletion) ==
@@ -100,10 +95,7 @@ MATCHER(VisibleOutsideFile, "") {
 MATCHER(RefRange, "") {
   const Ref &Pos = ::testing::get<0>(arg);
   const Range &Range = ::testing::get<1>(arg);
-  return std::make_tuple(Pos.Location.Start.line(), Pos.Location.Start.column(),
-                         Pos.Location.End.line(), Pos.Location.End.column()) ==
-         std::make_tuple(Range.start.line, Range.start.character,
-                         Range.end.line, Range.end.character);
+  return rangesMatch(Pos.Location, Range);
 }
 ::testing::Matcher<const std::vector<Ref> &>
 HaveRanges(const std::vector<Range> Ranges) {
@@ -738,6 +730,76 @@ TEST_F(SymbolCollectorTest, Refs) {
   EXPECT_THAT(Refs, Not(Contains(Pair(findSymbol(Symbols, "FUNC").ID, _))));
 }
 
+TEST_F(SymbolCollectorTest, RefContainers) {
+  Annotations Code(R"cpp(
+    int $toplevel1[[f1]](int);
+    void f2() {
+      (void) $ref1a[[f1]](1);
+      auto fptr = &$ref1b[[f1]];
+    }
+    int $toplevel2[[v1]] = $ref2[[f1]](2);
+    void f3(int arg = $ref3[[f1]](3));
+    struct S1 {
+      int $classscope1[[member1]] = $ref4[[f1]](4);
+      int $classscope2[[member2]] = 42;
+    };
+    constexpr int f4(int x) { return x + 1; }
+    template <int I = $ref5[[f4]](0)> struct S2 {};
+    S2<$ref6[[f4]](0)> v2;
+    S2<$ref7a[[f4]](0)> f5(S2<$ref7b[[f4]](0)>);
+    namespace N {
+      void $namespacescope1[[f6]]();
+      int $namespacescope2[[v3]];
+    }
+  )cpp");
+  CollectorOpts.RefFilter = RefKind::All;
+  CollectorOpts.CollectMainFileRefs = true;
+  runSymbolCollector("", Code.code());
+  auto FindRefWithRange = [&](Range R) -> Optional<Ref> {
+    for (auto &Entry : Refs) {
+      for (auto &Ref : Entry.second) {
+        if (rangesMatch(Ref.Location, R))
+          return Ref;
+      }
+    }
+    return llvm::None;
+  };
+  auto Container = [&](llvm::StringRef RangeName) {
+    auto Ref = FindRefWithRange(Code.range(RangeName));
+    EXPECT_TRUE(bool(Ref));
+    return Ref->Container;
+  };
+  EXPECT_EQ(Container("ref1a"),
+            findSymbol(Symbols, "f2").ID); // function body (call)
+  EXPECT_EQ(Container("ref1b"),
+            findSymbol(Symbols, "f2").ID); // function body (address-of)
+  EXPECT_EQ(Container("ref2"),
+            findSymbol(Symbols, "v1").ID); // variable initializer
+  EXPECT_EQ(Container("ref3"),
+            findSymbol(Symbols, "f3").ID); // function parameter default value
+  EXPECT_EQ(Container("ref4"),
+            findSymbol(Symbols, "S1::member1").ID); // member initializer
+  EXPECT_EQ(Container("ref5"),
+            findSymbol(Symbols, "S2").ID); // template parameter default value
+  EXPECT_EQ(Container("ref6"),
+            findSymbol(Symbols, "v2").ID); // type of variable
+  EXPECT_EQ(Container("ref7a"),
+            findSymbol(Symbols, "f5").ID); // return type of function
+  EXPECT_EQ(Container("ref7b"),
+            findSymbol(Symbols, "f5").ID); // parameter type of function
+
+  EXPECT_FALSE(Container("classscope1").isNull());
+  EXPECT_FALSE(Container("namespacescope1").isNull());
+
+  EXPECT_EQ(Container("toplevel1"), Container("toplevel2"));
+  EXPECT_EQ(Container("classscope1"), Container("classscope2"));
+  EXPECT_EQ(Container("namespacescope1"), Container("namespacescope2"));
+
+  EXPECT_NE(Container("toplevel1"), Container("namespacescope1"));
+  EXPECT_NE(Container("toplevel1"), Container("classscope1"));
+  EXPECT_NE(Container("classscope1"), Container("namespacescope1"));
+}
+
 TEST_F(SymbolCollectorTest, MacroRefInHeader) {
   Annotations Header(R"(
   #define $foo[[FOO]](X) (X + 1)