Create a separate checking function.

Author: bwendling

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">;
+  "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,11 +2518,17 @@ class Sema final : public SemaBase {
 
   bool BuiltinNonDeterministicValue(CallExpr *TheCall);
 
-  bool IsBuiltinCountedByRef(const Expr *E) const {
-    const CallExpr *CE =
-        E ? dyn_cast<CallExpr>(E->IgnoreParenImpCasts()) : nullptr;
-    return CE && CE->getBuiltinCallee() == Builtin::BI__builtin_counted_by_ref;
-  }
+  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,12 +14690,7 @@ void Sema::FinalizeDeclaration(Decl *ThisDecl) {
     }
   }
 
-  // The result of __builtin_counted_by_ref cannot be assigned to a variable.
-  // It allows leaking and modification of bounds safety information.
-  if (IsBuiltinCountedByRef(VD->getInit()))
-    Diag(VD->getInit()->getExprLoc(),
-         diag::err_builtin_counted_by_ref_cannot_leak_reference)
-        << VD->getInit()->getSourceRange();
+  CheckInvalidBuiltinCountedByRef(VD->getInit(), InitializerKind);
 
   checkAttributesAfterMerging(*this, *VD);
 

clang/lib/Sema/SemaExpr.cpp

@@ -4894,11 +4894,7 @@ ExprResult Sema::ActOnArraySubscriptExpr(Scope *S, Expr *base,
     return ExprError();
   }
 
-  // We cannot use __builtin_counted_by_ref in a binary expression. It's
-  // possible to leak the reference and violate bounds security.
-  if (IsBuiltinCountedByRef(base))
-    Diag(base->getExprLoc(), diag::err_builtin_counted_by_ref_invalid_use)
-        << 0 << base->getSourceRange();
+  CheckInvalidBuiltinCountedByRef(base, ArraySubscriptKind);
 
   // Handle any non-overload placeholder types in the base and index
   // expressions.  We can't handle overloads here because the other
@@ -6495,12 +6491,8 @@ ExprResult Sema::BuildCallExpr(Scope *Scope, Expr *Fn, SourceLocation LParenLoc,
   // 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 (IsBuiltinCountedByRef(Arg)) {
-      Diag(Arg->getExprLoc(),
-           diag::err_builtin_counted_by_ref_cannot_leak_reference)
-          << Arg->getSourceRange();
+    if (CheckInvalidBuiltinCountedByRef(Arg, FunctionArgKind))
       return ExprError();
-    }
 
   if (getLangOpts().CPlusPlus) {
     // If this is a pseudo-destructor expression, build the call immediately.
@@ -15222,22 +15214,11 @@ ExprResult Sema::ActOnBinOp(Scope *S, SourceLocation TokLoc,
   if (Kind == tok::TokenKind::slash)
     DetectPrecisionLossInComplexDivision(*this, TokLoc, LHSExpr);
 
-  // We cannot use __builtin_counted_by_ref in a binary expression. It's
-  // possible to leak the reference and violate bounds security.
-  auto CheckBuiltinCountedByRef = [&](const Expr *E) {
-    if (IsBuiltinCountedByRef(E)) {
-      if (BinaryOperator::isAssignmentOp(Opc))
-        Diag(E->getExprLoc(),
-             diag::err_builtin_counted_by_ref_cannot_leak_reference)
-            << E->getSourceRange();
-      else
-        Diag(E->getExprLoc(), diag::err_builtin_counted_by_ref_invalid_use)
-            << 1 << E->getSourceRange();
-    }
-  };
+  BuiltinCountedByRefKind K =
+      BinaryOperator::isAssignmentOp(Opc) ? AssignmentKind : BinaryExprKind;
 
-  CheckBuiltinCountedByRef(LHSExpr);
-  CheckBuiltinCountedByRef(RHSExpr);
+  CheckInvalidBuiltinCountedByRef(LHSExpr, K);
+  CheckInvalidBuiltinCountedByRef(RHSExpr, K);
 
   return BuildBinOp(S, TokLoc, Opc, LHSExpr, RHSExpr);
 }

clang/lib/Sema/SemaStmt.cpp

@@ -3765,10 +3765,7 @@ Sema::ActOnReturnStmt(SourceLocation ReturnLoc, Expr *RetValExp,
         << FSI->getFirstCoroutineStmtKeyword();
   }
 
-  if (IsBuiltinCountedByRef(RetVal.get()))
-    Diag(RetVal.get()->getExprLoc(),
-         diag::err_builtin_counted_by_ref_cannot_leak_reference)
-        << RetVal.get()->getSourceRange();
+  CheckInvalidBuiltinCountedByRef(RetVal.get(), ReturnArgKind);
 
   StmtResult R =
       BuildReturnStmt(ReturnLoc, RetVal.get(), /*AllowRecovery=*/true);

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

@@ -48,21 +48,21 @@ void test4(struct fam_struct *ptr, int idx) {
 }
 
 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}}
-  g(__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}}
-  g(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 (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) {