[AST] Improve traversal of concepts and concept requirements

- Do not traverse concept decl inside `AutoType`. We only traverse
  declaration and definitions, not references to a declaration.
- Do not visit implicit AST node the relevant traversal mode.
- Add traversal extension points for concept requirements.
- Renamed `TraverseConceptReference` to mark as helper to share
  the code. Having an extension point there seems confusing given that
  there are many concept refences in the AST that do not call the
  helper. Those are `AutoType`, `AutoTypeLoc` and constraint requirements.

Only clangd code requires an update.

There are no use-cases for concept requirement traversals yet, but
I added them in the earlier version of the patch and decided to keep
them for completeness.

Reviewed By: sammccall

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

Author: ilya-biryukov

clang-tools-extra/clangd/Selection.cpp

@@ -10,6 +10,7 @@
 #include "AST.h"
 #include "support/Logger.h"
 #include "support/Trace.h"
+#include "clang/AST/ASTConcept.h"
 #include "clang/AST/ASTTypeTraits.h"
 #include "clang/AST/Decl.h"
 #include "clang/AST/DeclCXX.h"
@@ -709,6 +710,15 @@ class SelectionVisitor : public RecursiveASTVisitor<SelectionVisitor> {
   bool TraversePseudoObjectExpr(PseudoObjectExpr *E) {
     return traverseNode(E, [&] { return TraverseStmt(E->getSyntacticForm()); });
   }
+  bool TraverseTypeConstraint(const TypeConstraint *C) {
+    if (auto *E = C->getImmediatelyDeclaredConstraint()) {
+      // Technically this expression is 'implicit' and not traversed by the RAV.
+      // However, the range is correct, so we visit expression to avoid adding
+      // an extra kind to 'DynTypeNode' that hold 'TypeConstraint'.
+      return TraverseStmt(E);
+    }
+    return Base::TraverseTypeConstraint(C);
+  }
 
 private:
   using Base = RecursiveASTVisitor<SelectionVisitor>;

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

@@ -2085,6 +2085,19 @@ TEST(FindReferences, ConceptsWithinAST) {
   checkFindRefs(Code);
 }
 
+TEST(FindReferences, ConceptReq) {
+  constexpr llvm::StringLiteral Code = R"cpp(
+    template <class T>
+    concept $def[[IsSmal^l]] = sizeof(T) <= 8;
+
+    template <class T>
+    concept IsSmallPtr = requires(T x) {
+      { *x } -> [[IsSmal^l]];
+    };
+  )cpp";
+  checkFindRefs(Code);
+}
+
 TEST(FindReferences, RequiresExprParameters) {
   constexpr llvm::StringLiteral Code = R"cpp(
     template <class T>

clang/include/clang/AST/RecursiveASTVisitor.h

@@ -13,18 +13,19 @@
 #ifndef LLVM_CLANG_AST_RECURSIVEASTVISITOR_H
 #define LLVM_CLANG_AST_RECURSIVEASTVISITOR_H
 
+#include "clang/AST/ASTConcept.h"
 #include "clang/AST/Attr.h"
 #include "clang/AST/Decl.h"
-#include "clang/AST/DeclarationName.h"
 #include "clang/AST/DeclBase.h"
 #include "clang/AST/DeclCXX.h"
 #include "clang/AST/DeclFriend.h"
 #include "clang/AST/DeclObjC.h"
 #include "clang/AST/DeclOpenMP.h"
 #include "clang/AST/DeclTemplate.h"
+#include "clang/AST/DeclarationName.h"
 #include "clang/AST/Expr.h"
-#include "clang/AST/ExprConcepts.h"
 #include "clang/AST/ExprCXX.h"
+#include "clang/AST/ExprConcepts.h"
 #include "clang/AST/ExprObjC.h"
 #include "clang/AST/ExprOpenMP.h"
 #include "clang/AST/LambdaCapture.h"
@@ -319,11 +320,6 @@ template <typename Derived> class RecursiveASTVisitor {
   bool TraverseSynOrSemInitListExpr(InitListExpr *S,
                                     DataRecursionQueue *Queue = nullptr);
 
-  /// Recursively visit a reference to a concept with potential arguments.
-  ///
-  /// \returns false if the visitation was terminated early, true otherwise.
-  bool TraverseConceptReference(const ConceptReference &C);
-
   /// Recursively visit an Objective-C protocol reference with location
   /// information.
   ///
@@ -475,11 +471,21 @@ template <typename Derived> class RecursiveASTVisitor {
   DEF_TRAVERSE_TMPL_INST(Function)
 #undef DEF_TRAVERSE_TMPL_INST
 
+  bool TraverseTypeConstraint(const TypeConstraint *C);
+
+  bool TraverseConceptRequirement(concepts::Requirement *R);
+  bool TraverseConceptTypeRequirement(concepts::TypeRequirement *R);
+  bool TraverseConceptExprRequirement(concepts::ExprRequirement *R);
+  bool TraverseConceptNestedRequirement(concepts::NestedRequirement *R);
+
   bool dataTraverseNode(Stmt *S, DataRecursionQueue *Queue);
 
 private:
   // These are helper methods used by more than one Traverse* method.
   bool TraverseTemplateParameterListHelper(TemplateParameterList *TPL);
+  /// Traverses the qualifier, name and template arguments of a concept
+  /// reference.
+  bool TraverseConceptReferenceHelper(const ConceptReference &C);
 
   // Traverses template parameter lists of either a DeclaratorDecl or TagDecl.
   template <typename T>
@@ -511,6 +517,54 @@ template <typename Derived> class RecursiveASTVisitor {
   bool PostVisitStmt(Stmt *S);
 };
 
+template <typename Derived>
+bool RecursiveASTVisitor<Derived>::TraverseTypeConstraint(
+    const TypeConstraint *C) {
+  if (!getDerived().shouldVisitImplicitCode()) {
+    TRY_TO(TraverseConceptReferenceHelper(*C));
+    return true;
+  }
+  if (Expr *IDC = C->getImmediatelyDeclaredConstraint()) {
+    TRY_TO(TraverseStmt(IDC));
+  } else {
+    // Avoid traversing the ConceptReference in the TypeConstraint
+    // if we have an immediately-declared-constraint, otherwise
+    // we'll end up visiting the concept and the arguments in
+    // the TC twice.
+    TRY_TO(TraverseConceptReferenceHelper(*C));
+  }
+  return true;
+}
+
+template <typename Derived>
+bool RecursiveASTVisitor<Derived>::TraverseConceptRequirement(
+    concepts::Requirement *R) {
+  switch (R->getKind()) {
+  case concepts::Requirement::RK_Type:
+    return getDerived().TraverseConceptTypeRequirement(
+        cast<concepts::TypeRequirement>(R));
+  case concepts::Requirement::RK_Simple:
+  case concepts::Requirement::RK_Compound:
+    return getDerived().TraverseConceptExprRequirement(
+        cast<concepts::ExprRequirement>(R));
+  case concepts::Requirement::RK_Nested:
+    return getDerived().TraverseConceptNestedRequirement(
+        cast<concepts::NestedRequirement>(R));
+  }
+}
+
+template <typename Derived>
+bool RecursiveASTVisitor<Derived>::TraverseConceptReferenceHelper(
+    const ConceptReference &C) {
+  TRY_TO(TraverseNestedNameSpecifierLoc(C.getNestedNameSpecifierLoc()));
+  TRY_TO(TraverseDeclarationNameInfo(C.getConceptNameInfo()));
+  if (C.hasExplicitTemplateArgs())
+    TRY_TO(TraverseTemplateArgumentLocsHelper(
+        C.getTemplateArgsAsWritten()->getTemplateArgs(),
+        C.getTemplateArgsAsWritten()->NumTemplateArgs));
+  return true;
+}
+
 template <typename Derived>
 bool RecursiveASTVisitor<Derived>::dataTraverseNode(Stmt *S,
                                                     DataRecursionQueue *Queue) {
@@ -530,6 +584,40 @@ bool RecursiveASTVisitor<Derived>::dataTraverseNode(Stmt *S,
 
 #undef DISPATCH_STMT
 
+template <typename Derived>
+bool RecursiveASTVisitor<Derived>::TraverseConceptTypeRequirement(
+    concepts::TypeRequirement *R) {
+  if (R->isSubstitutionFailure())
+    return true;
+  return getDerived().TraverseTypeLoc(R->getType()->getTypeLoc());
+}
+
+template <typename Derived>
+bool RecursiveASTVisitor<Derived>::TraverseConceptExprRequirement(
+    concepts::ExprRequirement *R) {
+  if (!R->isExprSubstitutionFailure())
+    TRY_TO(TraverseStmt(R->getExpr()));
+  auto &RetReq = R->getReturnTypeRequirement();
+  if (RetReq.isTypeConstraint()) {
+    if (getDerived().shouldVisitImplicitCode()) {
+      TRY_TO(TraverseTemplateParameterListHelper(
+          RetReq.getTypeConstraintTemplateParameterList()));
+    } else {
+      // Template parameter list is implicit, visit constraint directly.
+      TRY_TO(TraverseTypeConstraint(RetReq.getTypeConstraint()));
+    }
+  }
+  return true;
+}
+
+template <typename Derived>
+bool RecursiveASTVisitor<Derived>::TraverseConceptNestedRequirement(
+    concepts::NestedRequirement *R) {
+  if (!R->isSubstitutionFailure())
+    return getDerived().TraverseStmt(R->getConstraintExpr());
+  return true;
+}
+
 template <typename Derived>
 bool RecursiveASTVisitor<Derived>::PostVisitStmt(Stmt *S) {
   // In pre-order traversal mode, each Traverse##STMT method is responsible for
@@ -1007,7 +1095,6 @@ DEF_TRAVERSE_TYPE(UnaryTransformType, {
 DEF_TRAVERSE_TYPE(AutoType, {
   TRY_TO(TraverseType(T->getDeducedType()));
   if (T->isConstrained()) {
-    TRY_TO(TraverseDecl(T->getTypeConstraintConcept()));
     TRY_TO(TraverseTemplateArguments(T->getArgs(), T->getNumArgs()));
   }
 })
@@ -1838,17 +1925,8 @@ DEF_TRAVERSE_DECL(BuiltinTemplateDecl, {
 template <typename Derived>
 bool RecursiveASTVisitor<Derived>::TraverseTemplateTypeParamDeclConstraints(
     const TemplateTypeParmDecl *D) {
-  if (const auto *TC = D->getTypeConstraint()) {
-    if (Expr *IDC = TC->getImmediatelyDeclaredConstraint()) {
-      TRY_TO(TraverseStmt(IDC));
-    } else {
-      // Avoid traversing the ConceptReference in the TypeCosntraint
-      // if we have an immediately-declared-constraint, otherwise
-      // we'll end up visiting the concept and the arguments in
-      // the TC twice.
-      TRY_TO(TraverseConceptReference(*TC));
-    }
-  }
+  if (const auto *TC = D->getTypeConstraint())
+    TRY_TO(TraverseTypeConstraint(TC));
   return true;
 }
 
@@ -2435,18 +2513,6 @@ bool RecursiveASTVisitor<Derived>::TraverseSynOrSemInitListExpr(
   return true;
 }
 
-template<typename Derived>
-bool RecursiveASTVisitor<Derived>::TraverseConceptReference(
-    const ConceptReference &C) {
-  TRY_TO(TraverseNestedNameSpecifierLoc(C.getNestedNameSpecifierLoc()));
-  TRY_TO(TraverseDeclarationNameInfo(C.getConceptNameInfo()));
-  if (C.hasExplicitTemplateArgs())
-    TRY_TO(TraverseTemplateArgumentLocsHelper(
-        C.getTemplateArgsAsWritten()->getTemplateArgs(),
-        C.getTemplateArgsAsWritten()->NumTemplateArgs));
-  return true;
-}
-
 template <typename Derived>
 bool RecursiveASTVisitor<Derived>::TraverseObjCProtocolLoc(
     ObjCProtocolLoc ProtocolLoc) {
@@ -2825,31 +2891,15 @@ DEF_TRAVERSE_STMT(CoyieldExpr, {
   }
 })
 
-DEF_TRAVERSE_STMT(ConceptSpecializationExpr, {
-  TRY_TO(TraverseConceptReference(*S));
-})
+DEF_TRAVERSE_STMT(ConceptSpecializationExpr,
+                  { TRY_TO(TraverseConceptReferenceHelper(*S)); })
 
 DEF_TRAVERSE_STMT(RequiresExpr, {
   TRY_TO(TraverseDecl(S->getBody()));
   for (ParmVarDecl *Parm : S->getLocalParameters())
     TRY_TO(TraverseDecl(Parm));
   for (concepts::Requirement *Req : S->getRequirements())
-    if (auto *TypeReq = dyn_cast<concepts::TypeRequirement>(Req)) {
-      if (!TypeReq->isSubstitutionFailure())
-        TRY_TO(TraverseTypeLoc(TypeReq->getType()->getTypeLoc()));
-    } else if (auto *ExprReq = dyn_cast<concepts::ExprRequirement>(Req)) {
-      if (!ExprReq->isExprSubstitutionFailure())
-        TRY_TO(TraverseStmt(ExprReq->getExpr()));
-      auto &RetReq = ExprReq->getReturnTypeRequirement();
-      if (RetReq.isTypeConstraint()) {
-        TRY_TO(TraverseStmt(
-            RetReq.getTypeConstraint()->getImmediatelyDeclaredConstraint()));
-      }
-    } else {
-      auto *NestedReq = cast<concepts::NestedRequirement>(Req);
-      if (!NestedReq->isSubstitutionFailure())
-        TRY_TO(TraverseStmt(NestedReq->getConstraintExpr()));
-    }
+    TRY_TO(TraverseConceptRequirement(Req));
 })
 
 // These literals (all of them) do not need any action.

clang/lib/Index/IndexBody.cpp

@@ -483,13 +483,10 @@ class BodyIndexer : public RecursiveASTVisitor<BodyIndexer> {
     return true;
   }
 
-  bool VisitTemplateTypeParmDecl(TemplateTypeParmDecl *D) {
-    // This handles references in return type requirements of RequiresExpr.
-    // E.g. `requires (T x) { {*x} -> ConceptRef }`
-    if (auto *C = D->getTypeConstraint())
-      IndexCtx.handleReference(C->getNamedConcept(), C->getConceptNameLoc(),
-                               Parent, ParentDC);
-    return true;
+  bool TraverseTypeConstraint(const TypeConstraint *C) {
+    IndexCtx.handleReference(C->getNamedConcept(), C->getConceptNameLoc(),
+                             Parent, ParentDC);
+    return RecursiveASTVisitor::TraverseTypeConstraint(C);
   }
 };