clangd: Extend reference search with constructor calls through forwarding

clang-tools-extra/clangd/AST.cpp

@@ -1040,5 +1040,26 @@ bool isExpandedFromParameterPack(const ParmVarDecl *D) {
   return getUnderlyingPackType(D) != nullptr;
 }
 
+bool isLikelyForwardingFunction(FunctionTemplateDecl *FT) {
+  const auto *FD = FT->getTemplatedDecl();
+  const auto NumParams = FD->getNumParams();
+  // Check whether its last parameter is a parameter pack...
+  if (NumParams > 0) {
+    const auto *LastParam = FD->getParamDecl(NumParams - 1);
+    if (const auto *PET = dyn_cast<PackExpansionType>(LastParam->getType())) {
+      // ... of the type T&&... or T...
+      const auto BaseType = PET->getPattern().getNonReferenceType();
+      if (const auto *TTPT =
+              dyn_cast<TemplateTypeParmType>(BaseType.getTypePtr())) {
+        // ... whose template parameter comes from the function directly
+        if (FT->getTemplateParameters()->getDepth() == TTPT->getDepth()) {
+          return true;
+        }
+      }
+    }
+  }
+  return false;
+}
+
 } // namespace clangd
 } // namespace clang

clang-tools-extra/clangd/AST.h

@@ -253,6 +253,10 @@ resolveForwardingParameters(const FunctionDecl *D, unsigned MaxDepth = 10);
 /// reference to one (e.g. `Args&...` or `Args&&...`).
 bool isExpandedFromParameterPack(const ParmVarDecl *D);
 
+/// Heuristic that checks if FT is forwarding a parameter pack to another
+/// function. (e.g. `make_unique`).
+bool isLikelyForwardingFunction(FunctionTemplateDecl *FT);
+
 } // namespace clangd
 } // namespace clang
 

clang-tools-extra/clangd/Preamble.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "Preamble.h"
+#include "AST.h"
 #include "CollectMacros.h"
 #include "Compiler.h"
 #include "Config.h"
@@ -166,27 +167,6 @@ class CppFilePreambleCallbacks : public PreambleCallbacks {
         collectPragmaMarksCallback(*SourceMgr, Marks));
   }
 
