[Concurrency] TaskLocals lookup "skip" optimization

include/swift/ABI/MetadataValues.h

@@ -1938,8 +1938,7 @@ public:
 
     Task_IsChildTask    = 24,
     Task_IsFuture       = 25,
-    Task_IsTaskGroup    = 26,
+    Task_IsTaskGroup    = 26
-    Task_HasLocalValues = 27
   };
 
   explicit JobFlags(size_t bits) : FlagSet(bits) {}
@@ -1969,9 +1968,6 @@ public:
   FLAGSET_DEFINE_FLAG_ACCESSORS(Task_IsTaskGroup,
                                 task_isTaskGroup,
                                 task_setIsTaskGroup)
-  FLAGSET_DEFINE_FLAG_ACCESSORS(Task_HasLocalValues,
-                                task_hasLocalValues,
-                                task_setHasLocalValues)
 };
 
 /// Kinds of task status record.

include/swift/ABI/Task.h

@@ -139,13 +139,13 @@ public:
 ///
 ///    +--------------------------+
 ///    | childFragment?           |
-///    | taskLocalValuesFragment  |
+///    | taskLocalValuesFragment? |
 ///    | groupFragment?           |
 ///    | futureFragment?          |*
 ///    +--------------------------+
 ///
-/// The future fragment is dynamic in size, based on the future result type
+/// * The future fragment is dynamic in size, based on the future result type
-/// it can hold, and thus must be the *last* fragment.
+///   it can hold, and thus must be the *last* fragment.
 class AsyncTask : public HeapObject, public Job {
 public:
   /// The context for resuming the job.  When a task is scheduled
@@ -217,7 +217,7 @@ public:
     return reinterpret_cast<ChildFragment*>(this + 1);
   }
 
-  // ==== Task Locals Values-- -------------------------------------------------
+  // ==== Task Locals Values ---------------------------------------------------
 
   class TaskLocalValuesFragment {
   public:
@@ -230,8 +230,14 @@ public:
       IsTerminal = 0b00,
       /// The storage pointer points at the next TaskLocalChainItem in this task.
       IsNext     = 0b01,
-      /// The storage pointer points at a parent AsyncTask,
+      /// The storage pointer points at a parent AsyncTask, in which we should
-      /// in which we should continue the lookup.
+      /// continue the lookup.
+      ///
+      /// 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
     };
 
@@ -272,23 +278,42 @@ public:
       /// the TaskLocalItem linked list into the appropriate parent.
       static TaskLocalItem* createParentLink(AsyncTask *task, AsyncTask *parent) {
         assert(parent);
-        assert(parent->hasTaskLocalValues());
-        assert(task->hasTaskLocalValues());
         size_t amountToAllocate = TaskLocalItem::itemSize(/*valueType*/nullptr);
         // assert(amountToAllocate % MaximumAlignment == 0); // TODO: do we need this?
         void *allocation = malloc(amountToAllocate); // TODO: use task-local allocator
-        fprintf(stderr, "MALLOC parent link item: %d\n", allocation);
 
         TaskLocalItem *item =
             new(allocation) TaskLocalItem(nullptr, nullptr);
 
-        auto next = parent->localValuesFragment()->head;
+        auto parentHead = parent->localValuesFragment()->head;
-        auto nextLinkType = next ? NextLinkType::IsParent : NextLinkType::IsTerminal;
+        if (parentHead) {
-        item->next = reinterpret_cast<uintptr_t>(next) |
+          if (parentHead->isEmpty()) {
-                          static_cast<uintptr_t>(nextLinkType);
+            switch (parentHead->getNextLinkType()) {
-
+              case NextLinkType::IsParent:
-        fprintf(stderr, "error: %s [%s:%d] created parent item: task=%d -> parentTask=%d  :: item=%d -> item->getNext()=%d\n", __FUNCTION__, __FILE_NAME__, __LINE__,
+                // it has no values, and just points to its parent,
-                task, parent, item, item->getNext());
+                // therefore skip also skip pointing to that parent and point
+                // to whichever parent it was pointing to as well, it may be its
+                // immediate parent, or some super-parent.
+                item->next = reinterpret_cast<uintptr_t>(parentHead->getNext());
+                                  static_cast<uintptr_t>(NextLinkType::IsParent);
+                break;
+              case NextLinkType::IsNext:
+                assert(false && "empty taskValue head in parent task, yet parent's 'head' is `IsNext`, "
+                                "this should not happen, as it implies the parent must have stored some value.");
+                break;
+              case NextLinkType::IsTerminal:
+                item->next = reinterpret_cast<uintptr_t>(parentHead->getNext());
+                                  static_cast<uintptr_t>(NextLinkType::IsTerminal);
+                break;
+            }
+          } else {
+            item->next = reinterpret_cast<uintptr_t>(parentHead) |
+                         static_cast<uintptr_t>(NextLinkType::IsParent);
+          }
+        } else {
+          item->next = reinterpret_cast<uintptr_t>(parentHead) |
+                            static_cast<uintptr_t>(NextLinkType::IsTerminal);
+        }
 
         return item;
       }
