Merge pull request #77270 from kubamracek/embedded-concurrency-arm64e

[Concurrency] Fix arm64e ptrauth violations in Embedded Swift Concurrency

Author: kubamracek

include/swift/ABI/TaskOptions.h

@@ -194,9 +194,17 @@ class ResultTypeInfoTaskOptionRecord : public TaskOptionRecord {
  public:
   size_t size;
   size_t alignMask;
-  void (*initializeWithCopy)(OpaqueValue *, OpaqueValue *);
-  void (*storeEnumTagSinglePayload)(OpaqueValue *, unsigned, unsigned);
-  void (*destroy)(OpaqueValue *);
+
+  void (*__ptrauth_swift_value_witness_function_pointer(
+      SpecialPointerAuthDiscriminators::InitializeWithCopy)
+            initializeWithCopy)(OpaqueValue *, OpaqueValue *);
+
+  void (*__ptrauth_swift_value_witness_function_pointer(
+      SpecialPointerAuthDiscriminators::StoreEnumTagSinglePayload)
+            storeEnumTagSinglePayload)(OpaqueValue *, unsigned, unsigned);
+
+  void (*__ptrauth_swift_value_witness_function_pointer(
+      SpecialPointerAuthDiscriminators::Destroy) destroy)(OpaqueValue *);
 
   static bool classof(const TaskOptionRecord *record) {
     return record->getKind() == TaskOptionRecordKind::ResultTypeInfo;

lib/IRGen/GenConcurrency.cpp

@@ -20,6 +20,7 @@
 #include "BitPatternBuilder.h"
 #include "ExtraInhabitants.h"
 #include "GenCall.h"
+#include "GenPointerAuth.h"
 #include "GenProto.h"
 #include "GenType.h"
 #include "IRGenDebugInfo.h"
@@ -584,22 +585,50 @@ struct EmbeddedSwiftResultTypeOptionRecordTraits {
     IGF.Builder.CreateStore(
         TI.getStaticAlignmentMask(IGF.IGM),
         IGF.Builder.CreateStructGEP(optionsRecord, 2, Size()));
+
+    auto schema = IGF.getOptions().PointerAuth.ValueWitnesses;
     // initializeWithCopy witness
-    IGF.Builder.CreateStore(
-        IGF.IGM.getOrCreateValueWitnessFunction(
-            ValueWitness::InitializeWithCopy, packing, canType, lowered, TI),
-        IGF.Builder.CreateStructGEP(optionsRecord, 3, Size()));
+    {
+      auto gep = IGF.Builder.CreateStructGEP(optionsRecord, 3, Size());
+      llvm::Value *witness = IGF.IGM.getOrCreateValueWitnessFunction(
+          ValueWitness::InitializeWithCopy, packing, canType, lowered, TI);
+      auto discriminator = llvm::ConstantInt::get(
+          IGF.IGM.Int64Ty,
+          SpecialPointerAuthDiscriminators::InitializeWithCopy);
+      auto storageAddress = gep.getAddress();
+      auto info =
+          PointerAuthInfo::emit(IGF, schema, storageAddress, discriminator);
+      if (schema) witness = emitPointerAuthSign(IGF, witness, info);
+      IGF.Builder.CreateStore(witness, gep);
+    }
     // storeEnumTagSinglePayload witness
-    IGF.Builder.CreateStore(
-        IGF.IGM.getOrCreateValueWitnessFunction(
-            ValueWitness::StoreEnumTagSinglePayload, packing, canType, lowered,
-            TI),
-        IGF.Builder.CreateStructGEP(optionsRecord, 4, Size()));
+    {
+      auto gep = IGF.Builder.CreateStructGEP(optionsRecord, 4, Size());
+      llvm::Value *witness = IGF.IGM.getOrCreateValueWitnessFunction(
+          ValueWitness::StoreEnumTagSinglePayload, packing, canType, lowered,
+          TI);
+      auto discriminator = llvm::ConstantInt::get(
+          IGF.IGM.Int64Ty,
+          SpecialPointerAuthDiscriminators::StoreEnumTagSinglePayload);
+      auto storageAddress = gep.getAddress();
+      auto info =
+          PointerAuthInfo::emit(IGF, schema, storageAddress, discriminator);
+      if (schema) witness = emitPointerAuthSign(IGF, witness, info);
+      IGF.Builder.CreateStore(witness, gep);
+    }
     // destroy witness
-    IGF.Builder.CreateStore(
-        IGF.IGM.getOrCreateValueWitnessFunction(ValueWitness::Destroy, packing,
-                                                canType, lowered, TI),
-        IGF.Builder.CreateStructGEP(optionsRecord, 5, Size()));
+    {
+      auto gep = IGF.Builder.CreateStructGEP(optionsRecord, 5, Size());
+      llvm::Value *witness = IGF.IGM.getOrCreateValueWitnessFunction(
+          ValueWitness::Destroy, packing, canType, lowered, TI);
+      auto discriminator = llvm::ConstantInt::get(
+          IGF.IGM.Int64Ty, SpecialPointerAuthDiscriminators::Destroy);
+      auto storageAddress = gep.getAddress();
+      auto info =
+          PointerAuthInfo::emit(IGF, schema, storageAddress, discriminator);
+      if (schema) witness = emitPointerAuthSign(IGF, witness, info);
+      IGF.Builder.CreateStore(witness, gep);
+    }
   }
 };
 } // end anonymous namespace