-  static bool isLikelyForwardingFunction(FunctionTemplateDecl *FT) {
-    const auto *FD = FT->getTemplatedDecl();
-    const auto NumParams = FD->getNumParams();
-    // Check whether its last parameter is a parameter pack...
-    if (NumParams > 0) {
-      const auto *LastParam = FD->getParamDecl(NumParams - 1);
-      if (const auto *PET = dyn_cast<PackExpansionType>(LastParam->getType())) {
-        // ... of the type T&&... or T...
-        const auto BaseType = PET->getPattern().getNonReferenceType();
-        if (const auto *TTPT =
-                dyn_cast<TemplateTypeParmType>(BaseType.getTypePtr())) {
-          // ... whose template parameter comes from the function directly
-          if (FT->getTemplateParameters()->getDepth() == TTPT->getDepth()) {
-            return true;
-          }
-        }
-      }
-    }
-    return false;
-  }
-
   bool shouldSkipFunctionBody(Decl *D) override {
     // Usually we don't need to look inside the bodies of header functions
     // to understand the program. However when forwarding function like

clang-tools-extra/clangd/XRefs.cpp

@@ -916,6 +916,47 @@ std::vector<DocumentLink> getDocumentLinks(ParsedAST &AST) {
 
 namespace {
 
+class ForwardingToConstructorVisitor
+    : public RecursiveASTVisitor<ForwardingToConstructorVisitor> {
+public:
+  ForwardingToConstructorVisitor(
+      llvm::DenseSet<const CXXConstructorDecl *> &TargetConstructors)
+      : Targets(TargetConstructors) {}
+
+  bool VisitCXXConstructExpr(CXXConstructExpr *E) {
+    if (auto *Callee = E->getConstructor()) {
+      if (Targets.contains(Callee)) {
+        ConstructorFound = true;
+      }
+    }
+    // It is enough to find 1 constructor
+    return !ConstructorFound;
+  }
+
+  // Output of this visitor
+  bool ConstructorFound = false;
+
+private:
+  llvm::DenseSet<const CXXConstructorDecl *> &Targets;
+};
+
+bool forwardsToConstructor(
+    const Decl *D,
+    llvm::DenseSet<const CXXConstructorDecl *> &TargetConstructors) {
+  if (!TargetConstructors.empty()) {
+    if (auto *FD = llvm::dyn_cast<clang::FunctionDecl>(D);
+        FD && FD->isTemplateInstantiation()) {
+      if (auto *PT = FD->getPrimaryTemplate();
+          PT && isLikelyForwardingFunction(PT)) {
+        ForwardingToConstructorVisitor Visitor{TargetConstructors};
+        Visitor.TraverseStmt(FD->getBody());
+        return Visitor.ConstructorFound;
+      }
+    }
+  }
+  return false;
+}
+
 /// Collects references to symbols within the main file.
 class ReferenceFinder : public index::IndexDataConsumer {
 public:
@@ -933,8 +974,12 @@ class ReferenceFinder : public index::IndexDataConsumer {
                   const llvm::ArrayRef<const NamedDecl *> Targets,
                   bool PerToken)
       : PerToken(PerToken), AST(AST) {
-    for (const NamedDecl *ND : Targets)
+    for (const NamedDecl *ND : Targets) {
       TargetDecls.insert(ND->getCanonicalDecl());
+      if (auto *Constructor = llvm::dyn_cast<clang::CXXConstructorDecl>(ND)) {
+        TargetConstructors.insert(Constructor);
+      }
+    }
   }
 
   std::vector<Reference> take() && {
@@ -960,8 +1005,10 @@ class ReferenceFinder : public index::IndexDataConsumer {
                        llvm::ArrayRef<index::SymbolRelation> Relations,
                        SourceLocation Loc,
                        index::IndexDataConsumer::ASTNodeInfo ASTNode) override {
-    if (!TargetDecls.contains(D->getCanonicalDecl()))
+    if (!TargetDecls.contains(D->getCanonicalDecl()) &&
+        !forwardsToConstructor(ASTNode.OrigD, TargetConstructors)) {
       return true;
+    }
     const SourceManager &SM = AST.getSourceManager();
     if (!isInsideMainFile(Loc, SM))
       return true;
@@ -1002,6 +1049,8 @@ class ReferenceFinder : public index::IndexDataConsumer {
   std::vector<Reference> References;
   const ParsedAST &AST;
   llvm::DenseSet<const Decl *> TargetDecls;
+  // Constructors need special handling since they can be hidden behind forwards
+  llvm::DenseSet<const CXXConstructorDecl *> TargetConstructors;
 };
 
 std::vector<ReferenceFinder::Reference>

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

@@ -29,6 +29,7 @@
 #include "clang/AST/DeclTemplate.h"
 #include "clang/AST/DeclarationName.h"
 #include "clang/AST/Expr.h"
+#include "clang/AST/RecursiveASTVisitor.h"
 #include "clang/Basic/FileEntry.h"
 #include "clang/Basic/LangOptions.h"
 #include "clang/Basic/SourceLocation.h"
@@ -576,6 +577,22 @@ SymbolCollector::getRefContainer(const Decl *Enclosing,
   return Enclosing;
 }
 
+class ForwardingToConstructorVisitor
+    : public RecursiveASTVisitor<ForwardingToConstructorVisitor> {
+public:
+  ForwardingToConstructorVisitor() {}
+
+  bool VisitCXXConstructExpr(CXXConstructExpr *E) {
+    if (auto *Callee = E->getConstructor()) {
+      Constructors.push_back(Callee);
+    }
+    return true;
+  }
+
+  // Output of this visitor
+  std::vector<CXXConstructorDecl *> Constructors{};
+};
+
 // Always return true to continue indexing.
 bool SymbolCollector::handleDeclOccurrence(
     const Decl *D, index::SymbolRoleSet Roles,
@@ -669,6 +686,21 @@ bool SymbolCollector::handleDeclOccurrence(
       addRef(ID, SymbolRef{FileLoc, FID, Roles, index::getSymbolInfo(ND).Kind,
                            getRefContainer(ASTNode.Parent, Opts),
                            isSpelled(FileLoc, *ND)});
+      // Also collect indirect constructor calls like `make_unique`
+      if (auto *FD = llvm::dyn_cast<clang::FunctionDecl>(ASTNode.OrigD);
+          FD && FD->isTemplateInstantiation()) {
+        if (auto *PT = FD->getPrimaryTemplate();
+            PT && isLikelyForwardingFunction(PT)) {
+          ForwardingToConstructorVisitor Visitor{};
+          Visitor.TraverseStmt(FD->getBody());
+          for (auto *Constructor : Visitor.Constructors) {
+            addRef(getSymbolIDCached(Constructor),
+                   SymbolRef{FileLoc, FID, Roles, index::getSymbolInfo(ND).Kind,
+                             getRefContainer(ASTNode.Parent, Opts),
+                             isSpelled(FileLoc, *ND)});
+          }
+        }
+      }
     }
   }
   // Don't continue indexing if this is a mere reference.

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

@@ -5,14 +5,15 @@
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
-#include "Annotations.h"
 #include "AST.h"
+#include "Annotations.h"
 #include "ParsedAST.h"
 #include "Protocol.h"
 #include "SourceCode.h"
 #include "SyncAPI.h"
 #include "TestFS.h"
 #include "TestTU.h"
+#include "TestWorkspace.h"
 #include "XRefs.h"
 #include "index/MemIndex.h"
 #include "clang/AST/Decl.h"
@@ -2713,6 +2714,62 @@ TEST(FindReferences, NoQueryForLocalSymbols) {
   }
 }
 
+TEST(FindReferences, ForwardingInAST) {
+  Annotations Main(R"cpp(
+    namespace std {
+    template <class T> T &&forward(T &t);
+    template <class T, class... Args> T *make_unique(Args &&...args) {
+      return new T(std::forward<Args>(args)...);
+    }
+    }
+
+    struct Test {
+      $Constructor[[T^est]](){}
+    };
+
+    int main() {
+      auto a = std::$Caller[[make_unique]]<Test>();
+    }
+  )cpp");
+  TestTU TU;
+  TU.Code = std::string(Main.code());
+  auto AST = TU.build();
+
+  EXPECT_THAT(findReferences(AST, Main.point(), 0).References,
+              ElementsAre(rangeIs(Main.range("Constructor")),
+                          rangeIs(Main.range("Caller"))));
+}
+
+TEST(FindReferences, ForwardingInIndex) {
+  Annotations Header(R"cpp(
+    namespace std {
+    template <class T> T &&forward(T &t);
+    template <class T, class... Args> T *make_unique(Args &&...args) {
+      return new T(std::forward<Args>(args)...);
+    }
+    }
+    struct Test {
+      [[T^est]](){}
+    };
+  )cpp");
+  Annotations Main(R"cpp(
+    #include "header.hpp"
+    int main() {
+      auto a = std::[[make_unique]]<Test>();
+    }
+  )cpp");
+  TestWorkspace TW;
+  TW.addMainFile("header.hpp", Header.code());
+  TW.addMainFile("main.cpp", Main.code());
+  auto AST = TW.openFile("header.hpp").value();
+  auto Index = TW.index();
+
+  EXPECT_THAT(findReferences(AST, Header.point(), 0, Index.get(),
+                             /*AddContext*/ true)
+                  .References,
+              ElementsAre(rangeIs(Header.range()), rangeIs(Main.range())));
+}
+
 TEST(GetNonLocalDeclRefs, All) {
   struct Case {
     llvm::StringRef AnnotatedCode;