[clangd] Print template arguments helper

Summary:
Prepares ground for printing template arguments as written in the
source code, part of re-landing rC356541 with D59599 applied.

Reviewers: ioeric, ilya-biryukov

Subscribers: mgorny, MaskRay, jkorous, arphaman, cfe-commits

Tags: #clang

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

llvm-svn: 358272

clang-tools-extra/clangd/AST.cpp

@@ -11,15 +11,37 @@
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/Decl.h"
 #include "clang/AST/DeclTemplate.h"
+#include "clang/AST/TemplateBase.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Index/USRGeneration.h"
+#include "llvm/ADT/Optional.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/ScopedPrinter.h"
+#include "llvm/Support/raw_ostream.h"
 
 namespace clang {
 namespace clangd {
 
+namespace {
+llvm::Optional<llvm::ArrayRef<TemplateArgumentLoc>>
+getTemplateSpecializationArgLocs(const NamedDecl &ND) {
+  if (auto *Func = llvm::dyn_cast<FunctionDecl>(&ND)) {
+    if (const ASTTemplateArgumentListInfo *Args =
+            Func->getTemplateSpecializationArgsAsWritten())
+      return Args->arguments();
+  } else if (auto *Cls =
+                 llvm::dyn_cast<ClassTemplatePartialSpecializationDecl>(&ND)) {
+    if (auto *Args = Cls->getTemplateArgsAsWritten())
+      return Args->arguments();
+  } else if (auto *Var = llvm::dyn_cast<VarTemplateSpecializationDecl>(&ND))
+    return Var->getTemplateArgsInfo().arguments();
+  // We return None for ClassTemplateSpecializationDecls because it does not
+  // contain TemplateArgumentLoc information.
+  return llvm::None;
+}
+} // namespace
+
 // Returns true if the complete name of decl \p D is spelled in the source code.
 // This is not the case for:
 //   * symbols formed via macro concatenation, the spelling location will
@@ -65,17 +87,6 @@ std::string printQualifiedName(const NamedDecl &ND) {
   return QName;
 }
 
-static const TemplateArgumentList *
-getTemplateSpecializationArgs(const NamedDecl &ND) {
-  if (auto *Func = llvm::dyn_cast<FunctionDecl>(&ND))
-    return Func->getTemplateSpecializationArgs();
-  if (auto *Cls = llvm::dyn_cast<ClassTemplateSpecializationDecl>(&ND))
-    return &Cls->getTemplateInstantiationArgs();
-  if (auto *Var = llvm::dyn_cast<VarTemplateSpecializationDecl>(&ND))
-    return &Var->getTemplateInstantiationArgs();
-  return nullptr;
-}
-
 std::string printName(const ASTContext &Ctx, const NamedDecl &ND) {
   std::string Name;
   llvm::raw_string_ostream Out(Name);
@@ -90,9 +101,7 @@ std::string printName(const ASTContext &Ctx, const NamedDecl &ND) {
   }
   ND.getDeclName().print(Out, PP);
   if (!Out.str().empty()) {
-    // FIXME(ibiryukov): do not show args not explicitly written by the user.
+    Out << printTemplateSpecializationArgs(ND);
-    if (auto *ArgList = getTemplateSpecializationArgs(ND))
-      printTemplateArgumentList(Out, ArgList->asArray(), PP);
     return Out.str();
   }
   // The name was empty, so present an anonymous entity.
@@ -105,6 +114,35 @@ std::string printName(const ASTContext &Ctx, const NamedDecl &ND) {
   return "(anonymous)";
 }
 
+std::string printTemplateSpecializationArgs(const NamedDecl &ND) {
+  std::string TemplateArgs;
+  llvm::raw_string_ostream OS(TemplateArgs);
+  PrintingPolicy Policy(ND.getASTContext().getLangOpts());
+  if (llvm::Optional<llvm::ArrayRef<TemplateArgumentLoc>> Args =
+          getTemplateSpecializationArgLocs(ND)) {
+    printTemplateArgumentList(OS, *Args, Policy);
+  } else if (auto *Cls = llvm::dyn_cast<ClassTemplateSpecializationDecl>(&ND)) {
+    if (const TypeSourceInfo *TSI = Cls->getTypeAsWritten()) {
+      // ClassTemplateSpecializationDecls do not contain
+      // TemplateArgumentTypeLocs, they only have TemplateArgumentTypes. So we
+      // create a new argument location list from TypeSourceInfo.
+      auto STL = TSI->getTypeLoc().getAs<TemplateSpecializationTypeLoc>();
+      llvm::SmallVector<TemplateArgumentLoc, 8> ArgLocs;
+      ArgLocs.reserve(STL.getNumArgs());
+      for (unsigned I = 0; I < STL.getNumArgs(); ++I)
+        ArgLocs.push_back(STL.getArgLoc(I));
+      printTemplateArgumentList(OS, ArgLocs, Policy);
+    } else {
+      // FIXME: Fix cases when getTypeAsWritten returns null inside clang AST,
+      // e.g. friend decls. Currently we fallback to Template Arguments without
+      // location information.
+      printTemplateArgumentList(OS, Cls->getTemplateArgs().asArray(), Policy);
+    }
+  }
+  OS.flush();
+  return TemplateArgs;
+}
+
 std::string printNamespaceScope(const DeclContext &DC) {
   for (const auto *Ctx = &DC; Ctx != nullptr; Ctx = Ctx->getParent())
     if (const auto *NS = dyn_cast<NamespaceDecl>(Ctx))

clang-tools-extra/clangd/AST.h

@@ -47,6 +47,12 @@ std::string printNamespaceScope(const DeclContext &DC);
 /// "(anonymous struct)" or "(anonymous namespace)".
 std::string printName(const ASTContext &Ctx, const NamedDecl &ND);
 
+/// Prints template arguments of a decl as written in the source code, including
+/// enclosing '<' and '>', e.g for a partial specialization like: template
+/// <typename U> struct Foo<int, U> will return '<int, U>'. Returns an empty
+/// string if decl is not a template specialization.
+std::string printTemplateSpecializationArgs(const NamedDecl &ND);
+
 /// Gets the symbol ID for a declaration, if possible.
 llvm::Optional<SymbolID> getSymbolID(const Decl *D);
 

clang-tools-extra/unittests/clangd/CMakeLists.txt

@@ -13,6 +13,7 @@ include_directories(
 
 add_extra_unittest(ClangdTests
   Annotations.cpp
+  PrintASTTests.cpp
   BackgroundIndexTests.cpp
   CancellationTests.cpp
   ClangdTests.cpp

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

@@ -525,11 +525,9 @@ TEST_F(DocumentSymbolsTest, Template) {
           AllOf(WithName("Tmpl<double>"), WithKind(SymbolKind::Struct),
                 Children()),
           AllOf(WithName("funcTmpl"), Children()),
-          // FIXME(ibiryukov): template args should be <int> to match the code.
+          AllOf(WithName("funcTmpl<int>"), Children()),
-          AllOf(WithName("funcTmpl<int, double, float>"), Children()),
           AllOf(WithName("varTmpl"), Children()),
-          // FIXME(ibiryukov): template args should be <int> to match the code.
+          AllOf(WithName("varTmpl<int>"), Children())));
-          AllOf(WithName("varTmpl<int, double>"), Children())));
 }
 
 TEST_F(DocumentSymbolsTest, Namespaces) {

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

@@ -0,0 +1,100 @@
+//===--- PrintASTTests.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 "AST.h"
+#include "Annotations.h"
+#include "Protocol.h"
+#include "SourceCode.h"
+#include "TestTU.h"
+#include "clang/AST/RecursiveASTVisitor.h"
+#include "gmock/gmock.h"
+#include "gtest/gtest-param-test.h"
+#include "gtest/gtest.h"
+#include "gtest/internal/gtest-param-util-generated.h"
+
+namespace clang {
+namespace clangd {
+namespace {
+
+using testing::ElementsAreArray;
+
+struct Case {
+  const char *AnnotatedCode;
+  std::vector<const char *> Expected;
+};
+class ASTUtils : public testing::Test,
+                 public ::testing::WithParamInterface<Case> {};
+
+TEST_P(ASTUtils, PrintTemplateArgs) {
+  auto Pair = GetParam();
+  Annotations Test(Pair.AnnotatedCode);
+  auto AST = TestTU::withCode(Test.code()).build();
+  struct Visitor : RecursiveASTVisitor<Visitor> {
+    Visitor(std::vector<Position> Points) : Points(std::move(Points)) {}
+    bool VisitNamedDecl(const NamedDecl *ND) {
+      auto Pos = sourceLocToPosition(ND->getASTContext().getSourceManager(),
+                                     ND->getLocation());
+      if (Pos != Points[TemplateArgsAtPoints.size()])
+        return true;
+      TemplateArgsAtPoints.push_back(printTemplateSpecializationArgs(*ND));
+      return true;
+    }
+    std::vector<std::string> TemplateArgsAtPoints;
+    const std::vector<Position> Points;
+  };
+  Visitor V(Test.points());
+  V.TraverseDecl(AST.getASTContext().getTranslationUnitDecl());
+  EXPECT_THAT(V.TemplateArgsAtPoints, ElementsAreArray(Pair.Expected));
+}
+
+INSTANTIATE_TEST_CASE_P(ASTUtilsTests, ASTUtils,
+                        testing::ValuesIn(std::vector<Case>({
+                            {
+                                R"cpp(
+                                  template <class X> class Bar {};
+                                  template <> class ^Bar<double> {};)cpp",
+                                {"<double>"}},
+                            {
+                                R"cpp(
+                                  template <class X> class Bar {};
+                                  template <class T, class U,
+                                  template<typename> class Z, int Q>
+                                  struct Foo {};
+                                  template struct ^Foo<int, bool, Bar, 8>;
+                                  template <typename T>
+                                  struct ^Foo<T *, T, Bar, 3> {};)cpp",
+                                {"<int, bool, Bar, 8>", "<T *, T, Bar, 3>"}},
+                            {
+                                R"cpp(
+                                  template <int ...> void Foz() {};
+                                  template <> void ^Foz<3, 5, 8>() {};)cpp",
+                                {"<3, 5, 8>"}},
+                            {
+                                R"cpp(
+                                  template <class X> class Bar {};
+                                  template <template <class> class ...>
+                                  class Aux {};
+                                  template <> class ^Aux<Bar, Bar> {};
+                                  template <template <class> T>
+                                  class ^Aux<T, T> {};)cpp",
+                                {"<Bar, Bar>", "<T, T>"}},
+                            {
+                                R"cpp(
+                                  template <typename T> T var = 1234;
+                                  template <> int ^var<int> = 1;)cpp",
+                                {"<int>"}},
+                            {
+                                R"cpp(
+                                  template <typename T> struct Foo;
+                                  struct Bar { friend class Foo<int>; };
+                                  template <> struct ^Foo<int> {};)cpp",
+                                {"<int>"}},
+                        })));
+} // namespace
+} // namespace clangd
+} // namespace clang

