[GCAtomic] Treat single capability structs like a capability

clang/include/clang/AST/Type.h

@@ -2145,6 +2145,9 @@ class alignas(8) Type : public ExtQualsTypeCommonBase {
   /// pointers.
   bool isCHERICapabilityType(const ASTContext &Context,
                              bool IncludeIntCap = true) const;
+  /// Returns true if this is a struct/union type that contains exactly one
+  /// capability element.
+  bool isSingleCapabilityRecord(const ASTContext &Context) const;
   /// Returns true for __uintcap_t or __intcap_t (and enums/_Atomic with that
   /// underlying type)
   bool isIntCapType() const;

clang/lib/AST/Type.cpp

@@ -649,6 +649,18 @@ bool Type::isCHERICapabilityType(const ASTContext &Context,
   return false;
 }
 
+bool Type::isSingleCapabilityRecord(const ASTContext &Context) const {
+  if (auto *RT = getAs<RecordType>())
+    return Context.containsCapabilities(RT->getDecl()) &&
+           Context.getTypeSize(this) ==
+               Context.getTargetInfo().getCHERICapabilityWidth();
+  if (const AtomicType *AT = getAs<AtomicType>())
+    return AT->getValueType()->isSingleCapabilityRecord(Context);
+  if (const AttributedType *AT = getAs<AttributedType>())
+    return AT->getModifiedType()->isSingleCapabilityRecord(Context);
+  return false;
+}
+
 bool Type::isIntCapType() const {
   if (const BuiltinType *BT = dyn_cast<BuiltinType>(CanonicalType))
     return BT->getKind() == BuiltinType::IntCap ||

clang/lib/CodeGen/CGAtomic.cpp

@@ -146,7 +146,8 @@ bool isAtomicStoreOp(AtomicExpr::AtomicOp Op) {
       }
       UseLibcall = !C.getTargetInfo().hasBuiltinAtomic(
           AtomicSizeInBits, C.toBits(lvalue.getAlignment()),
-          AtomicTy->isCHERICapabilityType(CGF.CGM.getContext()));
+          AtomicTy->isCHERICapabilityType(CGF.getContext()) ||
+              AtomicTy->isSingleCapabilityRecord(CGF.getContext()));
     }
 
     QualType getAtomicType() const { return AtomicTy; }
