[Concurrency] Reimplement @TaskLocal as a macro

CHANGELOG.md

@@ -5,6 +5,35 @@
 
 ## Swift 6.0
 
+* Since its introduction in Swift 5.1 the @TaskLocal property wrapper was used to   
+  create and access task-local value bindings. Property wrappers introduce mutable storage,
+  which was now properly flagged as potential source of concurrency unsafety.
+
+  In order for Swift 6 language mode to not flag task-locals as potentially thread-unsafe,
+  task locals are now implemented using a macro. The macro has the same general semantics 
+  and usage patterns, however there are two source-break situations which the Swift 6 
+  task locals cannot handle:
+
+  Using an implicit default `nil` value for task local initialization, when combined with a type alias:
+  ```swift
+  // allowed in Swift 5.x, not allowed in Swift 6.x
+
+  typealias MyValue = Optional<Int> 
+
+  @TaskLocal
+  static var number: MyValue // Swift 6: error, please specify default value explicitly
+
+  // Solution 1: Specify the default value
+  @TaskLocal
+  static var number: MyValue = nil
+
+  // Solution 2: Avoid the type-alias
+  @TaskLocal
+  static var number: Optional<Int>
+  ```
+
+  At the same time, task locals can now be declared as global properties, which wasn't possible before.
+
 * Swift 5.10 missed a semantic check from [SE-0309][]. In type context, a reference to a
   protocol `P` that has associated types or `Self` requirements should use
   the `any` keyword, but this was not enforced in nested generic argument positions.

lib/Macros/Sources/SwiftMacros/CMakeLists.txt