@@ -297,7 +322,6 @@ public:
                                        const Metadata *keyType,
                                        const Metadata *valueType) {
         assert(task);
-        assert(task->hasTaskLocalValues());
         size_t amountToAllocate = TaskLocalItem::itemSize(valueType);
         // assert(amountToAllocate % MaximumAlignment == 0); // TODO: do we need this?
         void *allocation = malloc(amountToAllocate); // TODO: use task-local allocator
@@ -327,9 +351,14 @@ public:
         return static_cast<NextLinkType>(next & statusMask);
       }
 
+      /// Item does not contain any actual value, and is only used to point at
+      /// a specific parent item.
+      bool isEmpty() {
+        return !valueType;
+      }
+
       /// Retrieve a pointer to the storage of the value.
       OpaqueValue *getStoragePtr() {
-//        assert(valueType && "valueType must be set before accessing storage pointer.");
         return reinterpret_cast<OpaqueValue *>(
             reinterpret_cast<char *>(this) + storageOffset(valueType));
       }
@@ -356,7 +385,8 @@ public:
     };
 
   private:
-    /// Single-linked list of task local values.
+    /// A stack (single-linked list) of task local values.
+    ///
     /// Once task local values within this task are traversed, the list continues
     /// to the "next parent that contributes task local values," or if no such
     /// parent exists it terminates with null.
@@ -366,11 +396,28 @@ public:
     /// parent that has values. If this task does not have any values, the head
     /// pointer MAY immediately point at this task's parent task which has values.
     ///
-    /// NOTE: Check the higher bits to know if this is a self or parent value.
+    /// ### Concurrency
+    /// Access to the head is only performed from the task itself, when it
+    /// creates child tasks, the child during creation will inspect its parent's
+    /// task local value stack head, and point to it. This is done on the calling
+    /// task, and thus needs not to be synchronized. Subsequent traversal is
+    /// performed by child tasks concurrently, however they use their own
+    /// pointers/stack and can never mutate the parent's stack.
+    ///
+    /// The stack is only pushed/popped by the owning task, at the beginning and
+    /// end a `body` block of `withLocal(_:boundTo:body:)` respectively.
+    ///
+    /// Correctness of the stack strongly relies on the guarantee that tasks
+    /// never outline a scope in which they are created. Thanks to this, if
+    /// tasks are created inside the `body` of `withLocal(_:,boundTo:body:)`
+    /// all tasks created inside the `withLocal` body must complete before it
+    /// returns, as such, any child tasks potentially accessing the value stack
+    /// are guaranteed to be completed by the time we pop values off the stack
+    /// (after the body has completed).
     TaskLocalItem *head = nullptr;
 
   public:
