[TaskLocals] property wrapper keys

Author: ktoso

include/swift/ABI/MetadataKind.def

@@ -93,6 +93,10 @@ METADATAKIND(Task,
 METADATAKIND(Job,
              3 | MetadataKindIsNonType | MetadataKindIsRuntimePrivate)
 
+/// A heap-allocated task-local value `TaskLocal::Item`.
+METADATAKIND(TaskLocalHeapItem,
+             4 | MetadataKindIsNonType | MetadataKindIsRuntimePrivate)
+
 // getEnumeratedMetadataKind assumes that all the enumerated values here
 // will be <= LastEnumeratedMetadataKind.
 

include/swift/ABI/Task.h

@@ -50,7 +50,7 @@ class alignas(2 * alignof(void*)) Job :
     NextWaitingTaskIndex = 0,
 
     // The Dispatch object header is one pointer and two ints, which is
-    // equvialent to three pointers on 32-bit and two pointers 64-bit. Set the
+    // equivalent to three pointers on 32-bit and two pointers 64-bit. Set the
     // indexes accordingly so that DispatchLinkageIndex points to where Dispatch
     // expects.
     DispatchHasLongObjectHeader = sizeof(void *) == sizeof(int),
@@ -243,14 +243,14 @@ class AsyncTask : public Job {
 
   // ==== Task Local Values ----------------------------------------------------
 
-  void localValuePush(const Metadata *keyType,
+  void localValuePush(const HeapObject *key,
                       /* +1 */ OpaqueValue *value, const Metadata *valueType) {
-    Local.pushValue(this, keyType, value, valueType);
+    Local.pushValue(this, key, value, valueType);
   }
 
-  OpaqueValue* localValueGet(const Metadata *keyType,
+  OpaqueValue* localValueGet(const HeapObject *key,
                              TaskLocal::TaskLocalInheritance inherit) {
-    return Local.getValue(this, keyType, inherit);
+    return Local.getValue(this, key, inherit);
   }
 
   void localValuePop() {

include/swift/ABI/TaskLocal.h

@@ -28,27 +28,27 @@ struct SwiftError;
 class TaskStatusRecord;
 class TaskGroup;
 
+extern FullMetadata<HeapMetadata> taskLocalHeapItemHeapMetadata;
+
 // ==== Task Locals Values ---------------------------------------------------
 
 class TaskLocal {
 public:
   /// Type of the pointed at `next` task local item.
   enum class NextLinkType : uintptr_t {
-    /// This task is known to be a "terminal" node in the lookup of task locals.
-    /// In other words, even if it had a parent, the parent (and its parents)
-    /// are known to not contain any any more task locals, and thus any further
-    /// search beyond this task.
-    IsTerminal = 0b00,
     /// The storage pointer points at the next TaskLocal::Item in this task.
-    IsNext = 0b01,
+    IsNext = 0b00,
     /// The storage pointer points at a item stored by another AsyncTask.
     ///
     /// Note that this may not necessarily be the same as the task's parent
     /// task -- we may point to a super-parent if we know / that the parent
     /// does not "contribute" any task local values. This is to speed up
     /// lookups by skipping empty parent tasks during get(), and explained
     /// in depth in `createParentLink`.
-    IsParent = 0b11
+    IsParent = 0b01,
+
+    /// The next item is "unstructured" and must be ref-counted.
+    ThisItemIsHeapItem = 0b10
   };
 
   /// Values must match `TaskLocalInheritance` declared in `TaskLocal.swift`.
@@ -70,37 +70,41 @@ class TaskLocal {
     /// - specifically from the current task itself
     ///
     /// Such values are *not* accessible from child tasks or detached tasks.
-    Never   = 1
+    Never = 1
   };
 
   class Item {
   private:
     /// Mask used for the low status bits in a task local chain item.
     static const uintptr_t statusMask = 0x03;
 
-    /// Pointer to the next task local item; be it in this task or in a parent.
-    /// Low bits encode `NextLinkType`.
-    /// Item *next = nullptr;
+    /// Pointer to one of the following:
+    /// - next task local item as OpaqueValue* if it is task-local allocated
+    /// - next task local item as HeapObject* if it is heap allocated "heavy"
+    /// - the parent task's TaskLocal::Storage
+    ///
+    /// Low bits encode `NextLinkType`, based on which the type of the pointer
+    /// is determined.
     uintptr_t next;
 
   public:
     /// The type of the key with which this value is associated.
-    const Metadata *keyType;
+    const HeapObject *key;
     /// The type of the value stored by this item.
     const Metadata *valueType;
 
     // Trailing storage for the value itself. The storage will be
     // uninitialized or contain an instance of \c valueType.
 
-  private:
+  protected:
     explicit Item()
       : next(0),
-        keyType(nullptr),
+        key(nullptr),
         valueType(nullptr) {}
 
-    explicit Item(const Metadata *keyType, const Metadata *valueType)
+    explicit Item(const HeapObject *key, const Metadata *valueType)
       : next(0),
-        keyType(keyType),
+        key(key),
         valueType(valueType) {}
 
   public:
@@ -116,7 +120,7 @@ class TaskLocal {
     static Item *createParentLink(AsyncTask *task, AsyncTask *parent);
 
     static Item *createLink(AsyncTask *task,
-                            const Metadata *keyType,
+                            const HeapObject *key,
                             const Metadata *valueType);
 
     void destroy(AsyncTask *task);
@@ -125,25 +129,32 @@ class TaskLocal {
       return reinterpret_cast<Item *>(next & ~statusMask);
     }
 
-    NextLinkType getNextLinkType() {
+    NextLinkType getNextLinkType() const {
       return static_cast<NextLinkType>(next & statusMask);
     }
 
+    bool isHeapItem() const {
+      return getNextLinkType() == NextLinkType::ThisItemIsHeapItem;
+    }
+
     /// Item does not contain any actual value, and is only used to point at
     /// a specific parent item.
-    bool isEmpty() {
+    bool isEmpty() const {
       return !valueType;
     }
 
     /// Retrieve a pointer to the storage of the value.
     OpaqueValue *getStoragePtr() {
+      fprintf(stderr, "[%s:%d] (%s) GET STORAGE\n", __FILE__, __LINE__, __FUNCTION__);
       return reinterpret_cast<OpaqueValue *>(
-        reinterpret_cast<char *>(this) + storageOffset(valueType));
+        reinterpret_cast<char *>(this) + storageOffset(isHeapItem(), valueType));
     }
 
     /// Compute the offset of the storage from the base of the item.
-    static size_t storageOffset(const Metadata *valueType) {
+    static size_t storageOffset(bool isHeapItem, const Metadata *valueType) {
       size_t offset = sizeof(Item);
+
+
       if (valueType) {
         size_t alignment = valueType->vw_alignment();
         return (offset + alignment - 1) & ~(alignment - 1);
@@ -153,15 +164,53 @@ class TaskLocal {
     }
 
     /// Determine the size of the item given a particular value type.
-    static size_t itemSize(const Metadata *valueType) {
-      size_t offset = storageOffset(valueType);
+    static size_t itemSize(bool isHeapItem, const Metadata *valueType) {
+      size_t offset = storageOffset(isHeapItem, valueType);
+      size_t headerSize = isHeapItem ? sizeof(HeapObject) : 0;
+      offset += headerSize;
       if (valueType) {
         offset += valueType->vw_size();
       }
       return offset;
     }
   };
 
+  /// Heap allocated version of TaskLocal::Item.
+  /// It is used in limited un-structured concurrency cases, where the task
+  /// local must be stored on the heap for the time being.
+  ///
+  /// MUST HAVE THE SAME LAYOUT AS TaskLocal::Item.
+  /// FIXME: how to actually make this work...
+  ///        get the address of "after heap object" and cast that to Item?
+  class HeapItem : public HeapObject, public Item {
+  public:
+
+    explicit HeapItem(const HeapMetadata *metadata = &taskLocalHeapItemHeapMetadata)
+        : HeapObject(metadata),
+          Item() {}
+
+    explicit HeapItem(const HeapObject *key, const Metadata *valueType)
+        : HeapItem(&taskLocalHeapItemHeapMetadata, key, valueType) {}
+
+    explicit HeapItem(const HeapMetadata *metadata,
+                      const HeapObject *key, const Metadata *valueType)
+        : HeapObject(metadata),
+          Item(key, valueType) {}
+
+
+    /// Retrieve a pointer to the `Item` that is part of this object.
+    TaskLocal::Item *getItem() {
+      fprintf(stderr, "[%s:%d] (%s) GET ITEM PTR, heapItem:%p\n", __FILE__, __LINE__, __FUNCTION__, this);
+      return reinterpret_cast<TaskLocal::Item *>(
+          reinterpret_cast<char *>(this) + itemOffset());
+    }
+
+    /// Compute the offset of `Item` part of this object.
+    static size_t itemOffset() {
+      size_t offset = sizeof(HeapObject);
+      return offset;
+    }
+  };
 
   class Storage {
     friend class TaskLocal::Item;
@@ -202,11 +251,11 @@ class TaskLocal {
     void initializeLinkParent(AsyncTask *task, AsyncTask *parent);
 
     void pushValue(AsyncTask *task,
-                   const Metadata *keyType,
+                   const HeapObject *key,
                    /* +1 */ OpaqueValue *value, const Metadata *valueType);
 
     OpaqueValue* getValue(AsyncTask *task,
-                          const Metadata *keyType,
+                          const HeapObject *key,
                           TaskLocalInheritance inheritance);
 
     void popValue(AsyncTask *task);

include/swift/AST/KnownSDKTypes.def

@@ -40,4 +40,6 @@ KNOWN_SDK_TYPE_DECL(ObjectiveC, ObjCBool, StructDecl, 0)
 KNOWN_SDK_TYPE_DECL(Concurrency, UnsafeContinuation, NominalTypeDecl, 2)
 KNOWN_SDK_TYPE_DECL(Concurrency, MainActor, NominalTypeDecl, 0)
 
+KNOWN_SDK_TYPE_DECL(Concurrency, TaskLocal, ClassDecl, 1)
+
 #undef KNOWN_SDK_TYPE_DECL

include/swift/Runtime/Concurrency.h

@@ -362,11 +362,20 @@ bool swift_task_tryAddStatusRecord(TaskStatusRecord *record);
 ///
 /// The given record need not be the last record added to
 /// the task, but the operation may be less efficient if not.
-///s
+///
 /// Returns false if the task has been cancelled.
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 bool swift_task_removeStatusRecord(TaskStatusRecord *record);
 
+/// Signifies whether the current task is in the middle of executing the
+/// operation block of a `with(Throwing)TaskGroup(...) { <operation> }`.
+///
+/// Task local values must use un-structured allocation for values bound in this
+/// scope, as they may be referred to by `group.spawn`-ed tasks and therefore
+/// out-life the scope of a task-local value binding.
+SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
+bool swift_task_hasTaskGroupStatusRecord();
+
 /// Attach a child task to its parent task and return the newly created
 /// `ChildTaskStatusRecord`.
 ///
@@ -409,7 +418,7 @@ void swift_task_removeCancellationHandler(
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 OpaqueValue*
 swift_task_localValueGet(AsyncTask* task,
-                         const Metadata *keyType,
+                         const HeapObject *key,
                          TaskLocal::TaskLocalInheritance inheritance);
 
 /// Add a task local value to the passed in task.
@@ -427,7 +436,7 @@ swift_task_localValueGet(AsyncTask* task,
 /// \endcode
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
 void swift_task_localValuePush(AsyncTask* task,
-                         const Metadata *keyType,
+                         const HeapObject *key,
                          /* +1 */ OpaqueValue *value,
                          const Metadata *valueType);
 

stdlib/public/CompatibilityOverride/CompatibilityOverrideConcurrency.def

@@ -221,16 +221,16 @@ OVERRIDE_TASK_GROUP(taskGroup_addPending, bool,
 OVERRIDE_TASK_LOCAL(task_localValuePush, void,
                     SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
                     swift::,
-                    (AsyncTask *task, const Metadata *keyType,
+                    (AsyncTask *task, const HeapObject *key,
                      OpaqueValue *value, const Metadata *valueType),
-                    (task, keyType, value, valueType))
+                    (task, key, value, valueType))
 
 OVERRIDE_TASK_LOCAL(task_localValueGet, OpaqueValue *,
                     SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
                     swift::,
-                    (AsyncTask *task, const Metadata *keyType,
+                    (AsyncTask *task, const HeapObject *key,
                      TaskLocal::TaskLocalInheritance inheritance),
-                    (task, keyType, inheritance))
+                    (task, key, inheritance))
 
 OVERRIDE_TASK_LOCAL(task_localValuePop, void,
                     SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
@@ -248,6 +248,10 @@ OVERRIDE_TASK_STATUS(task_removeStatusRecord, bool,
                      SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
                      swift::, (TaskStatusRecord *record), (record))
 
+OVERRIDE_TASK_STATUS(task_hasTaskGroupStatusRecord, bool,
+                     SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
+                     swift::, , )
+
 OVERRIDE_TASK_STATUS(task_attachChild, ChildTaskStatusRecord *,
                      SWIFT_EXPORT_FROM(swift_Concurrency), SWIFT_CC(swift),
                      swift::, (AsyncTask *child), (child))

stdlib/public/Concurrency/TaskGroup.swift

@@ -215,7 +215,7 @@ public struct TaskGroup<ChildTaskResult: Sendable> {
       priority: Task.Priority = .unspecified,
       operation: __owned @Sendable @escaping () async -> ChildTaskResult
   ) async -> Bool {
-    return try self.spawnUnlessCancelled(priority: priority) {
+    return self.spawnUnlessCancelled(priority: priority) {
       await operation()
     }
   }
@@ -467,7 +467,7 @@ public struct ThrowingTaskGroup<ChildTaskResult: Sendable, Failure: Error> {
     priority: Task.Priority = .unspecified,
     operation: __owned @Sendable @escaping () async throws -> ChildTaskResult
   ) async -> Bool {
-    return try self.spawnUnlessCancelled(priority: priority) {
+    return self.spawnUnlessCancelled(priority: priority) {
       try await operation()
     }
   }