[TaskLocal] TL macro must handle force unwrapped optional types

Resolves rdar://128225191

Author: ktoso

lib/Macros/Sources/SwiftMacros/SyntaxExtensions.swift

@@ -30,4 +30,17 @@ extension VariableDeclSyntax {
       modifier.isAccessControl
     }
   }
+}
+
+extension ImplicitlyUnwrappedOptionalTypeSyntax {
+  internal var asOptionalTypeSyntax: some TypeSyntaxProtocol {
+    OptionalTypeSyntax(
+      leadingTrivia: leadingTrivia,
+      unexpectedBeforeWrappedType,
+      wrappedType: wrappedType,
+      self.unexpectedBetweenWrappedTypeAndExclamationMark,
+      self.unexpectedAfterExclamationMark,
+      trailingTrivia: self.trailingTrivia
+    )
+  }
 }

lib/Macros/Sources/SwiftMacros/TaskLocalMacro.swift

@@ -50,18 +50,24 @@ extension TaskLocalMacro: PeerMacro {
         message: "'@TaskLocal' property must have name", id: .incompatibleDecl)
     }
 
-    let type = firstBinding.typeAnnotation?.type
+    let originalType = firstBinding.typeAnnotation?.type
     let explicitTypeAnnotation: TypeAnnotationSyntax?
-    if let type {
-      explicitTypeAnnotation = TypeAnnotationSyntax(type: TypeSyntax("TaskLocal<\(type.trimmed)>"))
+    let taskLocalWrappedType: TypeSyntax?
+    if let forceUnwrappedOptionalType = firstBinding.typeAnnotation?.type.as(ImplicitlyUnwrappedOptionalTypeSyntax.self) {
+      taskLocalWrappedType = TypeSyntax("\(forceUnwrappedOptionalType.wrappedType.trimmed)?")
+    } else {
+      taskLocalWrappedType = originalType?.trimmed
+    }
+    if let taskLocalWrappedType {
+      explicitTypeAnnotation = TypeAnnotationSyntax(type: TypeSyntax("TaskLocal<\(taskLocalWrappedType)>"))
     } else {
       explicitTypeAnnotation = nil
     }
 
     let initialValue: ExprSyntax
     if let initializerValue = firstBinding.initializer?.value {
       initialValue = ExprSyntax(initializerValue)
-    } else if let type, type.isOptional {
+    } else if let originalType, originalType.isOptional {
       initialValue = ExprSyntax(NilLiteralExprSyntax())
     } else {
       throw DiagnosticsError(
@@ -162,7 +168,7 @@ extension TypeSyntax {
   // and T? we can detect the optional.
   fileprivate var isOptional: Bool {
     switch self.as(TypeSyntaxEnum.self) {
-    case .optionalType:
+    case .optionalType, .implicitlyUnwrappedOptionalType:
       return true
     case .identifierType(let identifierType):
       return identifierType.name.text == "Optional"
@@ -172,7 +178,8 @@ extension TypeSyntax {
         return false
       }
       return memberType.name.text == "Optional"
-    default: return false
+    default:
+      return false
     }
   }
 }

test/Concurrency/Runtime/async_task_locals_basic.swift

@@ -34,6 +34,9 @@ enum TL {
 
   @TaskLocal
   static var clazz: ClassTaskLocal?
+
+  @TaskLocal
+  static var force: ForceUnwrapMe!
 }
 
 @TaskLocal
@@ -43,6 +46,8 @@ var globalTaskLocal: StringLike = StringLike("<not-set>")
 @available(SwiftStdlib 5.10, *)
 struct LessAvailable {}
 
+struct ForceUnwrapMe {}
+
 @TaskLocal
 @available(SwiftStdlib 5.10, *)
 var globalLessAvailable: LessAvailable?