[Clang] Improve Sema diagnostic performance for __builtin_counted_by_ref

clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -6670,14 +6670,15 @@ def warn_counted_by_attr_elt_type_unknown_size :
 // __builtin_counted_by_ref diagnostics:
 def err_builtin_counted_by_ref_must_be_flex_array_member : Error<
   "'__builtin_counted_by_ref' argument must reference a flexible array member">;
+def err_builtin_counted_by_ref_has_side_effects : Error<
+  "'__builtin_counted_by_ref' argument cannot have side-effects">;
+
 def err_builtin_counted_by_ref_cannot_leak_reference : Error<
-  "value returned by '__builtin_counted_by_ref' cannot be assigned to a "
-  "variable, have its address taken, or passed into or returned from a function">;
-def err_builtin_counted_by_ref_invalid_lhs_use : Error<
+  "value returned by '__builtin_counted_by_ref' cannot be %select{assigned to a "
+  "variable|passed into a function|returned from a function}0">;
+def err_builtin_counted_by_ref_invalid_use : Error<
   "value returned by '__builtin_counted_by_ref' cannot be used in "
   "%select{an array subscript|a binary}0 expression">;
-def err_builtin_counted_by_ref_has_side_effects : Error<
-  "'__builtin_counted_by_ref' argument cannot have side-effects">;
 
 let CategoryName = "ARC Semantic Issue" in {
 

clang/include/clang/Sema/Sema.h

@@ -2518,6 +2518,17 @@ class Sema final : public SemaBase {
 
   bool BuiltinNonDeterministicValue(CallExpr *TheCall);
 
+  enum BuiltinCountedByRefKind {
+    AssignmentKind,
+    InitializerKind,
+    FunctionArgKind,
+    ReturnArgKind,
+    ArraySubscriptKind,
+    BinaryExprKind,
+  };
+
+  bool CheckInvalidBuiltinCountedByRef(const Expr *E,
+                                       BuiltinCountedByRefKind K);
   bool BuiltinCountedByRef(CallExpr *TheCall);
 
   // Matrix builtin handling.

clang/lib/Sema/SemaChecking.cpp

@@ -5621,6 +5621,45 @@ bool Sema::BuiltinCountedByRef(CallExpr *TheCall) {
   return false;
 }
 
+/// The result of __builtin_counted_by_ref cannot be assigned to a variable.
+/// It allows leaking and modification of bounds safety information.
+bool Sema::CheckInvalidBuiltinCountedByRef(const Expr *E,
+                                           BuiltinCountedByRefKind K) {
+  const CallExpr *CE =
+      E ? dyn_cast<CallExpr>(E->IgnoreParenImpCasts()) : nullptr;
+  if (!CE || CE->getBuiltinCallee() != Builtin::BI__builtin_counted_by_ref)
+    return false;
+
+  switch (K) {
+  case AssignmentKind:
+  case InitializerKind:
+    Diag(E->getExprLoc(),
+         diag::err_builtin_counted_by_ref_cannot_leak_reference)
+        << 0 << E->getSourceRange();
+    break;
+  case FunctionArgKind:
+    Diag(E->getExprLoc(),
+         diag::err_builtin_counted_by_ref_cannot_leak_reference)
+        << 1 << E->getSourceRange();
+    break;
+  case ReturnArgKind:
+    Diag(E->getExprLoc(),
+         diag::err_builtin_counted_by_ref_cannot_leak_reference)
+        << 2 << E->getSourceRange();
+    break;
+  case ArraySubscriptKind:
+    Diag(E->getExprLoc(), diag::err_builtin_counted_by_ref_invalid_use)
+        << 0 << E->getSourceRange();
+    break;
+  case BinaryExprKind:
+    Diag(E->getExprLoc(), diag::err_builtin_counted_by_ref_invalid_use)
+        << 1 << E->getSourceRange();
+    break;
+  }
+
+  return true;
+}
+
 namespace {
 
 class UncoveredArgHandler {

clang/lib/Sema/SemaDecl.cpp

@@ -14690,6 +14690,8 @@ void Sema::FinalizeDeclaration(Decl *ThisDecl) {
     }
   }
 
+  CheckInvalidBuiltinCountedByRef(VD->getInit(), InitializerKind);
+
   checkAttributesAfterMerging(*this, *VD);
 
   if (VD->isStaticLocal())

clang/lib/Sema/SemaExpr.cpp

@@ -4894,6 +4894,8 @@ ExprResult Sema::ActOnArraySubscriptExpr(Scope *S, Expr *base,
     return ExprError();
   }
 
+  CheckInvalidBuiltinCountedByRef(base, ArraySubscriptKind);
+
   // Handle any non-overload placeholder types in the base and index
   // expressions.  We can't handle overloads here because the other
   // operand might be an overloadable type, in which case the overload
@@ -6486,6 +6488,12 @@ ExprResult Sema::BuildCallExpr(Scope *Scope, Expr *Fn, SourceLocation LParenLoc,
   if (CheckArgsForPlaceholders(ArgExprs))
     return ExprError();
 
+  // The result of __builtin_counted_by_ref cannot be used as a function
+  // argument. It allows leaking and modification of bounds safety information.
+  for (const Expr *Arg : ArgExprs)
+    if (CheckInvalidBuiltinCountedByRef(Arg, FunctionArgKind))
+      return ExprError();
+
   if (getLangOpts().CPlusPlus) {
     // If this is a pseudo-destructor expression, build the call immediately.
     if (isa<CXXPseudoDestructorExpr>(Fn)) {
@@ -9196,38 +9204,6 @@ Sema::CheckAssignmentConstraints(QualType LHSType, ExprResult &RHS,
   LHSType = Context.getCanonicalType(LHSType).getUnqualifiedType();
   RHSType = Context.getCanonicalType(RHSType).getUnqualifiedType();
 
-  // __builtin_counted_by_ref cannot be assigned to a variable, used in
-  // function call, or in a return.
-  auto FindBuiltinCountedByRefExpr = [&](Expr *E) -> CallExpr * {
-    struct BuiltinCountedByRefVisitor : DynamicRecursiveASTVisitor {
-      CallExpr *TheCall = nullptr;
-      bool VisitCallExpr(CallExpr *CE) override {
-        if (CE->getBuiltinCallee() == Builtin::BI__builtin_counted_by_ref) {
-          TheCall = CE;
-          return false;
-        }
-        return true;
-      }
-      bool
-      VisitUnaryExprOrTypeTraitExpr(UnaryExprOrTypeTraitExpr *UE) override {
-        // A UnaryExprOrTypeTraitExpr---e.g. sizeof, __alignof, etc.---isn't
-        // the same as a CallExpr, so if we find a __builtin_counted_by_ref()
-        // call in one, ignore it.
-        return false;
-      }
-    } V;
-    V.TraverseStmt(E);
-    return V.TheCall;
-  };
-  static llvm::SmallPtrSet<CallExpr *, 4> Diagnosed;
-  if (auto *CE = FindBuiltinCountedByRefExpr(RHS.get());
-      CE && !Diagnosed.count(CE)) {
-    Diagnosed.insert(CE);
-    Diag(CE->getExprLoc(),
-         diag::err_builtin_counted_by_ref_cannot_leak_reference)
-        << CE->getSourceRange();
-  }
-
   // Common case: no conversion required.
   if (LHSType == RHSType) {
     Kind = CK_NoOp;
@@ -13778,42 +13754,6 @@ QualType Sema::CheckAssignmentOperands(Expr *LHSExpr, ExprResult &RHS,
     ConvTy = CheckAssignmentConstraints(Loc, LHSType, RHSType);
   }
 
-  // __builtin_counted_by_ref can't be used in a binary expression or array
-  // subscript on the LHS.
-  int DiagOption = -1;
-  auto FindInvalidUseOfBoundsSafetyCounter = [&](Expr *E) -> CallExpr * {
-    struct BuiltinCountedByRefVisitor : DynamicRecursiveASTVisitor {
-      CallExpr *CE = nullptr;
-      bool InvalidUse = false;
-      int Option = -1;
-
-      bool VisitCallExpr(CallExpr *E) override {
-        if (E->getBuiltinCallee() == Builtin::BI__builtin_counted_by_ref) {
-          CE = E;
-          return false;
-        }
-        return true;
-      }
-
-      bool VisitArraySubscriptExpr(ArraySubscriptExpr *E) override {
-        InvalidUse = true;
-        Option = 0; // report 'array expression' in diagnostic.
-        return true;
-      }
-      bool VisitBinaryOperator(BinaryOperator *E) override {
-        InvalidUse = true;
-        Option = 1; // report 'binary expression' in diagnostic.
-        return true;
-      }
-    } V;
-    V.TraverseStmt(E);
-    DiagOption = V.Option;
-    return V.InvalidUse ? V.CE : nullptr;
-  };
-  if (auto *CE = FindInvalidUseOfBoundsSafetyCounter(LHSExpr))
-    Diag(CE->getExprLoc(), diag::err_builtin_counted_by_ref_invalid_lhs_use)
-        << DiagOption << CE->getSourceRange();
-
   if (DiagnoseAssignmentResult(ConvTy, Loc, LHSType, RHSType, RHS.get(),
                                AssignmentAction::Assigning))
     return QualType();
@@ -15274,6 +15214,12 @@ ExprResult Sema::ActOnBinOp(Scope *S, SourceLocation TokLoc,
   if (Kind == tok::TokenKind::slash)
     DetectPrecisionLossInComplexDivision(*this, TokLoc, LHSExpr);
 
+  BuiltinCountedByRefKind K =
+      BinaryOperator::isAssignmentOp(Opc) ? AssignmentKind : BinaryExprKind;
+
+  CheckInvalidBuiltinCountedByRef(LHSExpr, K);
+  CheckInvalidBuiltinCountedByRef(RHSExpr, K);
+
   return BuildBinOp(S, TokLoc, Opc, LHSExpr, RHSExpr);
 }
 

clang/lib/Sema/SemaStmt.cpp

@@ -3765,6 +3765,8 @@ Sema::ActOnReturnStmt(SourceLocation ReturnLoc, Expr *RetValExp,
         << FSI->getFirstCoroutineStmtKeyword();
   }
 
+  CheckInvalidBuiltinCountedByRef(RetVal.get(), ReturnArgKind);
+
   StmtResult R =
       BuildReturnStmt(ReturnLoc, RetVal.get(), /*AllowRecovery=*/true);
   if (R.isInvalid() || ExprEvalContexts.back().isDiscardedStatementContext())

clang/test/Sema/builtin-counted-by-ref.c

@@ -11,62 +11,63 @@ struct fam_struct {
   int array[] __attribute__((counted_by(count)));
 };
 
-void test1(struct fam_struct *ptr, int size, int idx) {
-  size_t size_of = sizeof(__builtin_counted_by_ref(ptr->array)); // ok
-
-  *__builtin_counted_by_ref(ptr->array) = size;             // ok
+void g(char *);
+void h(char);
 
-  {
-      size_t __ignored_assignment;
-      *_Generic(__builtin_counted_by_ref(ptr->array),
-               void *: &__ignored_assignment,
-               default: __builtin_counted_by_ref(ptr->array)) = 42; // ok
-  }
+void test1(struct fam_struct *ptr, int size, int idx) {
+  size_t size_of = sizeof(__builtin_counted_by_ref(ptr->array));    // ok
+  int align_of = __alignof(__builtin_counted_by_ref(ptr->array));   // ok
+  size_t __ignored_assignment;
+
+  *__builtin_counted_by_ref(ptr->array) = size;                     // ok
+  *_Generic(__builtin_counted_by_ref(ptr->array),
+           void *: &__ignored_assignment,
+           default: __builtin_counted_by_ref(ptr->array)) = 42;     // ok
+  h(*__builtin_counted_by_ref(ptr->array));                         // ok
 }
 
 void test2(struct fam_struct *ptr, int idx) {
-  __builtin_counted_by_ref();                               // expected-error {{too few arguments to function call, expected 1, have 0}}
-  __builtin_counted_by_ref(ptr->array, ptr->x, ptr->count); // expected-error {{too many arguments to function call, expected 1, have 3}}
+  __builtin_counted_by_ref();                                       // expected-error {{too few arguments to function call, expected 1, have 0}}
+  __builtin_counted_by_ref(ptr->array, ptr->x, ptr->count);         // expected-error {{too many arguments to function call, expected 1, have 3}}
 }
 
 void test3(struct fam_struct *ptr, int idx) {
-  __builtin_counted_by_ref(&ptr->array[0]);                 // expected-error {{'__builtin_counted_by_ref' argument must reference a flexible array member}}
-  __builtin_counted_by_ref(&ptr->array[idx]);               // expected-error {{'__builtin_counted_by_ref' argument must reference a flexible array member}}
-  __builtin_counted_by_ref(&ptr->array);                    // expected-error {{'__builtin_counted_by_ref' argument must reference a flexible array member}}
-  __builtin_counted_by_ref(ptr->x);                         // expected-error {{'__builtin_counted_by_ref' argument must reference a flexible array member}}
-  __builtin_counted_by_ref(&ptr->x);                        // expected-error {{'__builtin_counted_by_ref' argument must reference a flexible array member}}
-  __builtin_counted_by_ref(global_array);                   // expected-error {{'__builtin_counted_by_ref' argument must reference a flexible array member}}
-  __builtin_counted_by_ref(global_int);                     // expected-error {{'__builtin_counted_by_ref' argument must reference a flexible array member}}
-  __builtin_counted_by_ref(&global_int);                    // expected-error {{'__builtin_counted_by_ref' argument must reference a flexible array member}}
+  __builtin_counted_by_ref(&ptr->array[0]);                         // expected-error {{'__builtin_counted_by_ref' argument must reference a flexible array member}}
+  __builtin_counted_by_ref(&ptr->array[idx]);                       // expected-error {{'__builtin_counted_by_ref' argument must reference a flexible array member}}
+  __builtin_counted_by_ref(&ptr->array);                            // expected-error {{'__builtin_counted_by_ref' argument must reference a flexible array member}}
+  __builtin_counted_by_ref(ptr->x);                                 // expected-error {{'__builtin_counted_by_ref' argument must reference a flexible array member}}
+  __builtin_counted_by_ref(&ptr->x);                                // expected-error {{'__builtin_counted_by_ref' argument must reference a flexible array member}}
+  __builtin_counted_by_ref(global_array);                           // expected-error {{'__builtin_counted_by_ref' argument must reference a flexible array member}}
+  __builtin_counted_by_ref(global_int);                             // expected-error {{'__builtin_counted_by_ref' argument must reference a flexible array member}}
+  __builtin_counted_by_ref(&global_int);                            // expected-error {{'__builtin_counted_by_ref' argument must reference a flexible array member}}
 }
 
 void test4(struct fam_struct *ptr, int idx) {
-  __builtin_counted_by_ref(ptr++->array);                   // expected-error {{'__builtin_counted_by_ref' argument cannot have side-effects}}
-  __builtin_counted_by_ref(&ptr->array[idx++]);             // expected-error {{'__builtin_counted_by_ref' argument cannot have side-effects}}
+  __builtin_counted_by_ref(ptr++->array);                           // expected-error {{'__builtin_counted_by_ref' argument cannot have side-effects}}
+  __builtin_counted_by_ref(&ptr->array[idx++]);                     // expected-error {{'__builtin_counted_by_ref' argument cannot have side-effects}}
 }
 
-void foo(char *);
-
 void *test5(struct fam_struct *ptr, int size, int idx) {
-  char *ref = __builtin_counted_by_ref(ptr->array);         // expected-error {{value returned by '__builtin_counted_by_ref' cannot be assigned to a variable, have its address taken, or passed into or returned from a function}}
+  char *ref = __builtin_counted_by_ref(ptr->array);                 // expected-error {{value returned by '__builtin_counted_by_ref' cannot be assigned to a variable}}
+  char *int_ptr;
+  char *p;
 
-  ref = __builtin_counted_by_ref(ptr->array);               // expected-error {{value returned by '__builtin_counted_by_ref' cannot be assigned to a variable, have its address taken, or passed into or returned from a function}}
-  ref = (char *)(int *)(42 + &*__builtin_counted_by_ref(ptr->array)); // expected-error {{value returned by '__builtin_counted_by_ref' cannot be assigned to a variable, have its address taken, or passed into or returned from a function}}
-  foo(__builtin_counted_by_ref(ptr->array));                // expected-error {{value returned by '__builtin_counted_by_ref' cannot be assigned to a variable, have its address taken, or passed into or returned from a function}}
-  foo(ref = __builtin_counted_by_ref(ptr->array));          // expected-error {{value returned by '__builtin_counted_by_ref' cannot be assigned to a variable, have its address taken, or passed into or returned from a function}}
+  ref = __builtin_counted_by_ref(ptr->array);                       // expected-error {{value returned by '__builtin_counted_by_ref' cannot be assigned to a variable}}
+  g(__builtin_counted_by_ref(ptr->array));                          // expected-error {{value returned by '__builtin_counted_by_ref' cannot be passed into a function}}
+  g(ref = __builtin_counted_by_ref(ptr->array));                    // expected-error {{value returned by '__builtin_counted_by_ref' cannot be assigned to a variable}}
 
-  if ((ref = __builtin_counted_by_ref(ptr->array)))         // expected-error {{value returned by '__builtin_counted_by_ref' cannot be assigned to a variable, have its address taken, or passed into or returned from a function}}
+  if ((ref = __builtin_counted_by_ref(ptr->array)))                 // expected-error {{value returned by '__builtin_counted_by_ref' cannot be assigned to a variable}}
     ;
 
-  for (char *p = __builtin_counted_by_ref(ptr->array); p && *p; ++p) // expected-error {{value returned by '__builtin_counted_by_ref' cannot be assigned to a variable, have its address taken, or passed into or returned from a function}}
+  for (p = __builtin_counted_by_ref(ptr->array); p && *p; ++p)      // expected-error {{value returned by '__builtin_counted_by_ref' cannot be assigned to a variable}}
     ;
 
-  return __builtin_counted_by_ref(ptr->array);              // expected-error {{value returned by '__builtin_counted_by_ref' cannot be assigned to a variable, have its address taken, or passed into or returned from a function}}
+  return __builtin_counted_by_ref(ptr->array);                      // expected-error {{value returned by '__builtin_counted_by_ref' cannot be returned from a function}}
 }
 
 void test6(struct fam_struct *ptr, int size, int idx) {
-  *(__builtin_counted_by_ref(ptr->array) + 4) = 37;         // expected-error {{value returned by '__builtin_counted_by_ref' cannot be used in a binary expression}}
-  __builtin_counted_by_ref(ptr->array)[3] = 37;             // expected-error {{value returned by '__builtin_counted_by_ref' cannot be used in an array subscript expression}}
+  *(__builtin_counted_by_ref(ptr->array) + 4) = 37;                 // expected-error {{value returned by '__builtin_counted_by_ref' cannot be used in a binary expression}}
+  __builtin_counted_by_ref(ptr->array)[3] = 37;                     // expected-error {{value returned by '__builtin_counted_by_ref' cannot be used in an array subscript expression}}
 }
 
 struct non_fam_struct {
@@ -77,10 +78,10 @@ struct non_fam_struct {
 };
 
 void *test7(struct non_fam_struct *ptr, int size) {
-  *__builtin_counted_by_ref(ptr->array) = size              // expected-error {{'__builtin_counted_by_ref' argument must reference a flexible array member}}
-  *__builtin_counted_by_ref(&ptr->array[0]) = size;         // expected-error {{'__builtin_counted_by_ref' argument must reference a flexible array member}}
-  *__builtin_counted_by_ref(ptr->pointer) = size;           // expected-error {{'__builtin_counted_by_ref' argument must reference a flexible array member}}
-  *__builtin_counted_by_ref(&ptr->pointer[0]) = size;       // expected-error {{'__builtin_counted_by_ref' argument must reference a flexible array member}}
+  *__builtin_counted_by_ref(ptr->array) = size                      // expected-error {{'__builtin_counted_by_ref' argument must reference a flexible array member}}
+  *__builtin_counted_by_ref(&ptr->array[0]) = size;                 // expected-error {{'__builtin_counted_by_ref' argument must reference a flexible array member}}
+  *__builtin_counted_by_ref(ptr->pointer) = size;                   // expected-error {{'__builtin_counted_by_ref' argument must reference a flexible array member}}
+  *__builtin_counted_by_ref(&ptr->pointer[0]) = size;               // expected-error {{'__builtin_counted_by_ref' argument must reference a flexible array member}}
 }
 
 struct char_count {