[clang][frontend] Add support for attribute plugins for statement att…

…ributes (#110334)

We already have support for declaration attributes; this is just a matter of extending the plugin infrastructure to cover one more case.

clang/docs/ClangPlugins.rst

@@ -92,11 +92,6 @@ The members of ``ParsedAttrInfo`` that a plugin attribute must define are:
    attribute, each of which consists of an attribute syntax and how the
    attribute name is spelled for that syntax. If the syntax allows a scope then
    the spelling must be "scope::attr" if a scope is present or "::attr" if not.
- * ``handleDeclAttribute``, which is the function that applies the attribute to
-   a declaration. It is responsible for checking that the attribute's arguments
-   are valid, and typically applies the attribute by adding an ``Attr`` to the
-   ``Decl``. It returns either ``AttributeApplied``, to indicate that the
-   attribute was successfully applied, or ``AttributeNotApplied`` if it wasn't.
 
 The members of ``ParsedAttrInfo`` that may need to be defined, depending on the
 attribute, are:
@@ -105,6 +100,18 @@ attribute, are:
    arguments to the attribute.
  * ``diagAppertainsToDecl``, which checks if the attribute has been used on the
    right kind of declaration and issues a diagnostic if not.
+ * ``handleDeclAttribute``, which is the function that applies the attribute to
+   a declaration. It is responsible for checking that the attribute's arguments
+   are valid, and typically applies the attribute by adding an ``Attr`` to the
+   ``Decl``. It returns either ``AttributeApplied``, to indicate that the
+   attribute was successfully applied, or ``AttributeNotApplied`` if it wasn't.
+ * ``diagAppertainsToStmt``, which checks if the attribute has been used on the
+   right kind of statement and issues a diagnostic if not.
+ * ``handleStmtAttribute``, which is the function that applies the attribute to
+   a statement. It is responsible for checking that the attribute's arguments
+   are valid, and typically applies the attribute by adding an ``Attr`` to the
+   ``Stmt``. It returns either ``AttributeApplied``, to indicate that the
+   attribute was successfully applied, or ``AttributeNotApplied`` if it wasn't.
  * ``diagLangOpts``, which checks if the attribute is permitted for the current
    language mode and issues a diagnostic if not.
  * ``existsInTarget``, which checks if the attribute is permitted for the given

clang/docs/ReleaseNotes.rst

@@ -250,6 +250,8 @@ Non-comprehensive list of changes in this release
 - The floating point comparison builtins (``__builtin_isgreater``,
   ``__builtin_isgreaterequal``, ``__builtin_isless``, etc.) and
   ``__builtin_signbit`` can now be used in constant expressions.
+- Plugins can now define custom attributes that apply to statements
+  as well as declarations.
 
 New Compiler Flags
 ------------------

clang/examples/Attribute/Attribute.cpp

@@ -94,6 +94,54 @@ struct ExampleAttrInfo : public ParsedAttrInfo {
     }
     return AttributeApplied;
   }
+
+  bool diagAppertainsToStmt(Sema &S, const ParsedAttr &Attr,
+                            const Stmt *St) const override {
+    // This attribute appertains to for loop statements only.
+    if (!isa<ForStmt>(St)) {
+      S.Diag(Attr.getLoc(), diag::warn_attribute_wrong_decl_type_str)
+          << Attr << Attr.isRegularKeywordAttribute() << "for loop statements";
+      return false;
+    }
+    return true;
+  }
+
+  AttrHandling handleStmtAttribute(Sema &S, Stmt *St, const ParsedAttr &Attr,
+                                   class Attr *&Result) const override {
+    // We make some rules here:
+    // 1. Only accept at most 3 arguments here.
+    // 2. The first argument must be a string literal if it exists.
+    if (Attr.getNumArgs() > 3) {
+      unsigned ID = S.getDiagnostics().getCustomDiagID(
+          DiagnosticsEngine::Error,
+          "'example' attribute only accepts at most three arguments");
+      S.Diag(Attr.getLoc(), ID);
+      return AttributeNotApplied;
+    }
+    // If there are arguments, the first argument should be a string literal.
+    if (Attr.getNumArgs() > 0) {
+      auto *Arg0 = Attr.getArgAsExpr(0);
+      StringLiteral *Literal =
+          dyn_cast<StringLiteral>(Arg0->IgnoreParenCasts());
+      if (!Literal) {
+        unsigned ID = S.getDiagnostics().getCustomDiagID(
+            DiagnosticsEngine::Error, "first argument to the 'example' "
+                                      "attribute must be a string literal");
+        S.Diag(Attr.getLoc(), ID);
+        return AttributeNotApplied;
+      }
+      SmallVector<Expr *, 16> ArgsBuf;
+      for (unsigned i = 0; i < Attr.getNumArgs(); i++) {
+        ArgsBuf.push_back(Attr.getArgAsExpr(i));
+      }
+      Result = AnnotateAttr::Create(S.Context, "example", ArgsBuf.data(),
+                                    ArgsBuf.size(), Attr.getRange());
+    } else {
+      Result = AnnotateAttr::Create(S.Context, "example", nullptr, 0,
+                                    Attr.getRange());
+    }
+    return AttributeApplied;
+  }
 };
 
 } // namespace

