Initial support for importing address diversified pointer auth qualified field function pointers

include/swift/AST/Decl.h

@@ -613,13 +613,16 @@ class alignas(1 << DeclAlignInBits) Decl : public ASTAllocated<Decl> {
     IsActor : 1
   );
 
-  SWIFT_INLINE_BITFIELD(
-      StructDecl, NominalTypeDecl, 1 + 1,
-      /// True if this struct has storage for fields that aren't accessible in
-      /// Swift.
-      HasUnreferenceableStorage : 1,
-      /// True if this struct is imported from C++ and does not have trivial value witness functions.
-      IsCxxNonTrivial : 1);
+  SWIFT_INLINE_BITFIELD(StructDecl, NominalTypeDecl, 1 + 1 + 1,
+                        /// True if this struct has storage for fields that
+                        /// aren't accessible in Swift.
+                        HasUnreferenceableStorage : 1,
+                        /// True if this struct is imported from C++ and does
+                        /// not have trivial value witness functions.
+                        IsCxxNonTrivial : 1,
+                        /// True if this struct is imported from C and has
+                        /// address diversified ptrauth qualified field.
+                        IsNonTrivialPtrAuth : 1);
 
   SWIFT_INLINE_BITFIELD(EnumDecl, NominalTypeDecl, 2+1,
     /// True if the enum has cases and at least one case has associated values.
@@ -4267,6 +4270,14 @@ class StructDecl final : public NominalTypeDecl {
 
   void setIsCxxNonTrivial(bool v) { Bits.StructDecl.IsCxxNonTrivial = v; }
 
+  bool isNonTrivialPtrAuth() const {
+    return Bits.StructDecl.IsNonTrivialPtrAuth;
+  }
+
+  void setHasNonTrivialPtrAuth(bool v) {
+    Bits.StructDecl.IsNonTrivialPtrAuth = v;
+  }
+
   Type getTemplateInstantiationType() const { return TemplateInstantiationType; }
   void setTemplateInstantiationType(Type t) { TemplateInstantiationType = t; }
 };

lib/ClangImporter/ImportDecl.cpp

@@ -2024,26 +2024,45 @@ namespace {
         return nullptr;
       }
 
+      auto isNonTrivialDueToAddressDiversifiedPtrAuth =
+          [](const clang::RecordDecl *decl) {
+            for (auto *field : decl->fields()) {
+              if (!field->getType().isNonTrivialToPrimitiveCopy()) {
+                continue;
+              }
+              if (field->getType().isNonTrivialToPrimitiveCopy() !=
+                  clang::QualType::PCK_PtrAuth) {
+                return false;
+              }
+            }
+            return true;
+          };
+
+      bool isNonTrivialPtrAuth = false;
       // FIXME: We should actually support strong ARC references and similar in
       // C structs. That'll require some SIL and IRGen work, though.
       if (decl->isNonTrivialToPrimitiveCopy() ||
           decl->isNonTrivialToPrimitiveDestroy()) {
-        // Note that there is a third predicate related to these,
-        // isNonTrivialToPrimitiveDefaultInitialize. That one's not important
-        // for us because Swift never "trivially default-initializes" a struct
-        // (i.e. uses whatever bits were lying around as an initial value).
-
-        // FIXME: It would be nice to instead import the declaration but mark
-        // it as unavailable, but then it might get used as a type for an
-        // imported function and the developer would be able to use it without
-        // referencing the name, which would sidestep our availability
-        // diagnostics.
-        Impl.addImportDiagnostic(
-            decl, Diagnostic(
-                      diag::record_non_trivial_copy_destroy,
-                      Impl.SwiftContext.AllocateCopy(decl->getNameAsString())),
-            decl->getLocation());
-        return nullptr;
+        isNonTrivialPtrAuth = isNonTrivialDueToAddressDiversifiedPtrAuth(decl);
+        if (!isNonTrivialPtrAuth) {
+          // Note that there is a third predicate related to these,
+          // isNonTrivialToPrimitiveDefaultInitialize. That one's not important
+          // for us because Swift never "trivially default-initializes" a struct
+          // (i.e. uses whatever bits were lying around as an initial value).
+
+          // FIXME: It would be nice to instead import the declaration but mark
+          // it as unavailable, but then it might get used as a type for an
+          // imported function and the developer would be able to use it without
+          // referencing the name, which would sidestep our availability
+          // diagnostics.
+          Impl.addImportDiagnostic(
+              decl,
+              Diagnostic(
+                  diag::record_non_trivial_copy_destroy,
+                  Impl.SwiftContext.AllocateCopy(decl->getNameAsString())),
+              decl->getLocation());
+          return nullptr;
+        }
       }
 
       // Import the name.
@@ -2346,8 +2365,12 @@ namespace {
         }
       }
 
