[Concurrency] TaskLocal.withValue adopting #isolation

Author: ktoso

stdlib/public/Concurrency/CheckedContinuation.swift

@@ -284,9 +284,7 @@ extension CheckedContinuation {
 @inlinable
 @_unavailableInEmbedded
 @available(SwiftStdlib 5.1, *)
-#if !$Embedded
 @backDeployed(before: SwiftStdlib 6.0)
-#endif
 public func withCheckedContinuation<T>(
   isolation: isolated (any Actor)? = #isolation,
   function: String = #function,
@@ -343,9 +341,7 @@ internal func __abi_withCheckedContinuation<T>(
 @inlinable
 @_unavailableInEmbedded
 @available(SwiftStdlib 5.1, *)
-#if !$Embedded
 @backDeployed(before: SwiftStdlib 6.0)
-#endif
 public func withCheckedThrowingContinuation<T>(
   isolation: isolated (any Actor)? = #isolation,
   function: String = #function,

stdlib/public/Concurrency/TaskLocal.swift

@@ -181,10 +181,27 @@ public final class TaskLocal<Value: Sendable>: Sendable, CustomStringConvertible
   /// the operation closure.
   @inlinable
   @discardableResult
-  @_unsafeInheritExecutor
-  @backDeployed(before: SwiftStdlib 5.8)
-  public func withValue<R>(_ valueDuringOperation: Value, operation: () async throws -> R,
+  @available(SwiftStdlib 5.1, *)
+  @backDeployed(before: SwiftStdlib 6.0)
+  public func withValue<R>(_ valueDuringOperation: Value,
+                           operation: () async throws -> R,
+                           isolation: isolated (any Actor)? = #isolation,
                            file: String = #fileID, line: UInt = #line) async rethrows -> R {
+    return try await withValueImpl(
+      valueDuringOperation,
+      operation: operation,
+      isolation: isolation,
+      file: file, line: line)
+  }
+
+  @usableFromInline
+  @discardableResult
+  @_unsafeInheritExecutor // ABI compatibility with Swift 5.1
+  @available(SwiftStdlib 5.1, *)
+  @_silgen_name("$ss9TaskLocalC9withValue_9operation4file4lineqd__x_qd__yYaKXESSSutYaKlF")
+  internal func __abi_withValue<R>(_ valueDuringOperation: Value,
+                                   operation: () async throws -> R,
+                                   file: String = #fileID, line: UInt = #line) async rethrows -> R {
     return try await withValueImpl(valueDuringOperation, operation: operation, file: file, line: line)
   }
 
@@ -206,11 +223,30 @@ public final class TaskLocal<Value: Sendable>: Sendable, CustomStringConvertible
   /// to swift_task_de/alloc for the copy as follows:
   /// - withValue contains the compiler-emitted calls swift_task_de/alloc.
   /// - withValueImpl contains the calls to _taskLocalValuePush/Pop
+  @inlinable
+  @discardableResult
+  @available(SwiftStdlib 5.1, *)
+  @backDeployed(before: SwiftStdlib 6.0)
+  internal func withValueImpl<R>(_ valueDuringOperation: __owned Value,
+                                 operation: () async throws -> R,
+                                 isolation: isolated (any Actor)?,
+                                 file: String = #fileID, line: UInt = #line) async rethrows -> R {
+    // check if we're not trying to bind a value from an illegal context; this may crash
+    _checkIllegalTaskLocalBindingWithinWithTaskGroup(file: file, line: line)
+
+    _taskLocalValuePush(key: key, value: consume valueDuringOperation)
+    defer { _taskLocalValuePop() }
+
+    return try await operation()
+  }
+
   @inlinable
   @discardableResult
   @_unsafeInheritExecutor
+  @available(SwiftStdlib 5.1, *)
   @backDeployed(before: SwiftStdlib 5.9)
-  internal func withValueImpl<R>(_ valueDuringOperation: __owned Value, operation: () async throws -> R,
+  internal func withValueImpl<R>(_ valueDuringOperation: __owned Value,
+                                 operation: () async throws -> R,
                                  file: String = #fileID, line: UInt = #line) async rethrows -> R {
     // check if we're not trying to bind a value from an illegal context; this may crash
     _checkIllegalTaskLocalBindingWithinWithTaskGroup(file: file, line: line)
@@ -221,6 +257,7 @@ public final class TaskLocal<Value: Sendable>: Sendable, CustomStringConvertible
     return try await operation()
   }
 
+
   /// Binds the task-local to the specific value for the duration of the
   /// synchronous operation.
   ///

test/Concurrency/async_task_locals_basic_warnings.swift

@@ -7,19 +7,36 @@
 // REQUIRES: concurrency
 // REQUIRES: asserts
 
-@available(SwiftStdlib 5.1, *)
 actor Test {
 
   @TaskLocal static var local: Int?
 
   func run() async {
     // This should NOT produce any warnings, the closure withValue uses is @Sendable:
-    await Test.$local.withValue(42) {
+    await Self.$local.withValue(42) {
       await work()
     }
   }
 
   func work() async {
-    print("Hello \(Test.local ?? 0)")
+    print("Hello \(Self.local ?? 0)")
+  }
+}
+
+class NonSendableValue {}
+
+@MainActor final class MainHelper {
+
+  @TaskLocal static var local: Int?
+
+  func run() async {
+    // This should NOT produce any warnings, the closure withValue uses is @Sendable:
+    await Self.$local.withValue(42) {
+      await work()
+    }
+  }
+
+  func work() async {
+    print("Hello \(Self.local ?? 0)")
   }
 }

test/Concurrency/async_task_locals_basic_warnings_bug_isolation.swift

@@ -0,0 +1,28 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend -emit-module -emit-module-path %t/OtherActors.swiftmodule -module-name OtherActors %S/Inputs/OtherActors.swift -disable-availability-checking
+
+// RUN: %target-swift-frontend -I %t  -disable-availability-checking -strict-concurrency=complete -parse-as-library %s -emit-sil -o /dev/null -verify
+// RUN: %target-swift-frontend -I %t  -disable-availability-checking -strict-concurrency=complete -parse-as-library %s -emit-sil -o /dev/null -verify -enable-upcoming-feature RegionBasedIsolation
+
+// REQUIRES: concurrency
+// REQUIRES: asserts
+
+// FIXME: rdar://125078448 is resolved
+// XFAIL: *
+
+actor Test {
+
+  @TaskLocal static var local: Int?
+
+  func withTaskLocal(isolation: isolated (any Actor)? = #isolation,
+                     _ body: (consuming NonSendableValue, isolated (any Actor)?) -> Void) async {
+    Self.$local.withValue(12) {
+      // Unexpected errors here:
+      //  error: unexpected warning produced: transferring 'body' may cause a race; this is an error in the Swift 6 language mode
+      //  error: unexpected note produced: actor-isolated 'body' is captured by a actor-isolated closure. actor-isolated uses in closure may race against later nonisolated uses
+      body(NonSendableValue(), isolation)
+    }
+  }
+}
+
+class NonSendableValue {}

test/abi/macOS/arm64/concurrency.swift

@@ -298,3 +298,11 @@ Added: _$sScfsE13checkIsolatedyyF
 Added: _$sScf13checkIsolatedyyFTj
 // method descriptor for Swift.SerialExecutor.checkIsolated() -> ()
 Added: _$sScf13checkIsolatedyyFTq
+
+// #isolated adoption in TaskLocal.withValue
+// Swift.TaskLocal.withValueImpl<A>(_: __owned A, operation: () async throws -> A1, isolation: isolated Swift.Actor?, file: Swift.String, line: Swift.UInt) async throws -> A1
+Added: _$ss9TaskLocalC13withValueImpl_9operation9isolation4file4lineqd__xn_qd__yYaKXEScA_pSgYiSSSutYaKlF
+Added: _$ss9TaskLocalC13withValueImpl_9operation9isolation4file4lineqd__xn_qd__yYaKXEScA_pSgYiSSSutYaKlFTu
+// Swift.TaskLocal.withValue<A>(_: A, operation: () async throws -> A1, isolation: isolated Swift.Actor?, file: Swift.String, line: Swift.UInt) async throws -> A1
+Added: _$ss9TaskLocalC9withValue_9operation9isolation4file4lineqd__x_qd__yYaKXEScA_pSgYiSSSutYaKlF
+Added: _$ss9TaskLocalC9withValue_9operation9isolation4file4lineqd__x_qd__yYaKXEScA_pSgYiSSSutYaKlFTu

test/api-digester/stability-concurrency-abi.test

@@ -100,8 +100,14 @@ Func SerialExecutor.enqueue(_:) has been added as a protocol requirement
 // rdar://106833284 (ABI checker confused with overload getting deprecated)
 Func Executor.enqueue(_:) is a new API without @available attribute
 
-// This function correctly inherits its availability from the TaskLocal struct.
-Func TaskLocal.withValueImpl(_:operation:file:line:) is a new API without @available attribute
+// Adopt #isolation in TaskLocal.withValue APIs
+Func TaskLocal.withValue(_:operation:file:line:) has been renamed to Func withValue(_:operation:isolation:file:line:)
+Func TaskLocal.withValue(_:operation:file:line:) has mangled name changing from '_Concurrency.TaskLocal.withValue<A>(_: A, operation: () async throws -> A1, file: Swift.String, line: Swift.UInt) async throws -> A1' to '_Concurrency.TaskLocal.withValue<A>(_: A, operation: () throws -> A1, file: Swift.String, line: Swift.UInt) throws -> A1'
+Func TaskLocal.withValue(_:operation:file:line:) has mangled name changing from '_Concurrency.TaskLocal.withValue<A>(_: A, operation: () throws -> A1, file: Swift.String, line: Swift.UInt) throws -> A1' to '_Concurrency.TaskLocal.withValue<A>(_: A, operation: () async throws -> A1, isolation: isolated Swift.Optional<Swift.Actor>, file: Swift.String, line: Swift.UInt) async throws -> A1'
+Func TaskLocal.withValue(_:operation:file:line:) has parameter 1 type change from () async throws -> τ_1_0 to () throws -> τ_1_0
+Func TaskLocal.withValue(_:operation:file:line:) has parameter 1 type change from () throws -> τ_1_0 to () async throws -> τ_1_0
+Func TaskLocal.withValue(_:operation:file:line:) has parameter 2 type change from Swift.String to (any _Concurrency.Actor)?
+Func TaskLocal.withValue(_:operation:file:line:) has parameter 3 type change from Swift.UInt to Swift.String
 
 // *** DO NOT DISABLE OR XFAIL THIS TEST. *** (See comment above.)