-      TaskLocalValuesFragment() {}
+    TaskLocalValuesFragment() {}
 
     void destroy();
 
@@ -386,13 +433,7 @@ public:
     OpaqueValue* get(const Metadata *keyType);
   };
 
-  bool hasTaskLocalValues() const {
-    return Flags.task_hasLocalValues();
-  }
-
   TaskLocalValuesFragment *localValuesFragment() {
-    assert(hasTaskLocalValues());
-
     auto offset = reinterpret_cast<char*>(this);
     offset += sizeof(AsyncTask);
 
@@ -404,14 +445,7 @@ public:
   }
 
   OpaqueValue* localValueGet(const Metadata *keyType) {
-    if (hasTaskLocalValues()) {
+    return localValuesFragment()->get(keyType);
-      return localValuesFragment()->get(keyType);
-    } else {
-      // We are guaranteed to have a task-local fragment even if this task has
-      // no bindings, but its parent tasks do. Thus, if no fragment, we can
-      // immediately return null.
-      return nullptr;
-    }
   }
 
   // ==== TaskGroup ------------------------------------------------------------
@@ -724,9 +758,7 @@ public:
       offset += sizeof(ChildFragment);
     }
 
-    if (hasTaskLocalValues()) {
+    offset += sizeof(TaskLocalValuesFragment);
-      offset += sizeof(TaskLocalValuesFragment);
-    }
 
     return reinterpret_cast<GroupFragment *>(offset);
   }
@@ -861,9 +893,7 @@ public:
       offset += sizeof(ChildFragment);
     }
 