-      if (auto structResult = dyn_cast<StructDecl>(result))
+      if (auto structResult = dyn_cast<StructDecl>(result)) {
         structResult->setHasUnreferenceableStorage(hasUnreferenceableStorage);
+        if (isNonTrivialPtrAuth) {
+          structResult->setHasNonTrivialPtrAuth(true);
+        }
+      }
 
       if (cxxRecordDecl) {
         if (auto structResult = dyn_cast<StructDecl>(result))

lib/IRGen/Callee.h

@@ -78,8 +78,9 @@ namespace irgen {
                                 llvm::Value *storageAddress,
                                 const PointerAuthEntity &entity);
 
-    static PointerAuthInfo emit(IRGenModule &IGM,
-                                clang::PointerAuthQualifier pointerAuthQual);
+    static PointerAuthInfo emit(IRGenFunction &IGF,
+                                clang::PointerAuthQualifier pointerAuthQual,
+                                llvm::Value *storageAddress);
 
     static PointerAuthInfo forFunctionPointer(IRGenModule &IGM,
                                               CanSILFunctionType fnType);

lib/IRGen/GenPointerAuth.cpp

@@ -239,12 +239,30 @@ PointerAuthInfo PointerAuthInfo::emit(IRGenFunction &IGF,
 }
 
 PointerAuthInfo
-PointerAuthInfo::emit(IRGenModule &IGM,
-                      clang::PointerAuthQualifier pointerAuthQual) {
-  return PointerAuthInfo(
-      pointerAuthQual.getKey(),
-      llvm::ConstantInt::get(IGM.Int64Ty,
-                             pointerAuthQual.getExtraDiscriminator()));
+PointerAuthInfo::emit(IRGenFunction &IGF,
+                      clang::PointerAuthQualifier pointerAuthQual,
+                      llvm::Value *storageAddress) {
+  unsigned key = pointerAuthQual.getKey();
+
+  // Produce the 'other' discriminator.
+  auto otherDiscriminator = pointerAuthQual.getExtraDiscriminator();
+  llvm::Value *discriminator =
+      llvm::ConstantInt::get(IGF.IGM.Int64Ty, otherDiscriminator);
+
+  // Factor in the address.
+  if (pointerAuthQual.isAddressDiscriminated()) {
+    assert(storageAddress &&
+           "no storage address for address-discriminated schema");
+
+    if (otherDiscriminator != 0) {
+      discriminator = emitPointerAuthBlend(IGF, storageAddress, discriminator);
+    } else {
+      discriminator =
+          IGF.Builder.CreatePtrToInt(storageAddress, IGF.IGM.Int64Ty);
+    }
+  }
+
+  return PointerAuthInfo(key, discriminator);
 }
 
 llvm::ConstantInt *

lib/IRGen/GenStruct.cpp

@@ -415,9 +415,9 @@ namespace {
     }
   };
 
-  class AddressOnlyClangRecordTypeInfo final
-      : public StructTypeInfoBase<AddressOnlyClangRecordTypeInfo, FixedTypeInfo,
-                                  ClangFieldInfo> {
+  class AddressOnlyCXXClangRecordTypeInfo final
+      : public StructTypeInfoBase<AddressOnlyCXXClangRecordTypeInfo,
+                                  FixedTypeInfo, ClangFieldInfo> {
     const clang::RecordDecl *ClangDecl;
 
     const clang::CXXConstructorDecl *findCopyConstructor() const {
@@ -490,10 +490,10 @@ namespace {
     }
 
   public:
-    AddressOnlyClangRecordTypeInfo(ArrayRef<ClangFieldInfo> fields,
-                                   llvm::Type *storageType, Size size,
-                                   Alignment align,
-                                   const clang::RecordDecl *clangDecl)
+    AddressOnlyCXXClangRecordTypeInfo(ArrayRef<ClangFieldInfo> fields,
+                                      llvm::Type *storageType, Size size,
+                                      Alignment align,
+                                      const clang::RecordDecl *clangDecl)
         : StructTypeInfoBase(StructTypeInfoKind::AddressOnlyClangRecordTypeInfo,
                              fields, storageType, size,
                              // We can't assume any spare bits in a C++ type
@@ -513,7 +513,7 @@ namespace {
       if (!destructor || destructor->isTrivial()) {
         // If we didn't find a destructor to call, bail out to the parent
         // implementation.
-        StructTypeInfoBase<AddressOnlyClangRecordTypeInfo, FixedTypeInfo,
+        StructTypeInfoBase<AddressOnlyCXXClangRecordTypeInfo, FixedTypeInfo,
                            ClangFieldInfo>::destroy(IGF, address, T,
                                                     isOutlined);
         return;
@@ -595,7 +595,7 @@ namespace {
                                     destAddr.getAddress());
         return;
       }
-      StructTypeInfoBase<AddressOnlyClangRecordTypeInfo, FixedTypeInfo,
+      StructTypeInfoBase<AddressOnlyCXXClangRecordTypeInfo, FixedTypeInfo,
                          ClangFieldInfo>::initializeWithCopy(IGF, destAddr,
                                                              srcAddr, T,
                                                              isOutlined);
@@ -609,10 +609,9 @@ namespace {
                                     destAddr.getAddress());
         return;
       }
-      StructTypeInfoBase<AddressOnlyClangRecordTypeInfo, FixedTypeInfo,
-                         ClangFieldInfo>::assignWithCopy(IGF, destAddr,
-                                                         srcAddr, T,
-                                                         isOutlined);
+      StructTypeInfoBase<AddressOnlyCXXClangRecordTypeInfo, FixedTypeInfo,
+                         ClangFieldInfo>::assignWithCopy(IGF, destAddr, srcAddr,
+                                                         T, isOutlined);
     }
 
     void initializeWithTake(IRGenFunction &IGF, Address dest, Address src,
@@ -625,9 +624,8 @@ namespace {
         return;
       }
 
-      StructTypeInfoBase<AddressOnlyClangRecordTypeInfo, FixedTypeInfo,
-                         ClangFieldInfo>::initializeWithTake(IGF, dest,
-                                                             src, T,
+      StructTypeInfoBase<AddressOnlyCXXClangRecordTypeInfo, FixedTypeInfo,
+                         ClangFieldInfo>::initializeWithTake(IGF, dest, src, T,
                                                              isOutlined);
     }
 
@@ -641,9 +639,8 @@ namespace {
         return;
       }
 
-      StructTypeInfoBase<AddressOnlyClangRecordTypeInfo, FixedTypeInfo,
-                         ClangFieldInfo>::assignWithTake(IGF, dest,
-                                                         src, T,
+      StructTypeInfoBase<AddressOnlyCXXClangRecordTypeInfo, FixedTypeInfo,
+                         ClangFieldInfo>::assignWithTake(IGF, dest, src, T,
                                                          isOutlined);
     }
 
@@ -1045,9 +1042,11 @@ class ClangRecordLowering {
   const TypeInfo *createTypeInfo(llvm::StructType *llvmType) {
     llvmType->setBody(LLVMFields, /*packed*/ true);
     if (SwiftType.getStructOrBoundGenericStruct()->isCxxNonTrivial()) {
-      return AddressOnlyClangRecordTypeInfo::create(
+      return AddressOnlyCXXClangRecordTypeInfo::create(
           FieldInfos, llvmType, TotalStride, TotalAlignment, ClangDecl);
     }
+    // TODO: New AddressOnlyPtrAuthTypeInfo to handle address diversified field
+    // function ptrs in C structs.
     return LoadableClangRecordTypeInfo::create(
         FieldInfos, NextExplosionIndex, llvmType, TotalStride,
         std::move(SpareBits), TotalAlignment, ClangDecl);
@@ -1265,8 +1264,8 @@ class ClangRecordLowering {
     case StructTypeInfoKind::LoadableClangRecordTypeInfo:                      \
       return structTI.as<LoadableClangRecordTypeInfo>().op(IGF, __VA_ARGS__);  \
     case StructTypeInfoKind::AddressOnlyClangRecordTypeInfo:                   \
-      return structTI.as<AddressOnlyClangRecordTypeInfo>().op(IGF,             \
-                                                              __VA_ARGS__);    \
+      return structTI.as<AddressOnlyCXXClangRecordTypeInfo>().op(IGF,          \
+                                                                 __VA_ARGS__); \
     case StructTypeInfoKind::LoadableStructTypeInfo:                           \
       return structTI.as<LoadableStructTypeInfo>().op(IGF, __VA_ARGS__);       \
     case StructTypeInfoKind::FixedStructTypeInfo:                              \

lib/IRGen/IRGenSIL.cpp

@@ -5811,7 +5811,8 @@ void IRGenSILFunction::visitBeginAccessInst(BeginAccessInst *access) {
       auto *signedFptr = Builder.CreateLoad(pointerToIntPtr, Int64PtrTy,
                                             IGM.getPointerAlignment());
       auto *resignedFptr = emitPointerAuthResign(
-          *this, signedFptr, PointerAuthInfo::emit(IGM, pointerAuthQual),
+          *this, signedFptr,
+          PointerAuthInfo::emit(*this, pointerAuthQual, pointerToSignedFptr),
           PointerAuthInfo::emit(*this,
                                 IGM.getOptions().PointerAuth.FunctionPointers,
                                 pointerToSignedFptr, PointerAuthEntity()));
@@ -5926,6 +5927,8 @@ void IRGenSILFunction::visitEndAccessInst(EndAccessInst *i) {
     auto pointerAuthQual = cast<StructElementAddrInst>(access->getOperand())
                                ->getField()
                                ->getPointerAuthQualifier();
+    auto *pointerToSignedFptr =
+        getLoweredAddress(access->getOperand()).getAddress();
     auto tempAddress = getLoweredAddress(access);
     auto *tempAddressToIntPtr =
         Builder.CreateBitCast(tempAddress.getAddress(), Int64PtrPtrTy);
@@ -5936,10 +5939,8 @@ void IRGenSILFunction::visitEndAccessInst(EndAccessInst *i) {
         PointerAuthInfo::emit(*this,
                               IGM.getOptions().PointerAuth.FunctionPointers,
                               tempAddress.getAddress(), PointerAuthEntity()),
-        PointerAuthInfo::emit(IGM, pointerAuthQual));
+        PointerAuthInfo::emit(*this, pointerAuthQual, pointerToSignedFptr));
 
-    auto *pointerToSignedFptr =
-        getLoweredAddress(access->getOperand()).getAddress();
     auto *pointerToIntPtr =
         Builder.CreateBitCast(pointerToSignedFptr, Int64PtrPtrTy);
     Builder.CreateStore(signedFptr, pointerToIntPtr, IGM.getPointerAlignment());

lib/SIL/IR/TypeLowering.cpp

@@ -2271,6 +2271,10 @@ namespace {
 
       // Classify the type according to its stored properties.
       for (auto field : D->getStoredProperties()) {
+        auto pointerAuthQual = field->getPointerAuthQualifier();
+        if (pointerAuthQual && pointerAuthQual.isAddressDiscriminated()) {
+          properties.setAddressOnly();
+        }
         auto substFieldType =
           field->getInterfaceType().subst(subMap)
                ->getCanonicalType();

lib/SILGen/SILGenLValue.cpp

@@ -854,7 +854,7 @@ namespace {
       auto Res = SGF.B.createStructElementAddr(loc, base.getValue(),
                                                Field, SubstFieldType);
 
-      if (!Field->getPointerAuthQualifier() ||
+      if (!Field->getPointerAuthQualifier().isPresent() ||
           !SGF.getOptions().EnableImportPtrauthFieldFunctionPointers) {
         return ManagedValue::forLValue(Res);
       }

test/IRGen/Inputs/ptrauth_field_fptr_import.h

@@ -2,9 +2,16 @@
 #define TEST_C_FUNCTION
 
 int returnInt() { return 111; }
+
 struct SecureStruct {
-  int (*__ptrauth(2, 0, 88)(secure_func_ptr))();
+  int (*__ptrauth(1, 0, 88)(secure_func_ptr))();
+};
+
+struct AddressDiscriminatedSecureStruct {
+  int (*__ptrauth(1, 1, 88)(secure_func_ptr))();
 };
 
 struct SecureStruct *ptr_to_secure_struct;
+struct AddressDiscriminatedSecureStruct
+    *ptr_to_addr_discriminated_secure_struct;
 #endif