clang/include/clang/Basic/ParsedAttrInfo.h

@@ -24,6 +24,7 @@
 
 namespace clang {
 
+class Attr;
 class Decl;
 class LangOptions;
 class ParsedAttr;
@@ -154,6 +155,15 @@ struct ParsedAttrInfo {
                                            const ParsedAttr &Attr) const {
     return NotHandled;
   }
+  /// If this ParsedAttrInfo knows how to handle this ParsedAttr applied to this
+  /// Stmt then do so (referencing the resulting Attr in Result) and return
+  /// either AttributeApplied if it was applied or AttributeNotApplied if it
+  /// wasn't. Otherwise return NotHandled.
+  virtual AttrHandling handleStmtAttribute(Sema &S, Stmt *St,
+                                           const ParsedAttr &Attr,
+                                           class Attr *&Result) const {
+    return NotHandled;
+  }
 
   static const ParsedAttrInfo &get(const AttributeCommonInfo &A);
   static ArrayRef<const ParsedAttrInfo *> getAllBuiltin();

clang/lib/Sema/SemaStmtAttr.cpp

@@ -682,6 +682,10 @@ static Attr *ProcessStmtAttribute(Sema &S, Stmt *St, const ParsedAttr &A,
   case ParsedAttr::AT_Annotate:
     return S.CreateAnnotationAttr(A);
   default:
+    if (Attr *AT = nullptr; A.getInfo().handleStmtAttribute(S, St, A, AT) !=
+                            ParsedAttrInfo::NotHandled) {
+      return AT;
+    }
     // N.B., ClangAttrEmitter.cpp emits a diagnostic helper that ensures a
     // declaration attribute is not written on a statement, but this code is
     // needed for attributes in Attr.td that do not list any subjects.

clang/test/Frontend/plugin-attribute.cpp

@@ -4,19 +4,57 @@
 // REQUIRES: plugins, examples
 //--- good_attr.cpp
 // expected-no-diagnostics
-void fn1a() __attribute__((example)) {}
-[[example]] void fn1b() {}
-[[plugin::example]] void fn1c() {}
-void fn2() __attribute__((example("somestring", 1, 2.0))) {}
-// CHECK-COUNT-4: -AnnotateAttr 0x{{[0-9a-z]+}} {{<col:[0-9]+(, col:[0-9]+)?>}} "example"
+void fn1a() __attribute__((example)) {
+  __attribute__((example)) for (int i = 0; i < 9; ++i) {}
+}
+[[example]] void fn1b() {
+  [[example]] for (int i = 0; i < 9; ++i) {}
+}
+[[plugin::example]] void fn1c() {
+  [[plugin::example]] for (int i = 0; i < 9; ++i) {}
+}
+void fn2() __attribute__((example("somestring", 1, 2.0))) {
+  __attribute__((example("abc", 3, 4.0))) for (int i = 0; i < 9; ++i) {}
+}
+template <int N> void template_fn() __attribute__((example("template", N))) {
+  __attribute__((example("def", N + 1))) for (int i = 0; i < 9; ++i) {}
+}
+void fn3() { template_fn<5>(); }
+// CHECK: -AttributedStmt 0x{{[0-9a-z]+}} {{<line:[0-9]+:[0-9]+(, col:[0-9]+)?>}}
+// CHECK: -AnnotateAttr 0x{{[0-9a-z]+}} {{<col:[0-9]+(, col:[0-9]+)?>}} "example"
+// CHECK: -AnnotateAttr 0x{{[0-9a-z]+}} {{<line:[0-9]+:[0-9]+(, col:[0-9]+)?>}} "example"
+// CHECK: -AttributedStmt 0x{{[0-9a-z]+}} {{<line:[0-9]+:[0-9]+(, col:[0-9]+)?>}}
+// CHECK: -AnnotateAttr 0x{{[0-9a-z]+}} {{<col:[0-9]+(, col:[0-9]+)?>}} "example"
+// CHECK: -AnnotateAttr 0x{{[0-9a-z]+}} {{<line:[0-9]+:[0-9]+(, col:[0-9]+)?>}} "example"
+// CHECK: -AttributedStmt 0x{{[0-9a-z]+}} {{<line:[0-9]+:[0-9]+(, col:[0-9]+)?>}}
+// CHECK: -AnnotateAttr 0x{{[0-9a-z]+}} {{<col:[0-9]+(, col:[0-9]+)?>}} "example"
+// CHECK: -AnnotateAttr 0x{{[0-9a-z]+}} {{<line:[0-9]+:[0-9]+(, col:[0-9]+)?>}} "example"
+// CHECK: -AttributedStmt 0x{{[0-9a-z]+}} {{<line:[0-9]+:[0-9]+(, col:[0-9]+)?>}}
+// CHECK: -AnnotateAttr 0x{{[0-9a-z]+}} {{<col:[0-9]+(, col:[0-9]+)?>}} "example"
+// CHECK: -StringLiteral 0x{{[0-9a-z]+}} {{<col:[0-9]+(, col:[0-9]+)?>}} 'const char[{{[0-9]+}}]' lvalue "abc"
+// CHECK: -IntegerLiteral 0x{{[0-9a-z]+}} {{<col:[0-9]+(, col:[0-9]+)?>}} 'int' 3
+// CHECK: -FloatingLiteral 0x{{[0-9a-z]+}} {{<col:[0-9]+(, col:[0-9]+)?>}} 'double' 4.000000e+00
+// CHECK: -AnnotateAttr 0x{{[0-9a-z]+}} {{<line:[0-9]+:[0-9]+(, col:[0-9]+)?>}} "example"
 // CHECK: -StringLiteral 0x{{[0-9a-z]+}} {{<col:[0-9]+(, col:[0-9]+)?>}} 'const char[{{[0-9]+}}]' lvalue "somestring"
 // CHECK: -IntegerLiteral 0x{{[0-9a-z]+}} {{<col:[0-9]+(, col:[0-9]+)?>}} 'int' 1
 // CHECK: -FloatingLiteral 0x{{[0-9a-z]+}} {{<col:[0-9]+(, col:[0-9]+)?>}} 'double' 2.000000e+00
+// CHECK: -AnnotateAttr 0x{{[0-9a-z]+}} {{<col:[0-9]+(, col:[0-9]+)?>}} Implicit "example"
+// CHECK: -StringLiteral 0x{{[0-9a-z]+}} {{<col:[0-9]+(, col:[0-9]+)?>}} 'const char[{{[0-9]+}}]' lvalue "def"
+// CHECK: -BinaryOperator 0x{{[0-9a-z]+}} {{<col:[0-9]+(, col:[0-9]+)?>}} 'int' '+'
+// CHECK: -IntegerLiteral 0x{{[0-9a-z]+}} {{<line:[0-9]+:[0-9]+(, col:[0-9]+)?>}} 'int' 5
+// CHECK: -IntegerLiteral 0x{{[0-9a-z]+}} {{<col:[0-9]+(, col:[0-9]+)?>}} 'int' 1
+// CHECK: -AnnotateAttr 0x{{[0-9a-z]+}} {{<line:[0-9]+:[0-9]+(, col:[0-9]+)?>}} "example"
+// CHECK: -StringLiteral 0x{{[0-9a-z]+}} {{<col:[0-9]+(, col:[0-9]+)?>}} 'const char[{{[0-9]+}}]' lvalue "template"
+// CHECK: -IntegerLiteral 0x{{[0-9a-z]+}} {{<col:[0-9]+(, col:[0-9]+)?>}} 'int' 5
 
 //--- bad_attr.cpp
 int var1 __attribute__((example("otherstring"))) = 1; // expected-warning {{'example' attribute only applies to functions}}
 class Example {
   void __attribute__((example)) fn3(); // expected-error {{'example' attribute only allowed at file scope}}
 };
-void fn4() __attribute__((example(123))) { } // expected-error {{first argument to the 'example' attribute must be a string literal}}
-void fn5() __attribute__((example("a","b", 3, 4.0))) { } // expected-error {{'example' attribute only accepts at most three arguments}}
+void fn4() __attribute__((example(123))) { // expected-error {{first argument to the 'example' attribute must be a string literal}}
+  __attribute__((example("somestring"))) while (true); // expected-warning {{'example' attribute only applies to for loop statements}}
+}
+void fn5() __attribute__((example("a","b", 3, 4.0))) { // expected-error {{'example' attribute only accepts at most three arguments}}
+  __attribute__((example("a","b", 3, 4.0))) for (int i = 0; i < 10; ++i) {} // expected-error {{'example' attribute only accepts at most three arguments}}
+}