clang/lib/AST/TypePrinter.cpp

@@ -1632,6 +1632,19 @@ static const TemplateArgument &getArgument(const TemplateArgumentLoc &A) {
   return A.getArgument();
 }
 
+static void printArgument(const TemplateArgument &A, const PrintingPolicy &PP,
+                          llvm::raw_ostream &OS) {
+  A.print(PP, OS);
+}
+
+static void printArgument(const TemplateArgumentLoc &A,
+                          const PrintingPolicy &PP, llvm::raw_ostream &OS) {
+  const TemplateArgument::ArgKind &Kind = A.getArgument().getKind();
+  if (Kind == TemplateArgument::ArgKind::Type)
+    return A.getTypeSourceInfo()->getType().print(OS, PP);
+  return A.getArgument().print(PP, OS);
+}
+
 template<typename TA>
 static void printTo(raw_ostream &OS, ArrayRef<TA> Args,
                     const PrintingPolicy &Policy, bool SkipBrackets) {
@@ -1653,7 +1666,8 @@ static void printTo(raw_ostream &OS, ArrayRef<TA> Args,
     } else {
       if (!FirstArg)
         OS << Comma;
-      Argument.print(Policy, ArgOS);
+      // Tries to print the argument with location info if exists.
+      printArgument(Arg, Policy, ArgOS);
     }
     StringRef ArgString = ArgOS.str();