-    if (hasTaskLocalValues()) {
+    offset += sizeof(TaskLocalValuesFragment);
-      offset += sizeof(TaskLocalValuesFragment);
-    }
 
     if (isTaskGroup()) {
       offset += sizeof(GroupFragment);

include/swift/Runtime/Concurrency.h

@@ -257,8 +257,7 @@ void swift_task_localValuePush(AsyncTask* task,
                          /* +1 */ OpaqueValue *value,
                          const Metadata *valueType);
 
-/// Remove task `count` local bindings from the task local binding stack.
+/// Remove task a local binding from the task local values stack.
-/// Crashes if `count` is greater if the number of task locals stored in the task.
 ///
 /// This must be only invoked by the task itself to avoid concurrent writes.
 ///
@@ -266,22 +265,11 @@ void swift_task_localValuePush(AsyncTask* task,
 ///
 /// \code
 ///  public func _taskLocalValuePop(
-///    _ task: Builtin.NativeObject,
+///    _ task: Builtin.NativeObject
-///    count: Int
 ///  )
 /// \endcode
 SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
-void swift_task_localValuePop(AsyncTask* task, int count);
+void swift_task_localValuePop(AsyncTask* task);
-
-/// Checks if task (or any of its parent tasks) has task local values.
-///
-/// \code
-/// func swift_task_hasTaskLocalValues<Key>(
-///   _ task: Builtin.NativeObject,
-/// ) -> Bool
-/// \endcode
-SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(swift)
-bool swift_task_hasTaskLocalValues(AsyncTask* task);
 
 /// This should have the same representation as an enum like this:
 ///    enum NearestTaskDeadline {

stdlib/public/Concurrency/Task.cpp

@@ -135,11 +135,7 @@ SWIFT_CC(swift)
 static void destroyTask(SWIFT_CONTEXT HeapObject *obj) {
   auto task = static_cast<AsyncTask*>(obj);
 
-  fprintf(stderr, "destroy task (%d): %d\n",
+  // For a group, destroy the queues and results.
-          task->hasTaskLocalValues(), task);
-
-
-// For a group, destroy the queues and results.
   if (task->isTaskGroup()) {
     task->groupFragment()->destroy();
   }
@@ -150,9 +146,7 @@ static void destroyTask(SWIFT_CONTEXT HeapObject *obj) {
   }
 
   // release any objects potentially held as task local values.
-  if (task->hasTaskLocalValues()) {
+  task->localValuesFragment()->destroy();
-    task->localValuesFragment()->destroy();
-  }
 
   // The task execution itself should always hold a reference to it, so
   // if we get here, we know the task has finished running, which means
@@ -242,15 +236,9 @@ AsyncTaskAndContext swift::swift_task_create_future_f(
     headerSize += sizeof(AsyncTask::ChildFragment);
   }
 
-  bool needsTaskLocalsFragment =
+  headerSize += sizeof(AsyncTask::TaskLocalValuesFragment);
-      flags.task_hasLocalValues() || (parent && parent->hasTaskLocalValues());
+  fprintf(stderr, "error: %s [%s:%d] adding values fragment size=%d\n", __FUNCTION__, __FILE_NAME__, __LINE__,
-  fprintf(stderr, "error: %s [%s:%d] prepare task taskHasTaskLocals=%d parentHasLocals=%d needsLocals=%d\n", __FUNCTION__, __FILE_NAME__, __LINE__,
+      sizeof(AsyncTask::TaskLocalValuesFragment));
-          flags.task_hasLocalValues(), (parent && parent->hasTaskLocalValues()), needsTaskLocalsFragment);
-  if (needsTaskLocalsFragment) {
-    headerSize += sizeof(AsyncTask::TaskLocalValuesFragment);
-    fprintf(stderr, "error: %s [%s:%d] adding values fragment size=%d\n", __FUNCTION__, __FILE_NAME__, __LINE__,
-        sizeof(AsyncTask::TaskLocalValuesFragment));
-  }
 
   if (flags.task_isTaskGroup()) {
     headerSize += sizeof(AsyncTask::GroupFragment);
@@ -287,12 +275,9 @@ AsyncTaskAndContext swift::swift_task_create_future_f(
     new (childFragment) AsyncTask::ChildFragment(parent);
   }
 
-  if (needsTaskLocalsFragment) {
+  auto taskLocalsFragment = task->localValuesFragment();
-    assert(task->hasTaskLocalValues());
+  new (taskLocalsFragment) AsyncTask::TaskLocalValuesFragment();
-    auto taskLocalsFragment = task->localValuesFragment();
+  taskLocalsFragment->initializeLinkParent(task, parent);
-    new (taskLocalsFragment) AsyncTask::TaskLocalValuesFragment();
-    taskLocalsFragment->initializeLinkParent(task, parent);
-  }
 
   // Initialize the task group fragment if applicable.
     if (flags.task_isTaskGroup()) {
@@ -494,7 +479,6 @@ void swift::swift_task_runAndBlockThread(const void *function,
 
   // Set up a task that runs the runAndBlock async function above.
   auto flags = JobFlags(JobKind::Task, JobPriority::Default);
-  flags.task_setHasLocalValues(true);
   auto pair = swift_task_create_f(flags,
                                   /*parent*/ nullptr,
                                   &runAndBlock_start,
@@ -515,27 +499,22 @@ size_t swift::swift_task_getJobFlags(AsyncTask *task) {
   return task->Flags.getOpaqueValue();
 }
 
-void swift::swift_task_localValuePush(AsyncTask *task,
+void swift::swift_task_local_value_push(AsyncTask *task,
                        const Metadata *keyType,
                        /* +1 */ OpaqueValue *value, const Metadata *valueType) {
   fprintf(stderr, "error: %s [%s:%d] PUSH keyType=%d value=%d *value=%d\n",
           __FUNCTION__, __FILE_NAME__, __LINE__,
           keyType, value, *reinterpret_cast<int*>(value));
-  assert(task->hasTaskLocalValues());
   task->localValuesFragment()->pushValue(task, keyType, value, valueType);
 }
 
-void swift::swift_task_localValuePop(AsyncTask *task, int count) {
+void swift::swift_task_local_value_pop(AsyncTask *task) {
-  assert(task->hasTaskLocalValues());
+  fprintf(stderr, "error: %s [%s:%d] POP task=%d\n", __FUNCTION__, __FILE_NAME__, __LINE__,
-  auto fragment = task->localValuesFragment();
+          task);
-  for (int i = 0; i < count; i++) {
+  task->localValuesFragment()->popValue(task);
-    fprintf(stderr, "error: %s [%s:%d] POP task=%d %d / %d\n", __FUNCTION__, __FILE_NAME__, __LINE__,
-            task, i, count);
-    fragment->popValue(task);
-  }
 }
 
-OpaqueValue* swift::swift_task_localValueGet(AsyncTask *task,
+OpaqueValue* swift::swift_task_local_value_get(AsyncTask *task,
                                              const Metadata *keyType) {
   auto value = task->localValueGet(keyType);
   fprintf(stderr, "error: %s [%s:%d] lookup keyType=%d value=%d\n",
@@ -544,10 +523,6 @@ OpaqueValue* swift::swift_task_localValueGet(AsyncTask *task,
   return value;
 }
 
-bool swift::swift_task_hasTaskLocalValues(AsyncTask *task) {
-  return task->hasTaskLocalValues();
-}
-
 namespace {
   
 /// Structure that gets filled in when a task is suspended by `withUnsafeContinuation`.

stdlib/public/Concurrency/Task.swift

@@ -12,7 +12,16 @@
 
 import Swift
 @_implementationOnly import _SwiftConcurrencyShims
+#if canImport(Darwin)
 import Darwin
+#elseif canImport(Glibc)
+import Glibc
+#elseif os(Windows)
+import CRT
+#else
+#error("Unsupported platform")
+#endif
+
 
 // ==== Task -------------------------------------------------------------------
 
@@ -646,7 +655,6 @@ public func _runChildTask<T>(
   flags.priority = getJobFlags(currentTask).priority
   flags.isFuture = true
   flags.isChildTask = true
-  flags.hasLocalValues = true // _taskHasTaskLocalValues(currentTask)
 
   // Create the asynchronous task future.
   let (task, _) = Builtin.createAsyncTaskFuture(
@@ -678,7 +686,6 @@ public func _runGroupChildTask<T>(
   flags.priority = priorityOverride ?? getJobFlags(currentTask).priority
   flags.isFuture = true
   flags.isChildTask = true
-  flags.hasLocalValues = true // hasLocalValues || _taskHasTaskLocalValues(currentTask)
 
   // Create the asynchronous task future.
   let (task, _) = Builtin.createAsyncTaskFuture(

stdlib/public/Concurrency/TaskGroup.swift

@@ -74,7 +74,6 @@ extension Task {
     groupFlags.isChildTask = true
     groupFlags.isTaskGroup = true
     groupFlags.isFuture = true
-    groupFlags.hasLocalValues = true // _taskHasTaskLocalValues(parent)
 
     let (groupTask, _) =
       Builtin.createAsyncTaskFuture(groupFlags.bits, parent) { () async throws -> BodyResult in

stdlib/public/Concurrency/TaskLocal.cpp

@@ -55,38 +55,36 @@ void TaskLocalValuesFragment::destroy() {
 void TaskLocalValuesFragment::initializeLinkParent(AsyncTask* task,
                                                    AsyncTask* parent) {
   assert(!head && "fragment was already initialized");
-  if (parent && parent->hasTaskLocalValues()) {
+  if (parent) {
     head = TaskLocalItem::createParentLink(task, parent);
   }
 }
 
 // =============================================================================
-// ==== push / pop / get -----------------------------------------------------------
+// ==== push / pop / get -------------------------------------------------------
 
 void TaskLocalValuesFragment::pushValue(AsyncTask *task,
                                         const Metadata *keyType,
                                         /* +1 */ OpaqueValue *value,
                                         const Metadata *valueType) {
-  assert(task->hasTaskLocalValues());
   assert(value && "Task local value must not be nil");
 
   auto item = TaskLocalItem::createLink(task, keyType, valueType);
   valueType->vw_initializeWithTake(item->getStoragePtr(), value);
 
-  fprintf(stderr, "error: %s [%s:%d] bound item: task=%d item=%d -> item.next=%d keyType=%d item->getStoragePtr=%d\n", __FUNCTION__, __FILE_NAME__, __LINE__,
+  fprintf(stderr, "error: %s [%s:%d] PUSH bound item: task=%d item=%d -> item.next=%d keyType=%d item->getStoragePtr=%d\n", __FUNCTION__, __FILE_NAME__, __LINE__,
           task, item, item->getNext(), keyType, item->getStoragePtr());
 
   head = item;
 }
 
 void TaskLocalValuesFragment::popValue(AsyncTask *task) {
-  assert(task->hasTaskLocalValues());
   assert(head && "attempted to pop value off empty task-local stack");
   head->destroy();
   head = head->getNext();
 }
 
-OpaqueValue *TaskLocalValuesFragment::get(const Metadata *keyType) {
+OpaqueValue *TaskLocalValuesFragment::get(const Metadata *keyType, const TaskLocalInheritance inherit) {
   assert(keyType && "Task.Local key must not be null.");
 
   auto item = head;
@@ -98,6 +96,13 @@ OpaqueValue *TaskLocalValuesFragment::get(const Metadata *keyType) {
     if (item->keyType == keyType) {
       return item->getStoragePtr();
     }
+
+    // if the hey is an `inherit = .never` type, we stop our search the first
+    // time we would be jumping to a parent task to continue the search.
+    if (item->getNextLinkType() == NextLinkType::IsParent &&
+        inherit == TaskLocalInheritance::never)
+      return nullptr;
+
     item = item->getNext();
   }
 

stdlib/public/Concurrency/TaskLocal.swift

@@ -13,7 +13,15 @@
 import Swift
 @_implementationOnly import _SwiftConcurrencyShims
 
+#if canImport(Darwin)
 import Darwin
+#elseif canImport(Glibc)
+import Glibc
+#elseif os(Windows)
+import CRT
+#else
+#error("Unsupported platform")
+#endif
 
 /// Namespace for declaring `TaskLocalKey`s.
 public enum TaskLocalValues {}
@@ -34,6 +42,34 @@ public protocol TaskLocalKey {
   /// if the type itself does not have a good "undefined" or "zero" value that could
   /// be used here.
   static var defaultValue: Value { get }
+
+  /// Allows configuring specialized inheritance strategies for task local values.
+  ///
+  /// By default, task local values are accessible by the current or any of its
+  /// child tasks (with this rule applying recursively).
+  ///
+  /// Some, rare yet important, use-cases may require specialized inheritance
+  /// strategies, and this property allows them to configure these for their keys.
+  static var inherit: TaskLocalInheritance { get }
+}
+
+extension TaskLocalKey {
+  static var inherit: TaskLocalInheritance { .default }
+}
+
+///
+// TODO: should likely remain extensible
+public enum TaskLocalInheritance: Int {
+  /// The default inheritance strategy.
+  ///
+  /// Task local values whose keys are `default` inherited are available to the
+  /// task which declared them, as well as recursively by any child tasks
+  case `default` = 0
+
+  /// Causes task local values to never be inherited.
+  /// If the parent task has a value bound using this key, and a child task
+  /// attempts to look up a value of that key, it will return `defaultValue`.
+  case never = 1
 }
 
 extension Task {
@@ -76,7 +112,7 @@ extension Task {
     _taskLocalValuePush(task, keyType: Key.self, value: value)
   
     defer {
-      _taskLocalValuePop(task, count: 1)
+      _taskLocalValuePop(task)
     }
 
     return await body()
@@ -101,7 +137,7 @@ extension Task {
     _taskLocalValuePush(task, keyType: Key.self, value: value)
 
     defer {
-      _taskLocalValuePop(task, count: 1)
+      _taskLocalValuePop(task)
     }
 
     return try! await body()
@@ -133,26 +169,20 @@ extension AnyTaskLocalKey: Hashable {
 
 // ==== ------------------------------------------------------------------------
 
-@_silgen_name("swift_task_localValuePush")
+@_silgen_name("swift_task_local_value_push")
 public func _taskLocalValuePush<Value>(
   _ task: Builtin.NativeObject,
   keyType: Any.Type/*Key.Type*/,
   value: __owned Value
-)
+) // where Key: TaskLocalKey
 
-@_silgen_name("swift_task_localValuePop")
+@_silgen_name("swift_task_local_value_pop")
 public func _taskLocalValuePop(
-  _ task: Builtin.NativeObject,
+  _ task: Builtin.NativeObject
-  count: Int
 )
 
-@_silgen_name("swift_task_localValueGet")
+@_silgen_name("swift_task_local_value_get")
 public func _taskLocalValueGet(
   _ task: Builtin.NativeObject,
   keyType: Any.Type/*Key.Type*/
 ) -> UnsafeMutableRawPointer? // where Key: TaskLocalKey
-
-@_silgen_name("swift_task_hasTaskLocalValues")
-public func _taskHasTaskLocalValues(
-  _ task: Builtin.NativeObject
-) -> Bool

test/Concurrency/Runtime/progress_testing.swift

@@ -0,0 +1,127 @@
+// RUN: %target-typecheck-verify-swift -enable-experimental-concurrency
+// REQUIRES: concurrency
+
+import Swift
+
+import Darwin
+
+struct ProgressBox {
+  let callMe: (ProgressValue) -> ()
+
+  func claim(file: String, line: UInt) {
+    print("Progress box: claimed at \(file):\(line)")
+  }
+}
+
+struct ProgressReporter {
+  let callMe: () -> ()
+  let total: Int
+  init(parent: ProgressValue, callMe: () -> (), total: Int) {
+    self.callMe = callme
+    self.total = total
+  }
+
+  mutating func increment(by: Int = 1) {
+    callMe()
+  }
+}
+
+public struct ProgressValue {
+  public var total: Int
+  public var completed: Int
+  public var fractionCompleted: Double { get }
+
+  public enum Phase {
+    case active
+    case cancelled
+    case finished
+  }
+  public var phase: Phase
+}
+
+extension TaskLocalValues {
+  var progress: ProgressKey { .init() }
+  enum ProgressKey: TaskLocalKey {
+    static var defaultValue: ProgressValue? { nil }
+
+    static var inherit: TaskLocalInheritance { .never }
+  }
+}
+
+extension Task {
+
+  func withProgressObserver<T>(
+    _ onProgressUpdate: (ProgressValue) -> (),
+    operation: () async throws -> T
+  ) -> async rethrows T {
+    let box = ProgressBox(callMe: onProgressUpdate)
+    return try await Task.withLocal(\.progress, boundTo: box) {
+      try await operation
+    }
+  }
+
+  func withProgress(pending: Int) async {
+    if let parentBox = await Task.local(\.progress) {
+
+    }
+  }
+
+  func reportingProgress<T>(
+    pending: Int,
+    file: String = #file, line: UInt = #line,
+    body: (inout ProgressReporter) -> T) async -> T {
+    if let parentProgress = Task.local(\.progress) {
+      parentProgress.claim(file: file, line: line)
+
+      return await Task.withLocal(\.progress, boundTo: nil) {
+        // unbind the progress!
+        // As we're reporting things in this "leaf" task no other operation
+        // may report things.
+        var reporter = ProgressReporter(total: parentProgress)
+        return body(&reporter)
+      }
+    } else {
+      var reporter = ProgressReporter(total: 1) // or "noop"
+      return body(&reporter)
+    }
+  }
+}
+
+func main() async {
+  func makeDinner() async throws -> Meal {
+    var progress = await Task.reportingProgress(pending: 5)
+
+    async let veggies = progress.withPendingProgress(1) {
+      await chopVegetables()
+    }
+    async let meat = marinateMeat()
+
+    async let oven = progress.withPendingProgress(3) {
+      await preheatOven(temperature: 350)
+    }
+
+    let dish = Dish(ingredients: await [veggies, meat])
+    let dinner = await progress.withPendingProgress(1) {
+      await oven.cook(dish, duration: .hours(3))
+    }
+
+    return dinner
+  }
+
+  func chopVegetables() async -> [String] {
+    return []
+  }
+
+  func marinateMeat() async -> String {
+    ""
+  }
+
+  func preheatOven(temperature: Int) -> String {
+    ""
+  }
+
+  struct Dish {
+    init(ingredients: [String]) {}
+  }
+
+}

test/Concurrency/Runtime/task_locals_async_let.swift

@@ -30,34 +30,31 @@ extension TaskLocalValues {
   var number: NumberKey { .init() }
 }
 
+@discardableResult
 func printTaskLocal<Key>(
   _ key: KeyPath<TaskLocalValues, Key>,
   _ expected: Key.Value? = nil,
   file: String = #file, line: UInt = #line
-) async throws where Key: TaskLocalKey {
+) async throws -> Key.Value? where Key: TaskLocalKey {
   let value = await Task.local(key)
   print("\(Key.self): \(value) at \(file):\(line)")
   if let expected = expected {
     assert("\(expected)" == "\(value)",
       "Expected [\(expected)] but found: \(value), at \(file):\(line)")
   }
+  return expected
 }
 
 // ==== ------------------------------------------------------------------------
 
-
-func test2() async {
-  try! await printTaskLocal(\.number)
-}
-
 func async_let_nested() async {
-  _ = try! await printTaskLocal(\.number) // CHECK: NumberKey: 0 {{.*}}
+  _ = try! await printTaskLocal(\.number) // COM: NumberKey: 0 {{.*}}
   async let x1 = Task.withLocal(\.number, boundTo: 2) {
-    async let x2 = printTaskLocal(\.number) // CHECK: NumberKey: 2 {{.*}}
+    async let x2 = printTaskLocal(\.number) // COM: NumberKey: 2 {{.*}}
 
     func test() async {
-      try! await printTaskLocal(\.number) // CHECK: NumberKey: 2 {{.*}}
+      try! await printTaskLocal(\.number) // COM: NumberKey: 2 {{.*}}
-      async let x31 = test2() // CHECK: NumberKey: 2 {{.*}}
+      async let x31 = printTaskLocal(\.number) // COM: NumberKey: 2 {{.*}}
       try! await x31
     }
     async let x3 = test()
@@ -67,7 +64,29 @@ func async_let_nested() async {
   }
 
   _ = try! await x1
-  try! await printTaskLocal(\.number) // CHECK: NumberKey: 0 {{.*}}
+  try! await printTaskLocal(\.number) // COM: NumberKey: 0 {{.*}}
+}
+
+func async_let_nested_skip_optimization() async {
+  async let x1: Int? = Task.withLocal(\.number, boundTo: 2) {
+    async let x2: Int? = { () async -> Int? in
+      async let x3: Int? = { () async -> Int? in
+        async let x4: Int? = { () async -> Int? in
+          async let x5: Int? = { () async -> Int? in
+            async let xx = printTaskLocal(\.number) // CHECK: NumberKey: 2 {{.*}}
+            return try! await xx
+          }()
+          return try! await x5
+        }()
+        return try! await x4
+      }()
+      return try! await x3
+    }()
+    return try! await x2
+  }
+
+  _ = try! await x1
 }
 
 runAsyncAndBlock(async_let_nested)
+runAsyncAndBlock(async_let_nested_skip_optimization)