[CIR][CIRGen] Implement array cookie ABI for AppleARM64 targets (#1301)

This change introduces CIRGenCXXABI subclasses for ARM and AppleARM64
and implements ARM CXXABI handling for array cookies.

clang/lib/CIR/CodeGen/CIRGenItaniumCXXABI.cpp

@@ -360,6 +360,30 @@ class CIRGenItaniumCXXABI : public CIRGenCXXABI {
   classifyRTTIUniqueness(QualType CanTy, cir::GlobalLinkageKind Linkage) const;
   friend class CIRGenItaniumRTTIBuilder;
 };
+
+class CIRGenARMCXXABI : public CIRGenItaniumCXXABI {
+public:
+  CIRGenARMCXXABI(CIRGenModule &CGM) : CIRGenItaniumCXXABI(CGM) {
+    // TODO(cir): When implemented, /*UseARMMethodPtrABI=*/true,
+    //                              /*UseARMGuardVarABI=*/true) {}
+    assert(!cir::MissingFeatures::appleArm64CXXABI());
+  }
+  CharUnits getArrayCookieSizeImpl(QualType elementType) override;
+  Address initializeArrayCookie(CIRGenFunction &CGF, Address NewPtr,
+                                mlir::Value NumElements, const CXXNewExpr *E,
+                                QualType ElementType) override;
+};
+
+class CIRGenAppleARM64CXXABI : public CIRGenARMCXXABI {
+public:
+  CIRGenAppleARM64CXXABI(CIRGenModule &CGM) : CIRGenARMCXXABI(CGM) {
+    Use32BitVTableOffsetABI = true;
+  }
+
+  // ARM64 libraries are prepared for non-unique RTTI.
+  bool shouldRTTIBeUnique() const override { return false; }
+};
+
 } // namespace
 
 CIRGenCXXABI::AddedStructorArgs CIRGenItaniumCXXABI::getImplicitConstructorArgs(
@@ -404,12 +428,11 @@ CIRGenCXXABI *clang::CIRGen::CreateCIRGenItaniumCXXABI(CIRGenModule &CGM) {
   switch (CGM.getASTContext().getCXXABIKind()) {
   case TargetCXXABI::GenericItanium:
   case TargetCXXABI::GenericAArch64:
-  case TargetCXXABI::AppleARM64:
-    // TODO: this isn't quite right, clang uses AppleARM64CXXABI which inherits
-    // from ARMCXXABI. We'll have to follow suit.
-    assert(!cir::MissingFeatures::appleArm64CXXABI());
     return new CIRGenItaniumCXXABI(CGM);
 
+  case TargetCXXABI::AppleARM64:
+    return new CIRGenAppleARM64CXXABI(CGM);
+
   default:
     llvm_unreachable("bad or NYI ABI kind");
   }
@@ -2700,4 +2723,56 @@ Address CIRGenItaniumCXXABI::initializeArrayCookie(CIRGenFunction &CGF,
   auto OffsetOp = CGF.getBuilder().getSignedInt(Loc, Offset, /*width=*/32);
   auto DataPtr = CGF.getBuilder().createPtrStride(Loc, CastOp, OffsetOp);
   return Address(DataPtr, NewPtr.getType(), NewPtr.getAlignment());
-}
+}
+
+CharUnits CIRGenARMCXXABI::getArrayCookieSizeImpl(QualType elementType) {
+  // ARM says that the cookie is always:
+  //   struct array_cookie {
+  //     std::size_t element_size; // element_size != 0
+  //     std::size_t element_count;
+  //   };
+  // But the base ABI doesn't give anything an alignment greater than
+  // 8, so we can dismiss this as typical ABI-author blindness to
+  // actual language complexity and round up to the element alignment.
+  return std::max(CharUnits::fromQuantity(2 * CGM.SizeSizeInBytes),
+                  getContext().getTypeAlignInChars(elementType));
+}
+
+Address CIRGenARMCXXABI::initializeArrayCookie(CIRGenFunction &cgf,
+                                               Address newPtr,
+                                               mlir::Value numElements,
+                                               const CXXNewExpr *expr,
+                                               QualType elementType) {
+  assert(requiresArrayCookie(expr));
+
+  // The cookie is always at the start of the buffer.
+  auto cookiePtr =
+      cgf.getBuilder().createPtrBitcast(newPtr.getPointer(), cgf.SizeTy);
+  Address cookie = Address(cookiePtr, cgf.SizeTy, newPtr.getAlignment());
+
+  ASTContext &ctx = getContext();
+  CharUnits sizeSize = cgf.getSizeSize();
+  mlir::Location loc = cgf.getLoc(expr->getSourceRange());
+
+  // The first element is the element size.
+  mlir::Value elementSize = cgf.getBuilder().getConstInt(
+      loc, cgf.SizeTy, ctx.getTypeSizeInChars(elementType).getQuantity());
+  cgf.getBuilder().createStore(loc, elementSize, cookie);
+
+  // The second element is the element count.
+  auto offsetOp = cgf.getBuilder().getSignedInt(loc, 1, /*width=*/32);
+  auto dataPtr =
+      cgf.getBuilder().createPtrStride(loc, cookie.getPointer(), offsetOp);
+  cookie = Address(dataPtr, cgf.SizeTy, newPtr.getAlignment());
+  cgf.getBuilder().createStore(loc, numElements, cookie);
+
+  // Finally, compute a pointer to the actual data buffer by skipping
+  // over the cookie completely.
+  CharUnits cookieSize = CIRGenARMCXXABI::getArrayCookieSizeImpl(elementType);
+  offsetOp = cgf.getBuilder().getSignedInt(loc, cookieSize.getQuantity(),
+                                           /*width=*/32);
+  auto castOp = cgf.getBuilder().createPtrBitcast(
+      newPtr.getPointer(), cgf.getBuilder().getUIntNTy(8));
+  dataPtr = cgf.getBuilder().createPtrStride(loc, castOp, offsetOp);
+  return Address(dataPtr, newPtr.getType(), newPtr.getAlignment());
+}

clang/test/CIR/CodeGen/applearm64-array-cookies.cpp

@@ -0,0 +1,54 @@
+// RUN: %clang_cc1 -std=c++20 -triple=arm64e-apple-darwin -fclangir -emit-cir %s -o %t.cir
+// RUN: FileCheck --input-file=%t.cir %s
+
+class C {
+  public:
+    ~C();
+};
+
+void t_constant_size_nontrivial() {
+  auto p = new C[3];
+}
+
+// CHECK:  cir.func @_Z26t_constant_size_nontrivialv()
+// CHECK:    %[[#NUM_ELEMENTS:]] = cir.const #cir.int<3> : !u64i
+// CHECK:    %[[#SIZE_WITHOUT_COOKIE:]] = cir.const #cir.int<3> : !u64i
+// CHECK:    %[[#ALLOCATION_SIZE:]] = cir.const #cir.int<19> : !u64i
+// CHECK:    %[[#ALLOC_PTR:]] = cir.call @_Znam(%[[#ALLOCATION_SIZE]]) : (!u64i) -> !cir.ptr<!void>
+// CHECK:    %[[#COOKIE_PTR:]] = cir.cast(bitcast, %[[#ALLOC_PTR]] : !cir.ptr<!void>), !cir.ptr<!u64i>
+// CHECK:    %[[#ELEMENT_SIZE:]] = cir.const #cir.int<1> : !u64i
+// CHECK:    cir.store %[[#ELEMENT_SIZE]], %[[#COOKIE_PTR]] : !u64i, !cir.ptr<!u64i>
+// CHECK:    %[[#SECOND_COOKIE_OFFSET:]] = cir.const #cir.int<1> : !s32i
+// CHECK:    %[[#COOKIE_PTR2:]] = cir.ptr_stride(%[[#COOKIE_PTR]] : !cir.ptr<!u64i>, %[[#SECOND_COOKIE_OFFSET]] : !s32i), !cir.ptr<!u64i>
+// CHECK:    cir.store %[[#NUM_ELEMENTS]], %[[#COOKIE_PTR2]] : !u64i, !cir.ptr<!u64i>
+// CHECK:    %[[#COOKIE_SIZE:]] = cir.const #cir.int<16> : !s32i
+// CHECK:    %[[#ALLOC_AS_I8:]] = cir.cast(bitcast, %[[#ALLOC_PTR]] : !cir.ptr<!void>), !cir.ptr<!u8i>
+// CHECK:    cir.ptr_stride(%[[#ALLOC_AS_I8]] : !cir.ptr<!u8i>, %[[#COOKIE_SIZE]] : !s32i), !cir.ptr<!u8i>
+
+class D {
+  public:
+    int x;
+    ~D();
+};
+
+void t_constant_size_nontrivial2() {
+  auto p = new D[3];
+}
+
+// In this test SIZE_WITHOUT_COOKIE isn't used, but it would be if there were
+// an initializer.
+
+// CHECK:  cir.func @_Z27t_constant_size_nontrivial2v()
+// CHECK:    %[[#NUM_ELEMENTS:]] = cir.const #cir.int<3> : !u64i
+// CHECK:    %[[#SIZE_WITHOUT_COOKIE:]] = cir.const #cir.int<12> : !u64i
+// CHECK:    %[[#ALLOCATION_SIZE:]] = cir.const #cir.int<28> : !u64i
+// CHECK:    %[[#ALLOC_PTR:]] = cir.call @_Znam(%[[#ALLOCATION_SIZE]]) : (!u64i) -> !cir.ptr<!void>
+// CHECK:    %[[#COOKIE_PTR:]] = cir.cast(bitcast, %[[#ALLOC_PTR]] : !cir.ptr<!void>), !cir.ptr<!u64i>
+// CHECK:    %[[#ELEMENT_SIZE:]] = cir.const #cir.int<4> : !u64i
+// CHECK:    cir.store %[[#ELEMENT_SIZE]], %[[#COOKIE_PTR]] : !u64i, !cir.ptr<!u64i>
+// CHECK:    %[[#SECOND_COOKIE_OFFSET:]] = cir.const #cir.int<1> : !s32i
+// CHECK:    %[[#COOKIE_PTR2:]] = cir.ptr_stride(%[[#COOKIE_PTR]] : !cir.ptr<!u64i>, %[[#SECOND_COOKIE_OFFSET]] : !s32i), !cir.ptr<!u64i>
+// CHECK:    cir.store %[[#NUM_ELEMENTS]], %[[#COOKIE_PTR2]] : !u64i, !cir.ptr<!u64i>
+// CHECK:    %[[#COOKIE_SIZE:]] = cir.const #cir.int<16> : !s32i
+// CHECK:    %[[#ALLOC_AS_I8:]] = cir.cast(bitcast, %[[#ALLOC_PTR]] : !cir.ptr<!void>), !cir.ptr<!u8i>
+// CHECK:    cir.ptr_stride(%[[#ALLOC_AS_I8]] : !cir.ptr<!u8i>, %[[#COOKIE_SIZE]] : !s32i), !cir.ptr<!u8i>

clang/test/CIR/Lowering/applearm64-new.cpp

@@ -0,0 +1,41 @@
+// RUN: %clang_cc1 -triple=arm64e-apple-darwin -fclangir -emit-llvm %s -o %t.ll
+// RUN: FileCheck --input-file=%t.ll %s -check-prefix=LLVM
+
+class C {
+  public:
+    ~C();
+};
+
+void t_constant_size_nontrivial() {
+  auto p = new C[3];
+}
+
+// Note: The below differs from the IR emitted by clang without -fclangir in
+//       several respects. (1) The alloca here has an extra "i64 1"
+//       (2) The operator new call is missing "noalias noundef nonnull" on
+//       the call and "noundef" on the argument, (3) The getelementptr is
+//       missing "inbounds"
+
+// LLVM: @_Z26t_constant_size_nontrivialv()
+// LLVM:   %[[COOKIE_PTR:.*]] = call ptr @_Znam(i64 19)
+// LLVM:   store i64 1, ptr %[[COOKIE_PTR]], align 8
+// LLVM:   %[[NUM_ELEMENTS_PTR:.*]] = getelementptr i64, ptr %[[COOKIE_PTR]], i64 1
+// LLVM:   store i64 3, ptr %[[NUM_ELEMENTS_PTR]], align 8
+// LLVM:   %[[ALLOCATED_PTR:.*]] = getelementptr i8, ptr %[[COOKIE_PTR]], i64 16
+
+class D {
+  public:
+    int x;
+    ~D();
+};
+
+void t_constant_size_nontrivial2() {
+  auto p = new D[3];
+}
+
+// LLVM: @_Z27t_constant_size_nontrivial2v()
+// LLVM:   %[[COOKIE_PTR:.*]] = call ptr @_Znam(i64 28)
+// LLVM:   store i64 4, ptr %[[COOKIE_PTR]], align 8
+// LLVM:   %[[NUM_ELEMENTS_PTR:.*]] = getelementptr i64, ptr %[[COOKIE_PTR]], i64 1
+// LLVM:   store i64 3, ptr %[[NUM_ELEMENTS_PTR]], align 8
+// LLVM:   %[[ALLOCATED_PTR:.*]] = getelementptr i8, ptr %[[COOKIE_PTR]], i64 16