stdlib/public/Concurrency/Actor.cpp

@@ -1782,9 +1782,14 @@ void DefaultActorImpl::deallocate() {
 void DefaultActorImpl::deallocateUnconditional() {
   concurrency::trace::actor_deallocate(this);
 
+#if !SWIFT_CONCURRENCY_EMBEDDED
   auto metadata = cast<ClassMetadata>(this->metadata);
   swift_deallocClassInstance(this, metadata->getInstanceSize(),
                              metadata->getInstanceAlignMask());
+#else
+  // Embedded Swift's runtime doesn't actually use the size/mask values.
+  swift_deallocClassInstance(this, 0, 0);
+#endif
 }
 
 bool DefaultActorImpl::tryLock(bool asDrainer) {
@@ -2107,12 +2112,16 @@ static bool isDefaultActorClass(const ClassMetadata *metadata) {
 }
 
 void swift::swift_defaultActor_deallocateResilient(HeapObject *actor) {
+#if !SWIFT_CONCURRENCY_EMBEDDED
   auto metadata = cast<ClassMetadata>(actor->metadata);
   if (isDefaultActorClass(metadata))
     return swift_defaultActor_deallocate(static_cast<DefaultActor*>(actor));
 
   swift_deallocObject(actor, metadata->getInstanceSize(),
                       metadata->getInstanceAlignMask());
+#else
+  return swift_defaultActor_deallocate(static_cast<DefaultActor*>(actor));
+#endif
 }
 
 /// FIXME: only exists for the quick-and-dirty MainActor implementation.
@@ -2550,12 +2559,16 @@ swift::swift_distributedActor_remote_initialize(const Metadata *actorType) {
 }
 
 bool swift::swift_distributed_actor_is_remote(HeapObject *_actor) {
+#if !SWIFT_CONCURRENCY_EMBEDDED
   const ClassMetadata *metadata = cast<ClassMetadata>(_actor->metadata);
   if (isDefaultActorClass(metadata)) {
     return asImpl(reinterpret_cast<DefaultActor *>(_actor))->isDistributedRemote();
   } else {
     return asImpl(reinterpret_cast<NonDefaultDistributedActor *>(_actor))->isDistributedRemote();
   }
+#else
+  return false;
+#endif
 }
 
 bool DefaultActorImpl::isDistributedRemote() {