[libclang] Expose arguments of clang::annotate

This enables easy consumption of arbitrary data added
to this annotation in addition to the annotation category,
which was already exposed.

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

Author: fridtjof

clang/docs/ReleaseNotes.rst

@@ -245,6 +245,8 @@ clang-format
 libclang
 --------
 
+- Exposed arguments of ``clang::annotate``.
+
 Static Analyzer
 ---------------
 

clang/tools/libclang/CIndex.cpp

@@ -23,6 +23,7 @@
 #include "CursorVisitor.h"
 #include "clang-c/FatalErrorHandler.h"
 #include "clang/AST/Attr.h"
+#include "clang/AST/AttrVisitor.h"
 #include "clang/AST/DeclObjCCommon.h"
 #include "clang/AST/Expr.h"
 #include "clang/AST/ExprCXX.h"
@@ -575,6 +576,13 @@ bool CursorVisitor::VisitChildren(CXCursor Cursor) {
           A->getInterfaceLoc()->getTypeLoc().getBeginLoc(), TU));
   }
 
+  if (clang_isAttribute(Cursor.kind)) {
+    if (const Attr *A = getCursorAttr(Cursor))
+      return Visit(A);
+
+    return false;
+  }
+
   // If pointing inside a macro definition, check if the token is an identifier
   // that was ever defined as a macro. In such a case, create a "pseudo" macro
   // expansion cursor for that token.
@@ -2089,7 +2097,8 @@ class MemberRefVisit : public VisitorJob {
         (SourceLocation::UIntTy)(uintptr_t)data[1]);
   }
 };
