[Clang] Improve Sema diagnostic performance for __builtin_counted_by_ref

clang/include/clang/AST/ASTContext.h

@@ -1184,6 +1184,7 @@ class ASTContext : public RefCountedBase<ASTContext> {
   CanQualType DependentTy, OverloadTy, BoundMemberTy, UnresolvedTemplateTy,
       UnknownAnyTy;
   CanQualType BuiltinFnTy;
+  CanQualType BuiltinCountedByRefTy;
   CanQualType PseudoObjectTy, ARCUnbridgedCastTy;
   CanQualType ObjCBuiltinIdTy, ObjCBuiltinClassTy, ObjCBuiltinSelTy;
   CanQualType ObjCBuiltinBoolTy;

clang/include/clang/AST/BuiltinTypes.def

@@ -314,6 +314,9 @@ PLACEHOLDER_TYPE(UnknownAny, UnknownAnyTy)
 
 PLACEHOLDER_TYPE(BuiltinFn, BuiltinFnTy)
 
+// A placeholder type for __builtin_counted_by_ref.
+PLACEHOLDER_TYPE(BuiltinCountedByRef, BuiltinCountedByRefTy)
+
 // The type of a cast which, in ARC, would normally require a
 // __bridge, but which might be okay depending on the immediate
 // context.

clang/include/clang/Basic/DiagnosticSemaKinds.td

@@ -6673,7 +6673,7 @@ def err_builtin_counted_by_ref_must_be_flex_array_member : Error<
 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<
+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<

clang/include/clang/Serialization/ASTBitCodes.h

@@ -1111,6 +1111,9 @@ enum PredefinedTypeIDs {
   /// \brief The '__ibm128' type
   PREDEF_TYPE_IBM128_ID = 74,
 
+  /// \brief The placeholder type for __builtin_counted_by_ref.
+  PREDEF_TYPE_BUILTIN_COUNTED_BY_REF = 75,
+
 /// OpenCL image types with auto numeration
 #define IMAGE_TYPE(ImgType, Id, SingletonId, Access, Suffix)                   \
   PREDEF_TYPE_##Id##_ID,
@@ -1148,7 +1151,7 @@ enum PredefinedTypeIDs {
 ///
 /// Type IDs for non-predefined types will start at
 /// NUM_PREDEF_TYPE_IDs.
-const unsigned NUM_PREDEF_TYPE_IDS = 513;
+const unsigned NUM_PREDEF_TYPE_IDS = 514;
 
 // Ensure we do not overrun the predefined types we reserved
 // in the enum PredefinedTypeIDs above.

clang/lib/AST/ASTContext.cpp

@@ -1390,6 +1390,10 @@ void ASTContext::InitBuiltinTypes(const TargetInfo &Target,
   if (LangOpts.MatrixTypes)
     InitBuiltinType(IncompleteMatrixIdxTy, BuiltinType::IncompleteMatrixIdx);
 
+  // Placeholder for __builtin_counted_by_ref().
+  if (!LangOpts.CPlusPlus)
+    InitBuiltinType(BuiltinCountedByRefTy, BuiltinType::BuiltinCountedByRef);
+
   // Builtin types for 'id', 'Class', and 'SEL'.
   InitBuiltinType(ObjCBuiltinIdTy, BuiltinType::ObjCId);
   InitBuiltinType(ObjCBuiltinClassTy, BuiltinType::ObjCClass);

clang/lib/AST/NSAPI.cpp

@@ -466,6 +466,7 @@ NSAPI::getNSNumberFactoryMethodKind(QualType T) const {
   case BuiltinType::Half:
   case BuiltinType::PseudoObject:
   case BuiltinType::BuiltinFn:
+  case BuiltinType::BuiltinCountedByRef:
   case BuiltinType::IncompleteMatrixIdx:
   case BuiltinType::ArraySection:
   case BuiltinType::OMPArrayShaping:

clang/lib/AST/Type.cpp

@@ -3444,6 +3444,8 @@ StringRef BuiltinType::getName(const PrintingPolicy &Policy) const {
     return "<ARC unbridged cast type>";
   case BuiltinFn:
     return "<builtin fn type>";
+  case BuiltinCountedByRef:
+    return "<builtin counted by ref type>";
   case ObjCId:
     return "id";
   case ObjCClass:
@@ -4861,6 +4863,7 @@ bool Type::canHaveNullability(bool ResultIfUnknown) const {
 #define HLSL_INTANGIBLE_TYPE(Name, Id, SingletonId) case BuiltinType::Id:
 #include "clang/Basic/HLSLIntangibleTypes.def"
     case BuiltinType::BuiltinFn:
+    case BuiltinType::BuiltinCountedByRef:
     case BuiltinType::NullPtr:
     case BuiltinType::IncompleteMatrixIdx:
     case BuiltinType::ArraySection:

clang/lib/AST/TypeLoc.cpp

@@ -433,6 +433,7 @@ TypeSpecifierType BuiltinTypeLoc::getWrittenTypeSpec() const {
 #define HLSL_INTANGIBLE_TYPE(Name, Id, SingletonId) case BuiltinType::Id:
 #include "clang/Basic/HLSLIntangibleTypes.def"
   case BuiltinType::BuiltinFn:
+  case BuiltinType::BuiltinCountedByRef:
   case BuiltinType::IncompleteMatrixIdx:
   case BuiltinType::ArraySection:
   case BuiltinType::OMPArrayShaping:

clang/lib/Sema/SemaChecking.cpp

@@ -755,6 +755,7 @@ static ExprResult BuiltinDumpStruct(Sema &S, CallExpr *TheCall) {
     case BuiltinType::PseudoObject:
     case BuiltinType::UnknownAny:
     case BuiltinType::BuiltinFn:
+    case BuiltinType::BuiltinCountedByRef:
       // This might be a callable.
       break;
 

clang/lib/Sema/SemaDecl.cpp

@@ -14690,6 +14690,17 @@ 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 (const auto *CE = dyn_cast_if_present<CallExpr>(VD->getInit())) {
+    const Expr *E = CE->getCallee()->IgnoreParenImpCasts();
+    if (const BuiltinType *PTy = E->getType()->getAsPlaceholderType();
+        PTy && PTy->getKind() == BuiltinType::BuiltinCountedByRef)
+      Diag(E->getExprLoc(),
+           diag::err_builtin_counted_by_ref_cannot_leak_reference)
+          << E->getSourceRange();
+  }
+
   checkAttributesAfterMerging(*this, *VD);
 
   if (VD->isStaticLocal())

clang/lib/Sema/SemaExpr.cpp

@@ -3382,7 +3382,9 @@ ExprResult Sema::BuildDeclarationNameExpr(
   case Decl::Function: {
     if (unsigned BID = cast<FunctionDecl>(VD)->getBuiltinID()) {
       if (!Context.BuiltinInfo.isDirectlyAddressable(BID)) {
-        type = Context.BuiltinFnTy;
+        type = (BID == Builtin::BI__builtin_counted_by_ref)
+                   ? Context.BuiltinCountedByRefTy
+                   : Context.BuiltinFnTy;
         valueKind = VK_PRValue;
         break;
       }
@@ -4894,6 +4896,18 @@ 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.
+  Expr *E = base->IgnoreParenImpCasts();
+  if (auto *CE = dyn_cast<CallExpr>(E)) {
+    E = CE->getCallee()->IgnoreParenImpCasts();
+    if (const BuiltinType *PTy = E->getType()->getAsPlaceholderType();
+        PTy && PTy->getKind() == BuiltinType::BuiltinCountedByRef) {
+      Diag(E->getExprLoc(), diag::err_builtin_counted_by_ref_invalid_use)
+          << 0 << E->getSourceRange();
+    }
+  }
+
   // 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
@@ -6188,6 +6202,7 @@ static bool isPlaceholderToRemoveAsArg(QualType type) {
   // These are always invalid as call arguments and should be reported.
   case BuiltinType::BoundMember:
   case BuiltinType::BuiltinFn:
+  case BuiltinType::BuiltinCountedByRef:
   case BuiltinType::IncompleteMatrixIdx:
   case BuiltinType::ArraySection:
   case BuiltinType::OMPArrayShaping:
@@ -6205,10 +6220,27 @@ bool Sema::CheckArgsForPlaceholders(MultiExprArg args) {
   for (size_t i = 0, e = args.size(); i != e; i++) {
     if (isPlaceholderToRemoveAsArg(args[i]->getType())) {
       ExprResult result = CheckPlaceholderExpr(args[i]);
-      if (result.isInvalid()) hasInvalid = true;
-      else args[i] = result.get();
+      if (result.isInvalid())
+        hasInvalid = true;
+      else
+        args[i] = result.get();
+    }
+
+    // The result of __builtin_counted_by_ref cannot be used as a function
+    // argument. It allows leaking and modification of bounds safety
+    // information.
+    if (const auto *CE = dyn_cast<CallExpr>(args[i])) {
+      const Expr *E = CE->getCallee()->IgnoreParenImpCasts();
+      if (const BuiltinType *PTy = E->getType()->getAsPlaceholderType();
+          PTy && PTy->getKind() == BuiltinType::BuiltinCountedByRef) {
+        hasInvalid = true;
+        Diag(E->getExprLoc(),
+             diag::err_builtin_counted_by_ref_cannot_leak_reference)
+            << E->getSourceRange();
+      }
     }
   }
+
   return hasInvalid;
 }
 
@@ -6770,7 +6802,9 @@ ExprResult Sema::BuildResolvedCallExpr(Expr *Fn, NamedDecl *NDecl,
   ExprResult Result;
   QualType ResultTy;
   if (BuiltinID &&
-      Fn->getType()->isSpecificBuiltinType(BuiltinType::BuiltinFn)) {
+      (Fn->getType()->isSpecificBuiltinType(BuiltinType::BuiltinFn) ||
+       Fn->getType()->isSpecificBuiltinType(
+           BuiltinType::BuiltinCountedByRef))) {
     // Extract the return type from the (builtin) function pointer type.
     // FIXME Several builtins still have setType in
     // Sema::CheckBuiltinFunctionCall. One should review their definitions in
@@ -9196,38 +9230,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 +13780,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 +15240,27 @@ 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 CheckBuiltinCountedByRefPlaceholder = [&](const Expr *E) {
+    if (const auto *CE = dyn_cast<CallExpr>(E->IgnoreParenImpCasts())) {
+      E = CE->getCallee()->IgnoreParenImpCasts();
+      if (const BuiltinType *PTy = E->getType()->getAsPlaceholderType();
+          PTy && PTy->getKind() == BuiltinType::BuiltinCountedByRef) {
+        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();
+      }
+    }
+  };
+
+  CheckBuiltinCountedByRefPlaceholder(LHSExpr);
+  CheckBuiltinCountedByRefPlaceholder(RHSExpr);
+
   return BuildBinOp(S, TokLoc, Opc, LHSExpr, RHSExpr);
 }
 
@@ -21074,6 +21061,13 @@ ExprResult Sema::CheckPlaceholderExpr(Expr *E) {
     return ExprError();
   }
 
+  case BuiltinType::BuiltinCountedByRef: {
+    Diag(E->IgnoreParens()->getExprLoc(),
+         diag::err_builtin_counted_by_ref_cannot_leak_reference)
+        << E->IgnoreParens()->getSourceRange();
+    return ExprError();
+  }
+
   case BuiltinType::IncompleteMatrixIdx:
     Diag(cast<MatrixSubscriptExpr>(E->IgnoreParens())
              ->getRowIdx()

clang/lib/Sema/SemaStmt.cpp

@@ -3765,6 +3765,16 @@ Sema::ActOnReturnStmt(SourceLocation ReturnLoc, Expr *RetValExp,
         << FSI->getFirstCoroutineStmtKeyword();
   }
 
+  if (const auto *CE = dyn_cast_if_present<CallExpr>(RetVal.get())) {
+    const Expr *E = CE->getCallee()->IgnoreParenImpCasts();
+    if (const BuiltinType *PTy = E->getType()->getAsPlaceholderType();
+        PTy && PTy->getKind() == BuiltinType::BuiltinCountedByRef) {
+      Diag(E->getExprLoc(),
+           diag::err_builtin_counted_by_ref_cannot_leak_reference)
+          << E->getSourceRange();
+    }
+  }
+
   StmtResult R =
       BuildReturnStmt(ReturnLoc, RetVal.get(), /*AllowRecovery=*/true);
   if (R.isInvalid() || ExprEvalContexts.back().isDiscardedStatementContext())

clang/lib/Serialization/ASTCommon.cpp

@@ -274,6 +274,9 @@ serialization::TypeIdxFromBuiltin(const BuiltinType *BT) {
   case BuiltinType::BuiltinFn:
     ID = PREDEF_TYPE_BUILTIN_FN;
     break;
+  case BuiltinType::BuiltinCountedByRef:
+    ID = PREDEF_TYPE_BUILTIN_COUNTED_BY_REF;
+    break;
   case BuiltinType::IncompleteMatrixIdx:
     ID = PREDEF_TYPE_INCOMPLETE_MATRIX_IDX;
     break;

clang/lib/Serialization/ASTReader.cpp

@@ -7458,6 +7458,9 @@ QualType ASTReader::GetType(TypeID ID) {
     case PREDEF_TYPE_BUILTIN_FN:
       T = Context.BuiltinFnTy;
       break;
+    case PREDEF_TYPE_BUILTIN_COUNTED_BY_REF:
+      T = Context.BuiltinCountedByRefTy;
+      break;
     case PREDEF_TYPE_INCOMPLETE_MATRIX_IDX:
       T = Context.IncompleteMatrixIdxTy;
       break;

clang/test/Modules/no-external-type-id.cppm

@@ -23,7 +23,7 @@ export module b;
 import a;
 export int b();
 
-// CHECK: <DECL_FUNCTION {{.*}} op8=4120
+// CHECK: <DECL_FUNCTION {{.*}} op8=4128
 // CHECK: <TYPE_FUNCTION_PROTO
 
 //--- a.v1.cppm