@@ -14,6 +14,7 @@ add_swift_macro_library(SwiftMacros
   OptionSetMacro.swift
   DebugDescriptionMacro.swift
   DistributedResolvableMacro.swift
+  TaskLocalMacro.swift
   SWIFT_DEPENDENCIES
     SwiftDiagnostics
     SwiftOperators

lib/Macros/Sources/SwiftMacros/TaskLocalMacro.swift

@@ -0,0 +1,218 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2022-2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import SwiftSyntax
+import SwiftSyntaxMacros
+import SwiftDiagnostics
+
+/// Macro implementing the TaskLocal functionality.
+///
+/// It introduces a peer `static let $name: TaskLocal<Type>` as well as a getter
+/// that accesses the task local storage.
+public enum TaskLocalMacro {}
+
+extension TaskLocalMacro: PeerMacro {
+  public static func expansion(
+    of node: AttributeSyntax,
+    providingPeersOf declaration: some DeclSyntaxProtocol,
+    in context: some MacroExpansionContext
+  ) throws -> [DeclSyntax] {
+    guard let varDecl = try requireVar(declaration, diagnose: false) else {
+      return []
+    }
+    guard try requireStaticContext(varDecl, in: context, diagnose: false) else {
+      return []
+    }
+
+    guard varDecl.bindings.count == 1 else {
+      throw DiagnosticsError(
+        syntax: declaration,
+        message: "'@TaskLocal' property must have exactly one binding", id: .incompatibleDecl)
+    }
+    guard let firstBinding = varDecl.bindings.first else {
+      throw DiagnosticsError(
+        syntax: declaration,
+        message: "'@TaskLocal' property must have declared binding", id: .incompatibleDecl)
+    }
+
+    guard let name = firstBinding.pattern.as(IdentifierPatternSyntax.self)?.identifier else {
+      throw DiagnosticsError(
+        syntax: declaration,
+        message: "'@TaskLocal' property must have name", id: .incompatibleDecl)
+    }
+
+    let type = firstBinding.typeAnnotation?.type
+    let explicitTypeAnnotation: TypeAnnotationSyntax?
+    if let type {
+      explicitTypeAnnotation = TypeAnnotationSyntax(type: TypeSyntax("TaskLocal<\(type.trimmed)>"))
+    } else {
+      explicitTypeAnnotation = nil
+    }
+
+    let initialValue: ExprSyntax
+    if let initializerValue = firstBinding.initializer?.value {
+      initialValue = ExprSyntax(initializerValue)
+    } else if let type, type.isOptional {
+      initialValue = ExprSyntax(NilLiteralExprSyntax())
+    } else {
+      throw DiagnosticsError(
+        syntax: declaration,
+        message: "'@TaskLocal' property must have default value, or be optional", id: .mustBeVar)
+    }
+
+    // If the property is global, do not prefix the synthesised decl with 'static'
+    let isGlobal = context.lexicalContext.isEmpty
+    let staticKeyword: TokenSyntax?
+    if isGlobal {
+      staticKeyword = nil
+    } else {
+      staticKeyword = TokenSyntax.keyword(.static, trailingTrivia: .space)
+    }
+
+    return [
+      """
+      \(staticKeyword)let $\(name)\(explicitTypeAnnotation) = TaskLocal(wrappedValue: \(initialValue))
+      """
+    ]
+  }
+}
+
+extension TaskLocalMacro: AccessorMacro {
+  public static func expansion(
+    of node: AttributeSyntax,
+    providingAccessorsOf declaration: some DeclSyntaxProtocol,
+    in context: some MacroExpansionContext
+  ) throws -> [AccessorDeclSyntax] {
+    // We very specifically have to fail and diagnose in the accessor macro,
+    // rather than in the peer macro, since returning [] from the accessor
+    // macro adds another series of errors about it missing to emit a decl.
+    guard let varDecl = try requireVar(declaration) else {
+      return []
+    }
+    try requireStaticContext(varDecl, in: context)
+
+    guard let firstBinding = varDecl.bindings.first else {
+      return []
+    }
+
+    guard let name = firstBinding.pattern.as(IdentifierPatternSyntax.self)?.identifier else {
+      return []
+    }
+
+    return ["get { $\(name).get() }"]
+  }
+}
+
+@discardableResult
+private func requireVar(_ decl: some DeclSyntaxProtocol,
+                        diagnose: Bool = true) throws -> VariableDeclSyntax? {
+  if let varDecl = decl.as(VariableDeclSyntax.self) {
+    return varDecl
+  }
+  if diagnose {
+    throw DiagnosticsError(
+      syntax: decl,
+      message: "'@TaskLocal' can only be applied to properties", id: .mustBeVar)
+  }
+
+  return nil
+}
+
+@discardableResult
+private func requireStaticContext(_ decl: VariableDeclSyntax,
+                                  in context: some MacroExpansionContext,
+                                  diagnose: Bool = true) throws -> Bool {
+  let isStatic = decl.modifiers.contains { modifier in
+    modifier.name.text == "\(Keyword.static)"
+  }
+
+  if isStatic {
+    return true
+  }
+
+  let isGlobal = context.lexicalContext.isEmpty
+  if isGlobal {
+    return true
+  }
+
+  if diagnose {
+    throw DiagnosticsError(
+      syntax: decl,
+      message: "'@TaskLocal' can only be applied to 'static' property, or global variables", id: .mustBeStatic)
+  }
+
+  return false
+}
+
+extension TypeSyntax {
+  // This isn't great since we can't handle type aliases since the macro
+  // has no type information, but at least for the common case for Optional<T>
+  // and T? we can detect the optional.
+  fileprivate var isOptional: Bool {
+    switch self.as(TypeSyntaxEnum.self) {
+    case .optionalType:
+      return true
+    case .identifierType(let identifierType):
+      return identifierType.name.text == "Optional"
+    case .memberType(let memberType):
+      guard let baseIdentifier = memberType.baseType.as(IdentifierTypeSyntax.self),
+            baseIdentifier.name.text == "Swift" else {
+        return false
+      }
+      return memberType.name.text == "Optional"
+    default: return false
+    }
+  }
+}
+
+struct TaskLocalMacroDiagnostic: DiagnosticMessage {
+  enum ID: String {
+    case mustBeVar = "must be var"
+    case mustBeStatic = "must be static"
+    case incompatibleDecl = "incompatible declaration"
+  }
+
+  var message: String
+  var diagnosticID: MessageID
+  var severity: DiagnosticSeverity
+
+  init(message: String, diagnosticID: SwiftDiagnostics.MessageID, severity: SwiftDiagnostics.DiagnosticSeverity = .error) {
+    self.message = message
+    self.diagnosticID = diagnosticID
+    self.severity = severity
+  }
+
+  init(message: String, domain: String, id: ID, severity: SwiftDiagnostics.DiagnosticSeverity = .error) {
+    self.message = message
+    self.diagnosticID = MessageID(domain: domain, id: id.rawValue)
+    self.severity = severity
+  }
+}
+
+extension DiagnosticsError {
+  init(
+    syntax: some SyntaxProtocol,
+    message: String,
+    domain: String = "Swift",
+    id: TaskLocalMacroDiagnostic.ID,
+    severity: SwiftDiagnostics.DiagnosticSeverity = .error) {
+    self.init(diagnostics: [
+      Diagnostic(
+        node: Syntax(syntax),
+        message: TaskLocalMacroDiagnostic(
+          message: message,
+          domain: domain,
+          id: id,
+          severity: severity))
+    ])
+  }
+}

lib/Sema/TypeCheckConcurrency.cpp

@@ -4980,15 +4980,6 @@ ActorIsolation ActorIsolationRequest::evaluate(
       if (var->isGlobalStorage() && !isActorType) {
         auto *diagVar = var;
         if (auto *originalVar = var->getOriginalWrappedProperty()) {
-          // temporary 5.10 checking bypass for @TaskLocal <rdar://120907014>
-          // TODO: @TaskLocal should be a macro <rdar://120914014>
-          if (auto *classDecl =
-                  var->getInterfaceType()->getClassOrBoundGenericClass()) {
-            auto &ctx = var->getASTContext();
-            if (classDecl == ctx.getTaskLocalDecl()) {
-              return isolation;
-            }
-          }
           diagVar = originalVar;
         }
         if (var->isLet()) {

stdlib/public/Concurrency/TaskLocal.swift

@@ -13,22 +13,42 @@
 import Swift
 @_implementationOnly import _SwiftConcurrencyShims
 
-/// Property wrapper that defines a task-local value key.
+
+// Macros are disabled when Swift is built without swift-syntax.
+#if $Macros && hasAttribute(attached)
+
+/// Macro that introduces a ``TaskLocal-class`` binding.
+///
+/// For information about task-local bindings, see ``TaskLocal-class``.
+///
+/// - SeeAlso: ``TaskLocal-class``
+@available(SwiftStdlib 5.1, *)
+@attached(accessor)
+@attached(peer, names: prefixed(`$`))
+public macro TaskLocal() =
+  #externalMacro(module: "SwiftMacros", type: "TaskLocalMacro")
+
+#endif
+
+/// Wrapper type that defines a task-local value key.
 ///
 /// A task-local value is a value that can be bound and read in the context of a
-/// `Task`. It is implicitly carried with the task, and is accessible by any
-/// child tasks the task creates (such as TaskGroup or `async let` created tasks).
+/// ``Task``. It is implicitly carried with the task, and is accessible by any
+/// child tasks it creates (such as TaskGroup or `async let` created tasks).
 ///
 /// ### Task-local declarations
 ///
-/// Task locals must be declared as static properties (or global properties,
-/// once property wrappers support these), like this:
+/// Task locals must be declared as static properties or global properties, like this:
 ///
 ///     enum Example {
 ///         @TaskLocal
 ///         static var traceID: TraceID?
 ///     }
 ///
+///     // Global task local properties are supported since Swift 6.0:
+///     @TaskLocal
+///     var contextualNumber: Int = 12
+///
 /// ### Default values
 /// Reading a task local value when no value was bound to it results in returning
 /// its default value. For a task local declared as optional (such as e.g. `TraceID?`),
@@ -137,7 +157,8 @@ import Swift
 ///         read() // traceID: nil
 ///       }
 ///     }
-@propertyWrapper
+///
+/// - SeeAlso: ``TaskLocal-macro``
 @available(SwiftStdlib 5.1, *)
 public final class TaskLocal<Value: Sendable>: Sendable, CustomStringConvertible {
   let defaultValue: Value

test/Concurrency/Runtime/async_task_locals_spawn_let.swift → test/Concurrency/Runtime/async_task_locals_async_let.swift

@@ -1,4 +1,4 @@
-// RUN: %target-run-simple-swift( -Xfrontend -disable-availability-checking -parse-as-library) | %FileCheck %s
+// RUN: %target-run-simple-swift( -plugin-path %swift-plugin-dir -Xfrontend -disable-availability-checking -parse-as-library) | %FileCheck %s
 
 // REQUIRES: executable_test
 // REQUIRES: concurrency

test/Concurrency/Runtime/async_task_locals_basic.swift

@@ -1,4 +1,4 @@
-// RUN: %target-run-simple-swift( -Xfrontend -disable-availability-checking -parse-as-library %import-libdispatch) | %FileCheck %s
+// RUN: %target-run-simple-swift( -plugin-path %swift-plugin-dir -Xfrontend -disable-availability-checking -parse-as-library %import-libdispatch) | %FileCheck %s
 
 // REQUIRES: executable_test
 // REQUIRES: concurrency
@@ -8,11 +8,14 @@
 // REQUIRES: concurrency_runtime
 // UNSUPPORTED: back_deployment_runtime
 
-final class StringLike: Sendable, CustomStringConvertible {
+final class StringLike: Sendable, ExpressibleByStringLiteral, CustomStringConvertible {
   let value: String
   init(_ value: String) {
     self.value = value
   }
+  init(stringLiteral value: StringLiteralType) {
+    self.value = value
+  }
 
   var description: String { value }
 }
@@ -33,6 +36,9 @@ enum TL {
   static var clazz: ClassTaskLocal?
 }
 
+@TaskLocal
+var globalTaskLocal: StringLike = StringLike("<not-set>")
+
 @available(SwiftStdlib 5.1, *)
 final class ClassTaskLocal: Sendable {
   init() {
@@ -217,6 +223,13 @@ func inside_actor() async {
   await Worker().setAndRead()
 }
 
+@available(SwiftStdlib 5.1, *)
+func global_task_local() async {
+  await $globalTaskLocal.withValue("value-1") {
+    await printTaskLocalAsync($globalTaskLocal) // CHECK-NEXT: TaskLocal<StringLike>(defaultValue: <not-set>) (value-1)
+  }
+}
+
 @available(SwiftStdlib 5.1, *)
 @main struct Main {
   static func main() async {
@@ -229,5 +242,6 @@ func inside_actor() async {
     await nested_3_onlyTopContributesAsync()
     await nested_3_onlyTopContributesMixed()
     await inside_actor()
+    await global_task_local()
   }
 }