-class EnqueueVisitor : public ConstStmtVisitor<EnqueueVisitor, void> {
+class EnqueueVisitor : public ConstStmtVisitor<EnqueueVisitor, void>,
+                       public ConstAttrVisitor<EnqueueVisitor, void> {
   friend class OMPClauseEnqueue;
   VisitorWorkList &WL;
   CXCursor Parent;
@@ -2231,6 +2240,9 @@ class EnqueueVisitor : public ConstStmtVisitor<EnqueueVisitor, void> {
   void VisitOMPTargetTeamsDistributeSimdDirective(
       const OMPTargetTeamsDistributeSimdDirective *D);
 
+  // Attributes
+  void VisitAnnotateAttr(const AnnotateAttr *A);
+
 private:
   void AddDeclarationNameInfo(const Stmt *S);
   void AddNestedNameSpecifierLoc(NestedNameSpecifierLoc Qualifier);
@@ -2242,6 +2254,7 @@ class EnqueueVisitor : public ConstStmtVisitor<EnqueueVisitor, void> {
   void AddTypeLoc(TypeSourceInfo *TI);
   void EnqueueChildren(const Stmt *S);
   void EnqueueChildren(const OMPClause *S);
+  void EnqueueChildren(const AnnotateAttr *A);
 };
 } // namespace
 
@@ -2736,6 +2749,20 @@ void EnqueueVisitor::EnqueueChildren(const OMPClause *S) {
   VisitorWorkList::iterator I = WL.begin() + size, E = WL.end();
   std::reverse(I, E);
 }
+
+void EnqueueVisitor::EnqueueChildren(const AnnotateAttr *A) {
+  unsigned size = WL.size();
+  for (const Expr *Arg : A->args()) {
+    VisitStmt(Arg);
+  }
+  if (size == WL.size())
+    return;
+  // Now reverse the entries we just added.  This will match the DFS
+  // ordering performed by the worklist.
+  VisitorWorkList::iterator I = WL.begin() + size, E = WL.end();
+  std::reverse(I, E);
+}
+
 void EnqueueVisitor::VisitAddrLabelExpr(const AddrLabelExpr *E) {
   WL.push_back(LabelRefVisit(E->getLabel(), E->getLabelLoc(), Parent));
 }
@@ -3008,7 +3035,7 @@ void EnqueueVisitor::VisitOpaqueValueExpr(const OpaqueValueExpr *E) {
   // If the opaque value has a source expression, just transparently
   // visit that.  This is useful for (e.g.) pseudo-object expressions.
   if (Expr *SourceExpr = E->getSourceExpr())
-    return Visit(SourceExpr);
+    return ConstStmtVisitor::Visit(SourceExpr);
 }
 void EnqueueVisitor::VisitLambdaExpr(const LambdaExpr *E) {
   AddStmt(E->getBody());
@@ -3028,7 +3055,7 @@ void EnqueueVisitor::VisitCXXParenListInitExpr(const CXXParenListInitExpr *E) {
 }
 void EnqueueVisitor::VisitPseudoObjectExpr(const PseudoObjectExpr *E) {
   // Treat the expression like its syntactic form.
-  Visit(E->getSyntacticForm());
+  ConstStmtVisitor::Visit(E->getSyntacticForm());
 }
 
 void EnqueueVisitor::VisitOMPExecutableDirective(
@@ -3338,9 +3365,28 @@ void EnqueueVisitor::VisitOMPTargetTeamsDistributeSimdDirective(
   VisitOMPLoopDirective(D);
 }
 
+void EnqueueVisitor::VisitAnnotateAttr(const AnnotateAttr *A) {
+  EnqueueChildren(A);
+}
+
 void CursorVisitor::EnqueueWorkList(VisitorWorkList &WL, const Stmt *S) {
   EnqueueVisitor(WL, MakeCXCursor(S, StmtParent, TU, RegionOfInterest))
-      .Visit(S);
+      .ConstStmtVisitor::Visit(S);
+}
+
+void CursorVisitor::EnqueueWorkList(VisitorWorkList &WL, const Attr *A) {
+  // Parent is the attribute itself when this is indirectly called from
+  // VisitChildren. Because we need to make a CXCursor for A, we need *its*
+  // parent.
+  auto AttrCursor = Parent;
+
+  // Get the attribute's parent as stored in
+  // cxcursor::MakeCXCursor(const Attr *A, const Decl *Parent, CXTranslationUnit
+  // TU)
+  const Decl *AttrParent = static_cast<const Decl *>(AttrCursor.data[1]);
+
+  EnqueueVisitor(WL, MakeCXCursor(A, AttrParent, TU))
+      .ConstAttrVisitor::Visit(A);
 }
 
 bool CursorVisitor::IsInRegionOfInterest(CXCursor C) {
@@ -3605,6 +3651,22 @@ bool CursorVisitor::Visit(const Stmt *S) {
   return result;
 }
 
+bool CursorVisitor::Visit(const Attr *A) {
+  VisitorWorkList *WL = nullptr;
+  if (!WorkListFreeList.empty()) {
+    WL = WorkListFreeList.back();
+    WL->clear();
+    WorkListFreeList.pop_back();
+  } else {
+    WL = new VisitorWorkList();
+    WorkListCache.push_back(WL);
+  }
+  EnqueueWorkList(*WL, A);
+  bool result = RunVisitorWorkList(*WL);
+  WorkListFreeList.push_back(WL);
+  return result;
+}
+
 namespace {
 typedef SmallVector<SourceRange, 4> RefNamePieces;
 RefNamePieces buildPieces(unsigned NameFlags, bool IsMemberRefExpr,

clang/tools/libclang/CursorVisitor.h

@@ -276,7 +276,9 @@ class CursorVisitor : public DeclVisitor<CursorVisitor, bool>,
   bool IsInRegionOfInterest(CXCursor C);
   bool RunVisitorWorkList(VisitorWorkList &WL);
   void EnqueueWorkList(VisitorWorkList &WL, const Stmt *S);
+  void EnqueueWorkList(VisitorWorkList &WL, const Attr *A);
   LLVM_ATTRIBUTE_NOINLINE bool Visit(const Stmt *S);
+  LLVM_ATTRIBUTE_NOINLINE bool Visit(const Attr *A);
 
 private:
   std::optional<bool> handleDeclForVisitation(const Decl *D);

clang/unittests/libclang/LibclangTest.cpp

@@ -1246,6 +1246,50 @@ static_assert(true, message);
   EXPECT_EQ(fromCXString(clang_getCursorSpelling(*staticAssertCsr)), "");
 }
 
+TEST_F(LibclangParseTest, ExposesAnnotateArgs) {
+  const char testSource[] = R"cpp(
+[[clang::annotate("category", 42)]]
+void func() {}
+)cpp";
+  std::string fileName = "main.cpp";
+  WriteFile(fileName, testSource);
+
+  const char *Args[] = {"-xc++"};
+  ClangTU = clang_parseTranslationUnit(Index, fileName.c_str(), Args, 1,
+                                       nullptr, 0, TUFlags);
+
+  int attrCount = 0;
+
+  Traverse(
+      [&attrCount](CXCursor cursor, CXCursor parent) -> CXChildVisitResult {
+        if (cursor.kind == CXCursor_AnnotateAttr) {
+          int childCount = 0;
+          clang_visitChildren(
+              cursor,
+              [](CXCursor child, CXCursor,
+                 CXClientData data) -> CXChildVisitResult {
+                int *pcount = static_cast<int *>(data);
+
+                // we only expect one argument here, so bail otherwise
+                EXPECT_EQ(*pcount, 0);
+
+                auto *result = clang_Cursor_Evaluate(child);
+                EXPECT_NE(result, nullptr);
+                EXPECT_EQ(clang_EvalResult_getAsInt(result), 42);
+                ++*pcount;
+
+                return CXChildVisit_Recurse;
+              },
+              &childCount);
+          attrCount++;
+          return CXChildVisit_Continue;
+        }
+        return CXChildVisit_Recurse;
+      });
+
+  EXPECT_EQ(attrCount, 1);
+}
+
 class LibclangRewriteTest : public LibclangParseTest {
 public:
   CXRewriter Rew = nullptr;