@@ -547,7 +548,8 @@ static void EmitAtomicOp(CodeGenFunction &CGF, AtomicExpr *E, Address Dest,
   bool PostOpMinMax = false;
   unsigned PostOp = 0;
   QualType AtomicTy = E->getPtr()->getType()->getPointeeType();
-  bool IsCheriCap = AtomicTy->isCHERICapabilityType(CGF.CGM.getContext());
+  bool IsCheriCap = AtomicTy->isCHERICapabilityType(CGF.getContext()) ||
+                    AtomicTy->isSingleCapabilityRecord(CGF.getContext());
 
   switch (E->getOp()) {
   case AtomicExpr::AO__c11_atomic_init:
@@ -811,12 +813,14 @@ static void
 AddDirectArgument(CodeGenFunction &CGF, CallArgList &Args,
                   bool UseOptimizedLibcall, llvm::Value *Val, QualType ValTy,
                   SourceLocation Loc, CharUnits SizeInChars) {
+  bool IsCapTy = ValTy->isCHERICapabilityType(CGF.getContext()) ||
+                 ValTy->isSingleCapabilityRecord(CGF.getContext());
   if (UseOptimizedLibcall) {
     // Load value and pass it to the function directly.
     CharUnits Align = CGF.getContext().getTypeAlignInChars(ValTy);
     int64_t SizeInBits = CGF.getContext().toBits(SizeInChars);
     llvm::Type *ITy;
-    if (ValTy->isCHERICapabilityType(CGF.getContext())) {
+    if (IsCapTy) {
       ValTy = CGF.getContext().getPointerType(CGF.getContext().VoidTy,
                                               PIK_Capability);
       ITy = CGF.Int8CheriCapTy;
@@ -836,7 +840,7 @@ AddDirectArgument(CodeGenFunction &CGF, CallArgList &Args,
   } else {
     // Non-optimized functions always take a reference.
     // NB: Capabilities must be passed directly to the optimized libcall
-    assert(!ValTy->isCHERICapabilityType(CGF.getContext()) &&
+    assert(!IsCapTy &&
            "Capabilities should not be passed to the generic libcall");
     Args.add(RValue::get(CGF.EmitCastToVoidPtr(Val)),
                          CGF.getContext().VoidPtrTy);
@@ -866,7 +870,8 @@ RValue CodeGenFunction::EmitAtomicExpr(AtomicExpr *E) {
   uint64_t Size = TInfo.Width.getQuantity();
   unsigned MaxInlineWidthInBits = getTarget().getMaxAtomicInlineWidth();
 
-  bool IsCheriCap = AtomicTy->isCHERICapabilityType(CGM.getContext());
+  bool IsCheriCap = AtomicTy->isCHERICapabilityType(CGM.getContext()) ||
+                    AtomicTy->isSingleCapabilityRecord(CGM.getContext());
   bool Oversized = (!IsCheriCap &&
                     getContext().toBits(TInfo.Width) > MaxInlineWidthInBits) ||
                    (IsCheriCap && MaxInlineWidthInBits == 0);
@@ -1519,14 +1524,16 @@ RValue CodeGenFunction::EmitAtomicExpr(AtomicExpr *E) {
 
 Address AtomicInfo::emitCastToAtomicIntPointer(Address addr) const {
   llvm::Type *ty;
-  if (AtomicTy->isCHERICapabilityType(CGF.getContext())) {
+  if (AtomicTy->isCHERICapabilityType(CGF.getContext()) ||
+      AtomicTy->isSingleCapabilityRecord(CGF.getContext())) {
     // If capability atomics are natively supported the instruction expects
     // a capability type. We also pass capabilities directly to the atomic
     // libcalls (i.e. always use optimized ones) since this is required to
     // support the RMW operations and special-casing the load/store/xchg to
     // use the generic libcalls (with mutex+memcpy) adds unncessary complexity.
-    if (!UseLibcall) {
-      // If we aren't using a libcall there is no need to cast to i8*
+    if (!UseLibcall && !AtomicTy->isSingleCapabilityRecord(CGF.getContext())) {
+      // If we aren't using a libcall and aren't using a single-capability
+      // struct a there is no need to cast to i8*
       return CGF.Builder.CreateElementBitCast(
           addr, getAtomicAddress().getElementType());
     }

clang/lib/CodeGen/TargetInfo.cpp

@@ -11536,11 +11536,7 @@ ABIArgInfo RISCVABIInfo::classifyArgumentType(QualType Ty, bool IsFixed,
 
   uint64_t Size = getContext().getTypeSize(Ty);
 
-  bool IsSingleCapRecord = false;
-  if (auto *RT = Ty->getAs<RecordType>())
-    IsSingleCapRecord = Size == getTarget().getCHERICapabilityWidth() &&
-                        getContext().containsCapabilities(RT->getDecl());
-
+  bool IsSingleCapRecord = Ty->isSingleCapabilityRecord(getContext());
   bool IsCapability = Ty->isCHERICapabilityType(getContext()) ||
                       IsSingleCapRecord;
 
@@ -11687,13 +11683,7 @@ Address RISCVABIInfo::EmitVAArg(CodeGenFunction &CGF, Address VAListAddr,
 
   auto TInfo = getContext().getTypeInfoInChars(Ty);
 
-  bool IsSingleCapRecord = false;
-  CharUnits CapabilityWidth =
-    CharUnits::fromQuantity(getTarget().getCHERICapabilityWidth() / 8);
-  if (const auto *RT = Ty->getAs<RecordType>())
-    IsSingleCapRecord = TInfo.Width == CapabilityWidth &&
-                        getContext().containsCapabilities(RT->getDecl());
-
+  bool IsSingleCapRecord = Ty->isSingleCapabilityRecord(getContext());
   bool IsCapability = Ty->isCHERICapabilityType(getContext()) ||
                       IsSingleCapRecord;