codex: address warnings and deprecations

Author: Thibault Wittemberg

Sources/AsyncChannels/AsyncBufferedChannel.swift

@@ -157,27 +157,7 @@ public final class AsyncBufferedChannel<Element: Sendable>: AsyncSequence, Senda
     let awaitingId = self.generateId()
     let cancellation = ManagedCriticalState<Bool>(false)
 
-    return await withTaskCancellationHandler { [state] in
-      let awaiting = state.withCriticalRegion { state -> Awaiting? in
-        cancellation.withCriticalRegion { cancellation in
-          cancellation = true
-        }
-        switch state {
-          case .awaiting(var awaitings):
-            let awaiting = awaitings.remove(.placeHolder(id: awaitingId))
-            if awaitings.isEmpty {
-              state = .idle
-            } else {
-              state = .awaiting(awaitings)
-            }
-            return awaiting
-          default:
-            return nil
-        }
-      }
-
-      awaiting?.continuation?.resume(returning: nil)
-    } operation: {
+    return await withTaskCancellationHandler {
       await withUnsafeContinuation { [state] (continuation: UnsafeContinuation<Element?, Never>) in
         let decision = state.withCriticalRegion { state -> AwaitingDecision in
           let isCancelled = cancellation.withCriticalRegion { $0 }
@@ -218,6 +198,26 @@ public final class AsyncBufferedChannel<Element: Sendable>: AsyncSequence, Senda
             onSuspend?()
         }
       }
+    } onCancel: { [state] in
+      let awaiting = state.withCriticalRegion { state -> Awaiting? in
+        cancellation.withCriticalRegion { cancellation in
+          cancellation = true
+        }
+        switch state {
+          case .awaiting(var awaitings):
+            let awaiting = awaitings.remove(.placeHolder(id: awaitingId))
+            if awaitings.isEmpty {
+              state = .idle
+            } else {
+              state = .awaiting(awaitings)
+            }
+            return awaiting
+          default:
+            return nil
+        }
+      }
+
+      awaiting?.continuation?.resume(returning: nil)
     }
   }
 

Sources/AsyncChannels/AsyncThrowingBufferedChannel.swift

@@ -178,27 +178,7 @@ public final class AsyncThrowingBufferedChannel<Element, Failure: Error>: AsyncS
     let awaitingId = self.generateId()
     let cancellation = ManagedCriticalState<Bool>(false)
 
-    return try await withTaskCancellationHandler { [state] in
-      let awaiting = state.withCriticalRegion { state -> Awaiting? in
-        cancellation.withCriticalRegion { cancellation in
-          cancellation = true
-        }
-        switch state {
-          case .awaiting(var awaitings):
-            let awaiting = awaitings.remove(.placeHolder(id: awaitingId))
-            if awaitings.isEmpty {
-              state = .idle
-            } else {
-              state = .awaiting(awaitings)
-            }
-            return awaiting
-          default:
-            return nil
-        }
-      }
-
-      awaiting?.continuation?.resume(returning: nil)
-    } operation: {
+    return try await withTaskCancellationHandler {
       try await withUnsafeThrowingContinuation { [state] (continuation: UnsafeContinuation<Element?, Error>) in
         let decision = state.withCriticalRegion { state -> AwaitingDecision in
           let isCancelled = cancellation.withCriticalRegion { $0 }
@@ -245,6 +225,26 @@ public final class AsyncThrowingBufferedChannel<Element, Failure: Error>: AsyncS
             onSuspend?()
         }
       }
+    } onCancel: { [state] in
+      let awaiting = state.withCriticalRegion { state -> Awaiting? in
+        cancellation.withCriticalRegion { cancellation in
+          cancellation = true
+        }
+        switch state {
+          case .awaiting(var awaitings):
+            let awaiting = awaitings.remove(.placeHolder(id: awaitingId))
+            if awaitings.isEmpty {
+              state = .idle
+            } else {
+              state = .awaiting(awaitings)
+            }
+            return awaiting
+          default:
+            return nil
+        }
+      }
+
+      awaiting?.continuation?.resume(returning: nil)
     }
   }
 

Sources/Combiners/Merge/MergeStateMachine.swift

@@ -197,8 +197,6 @@ struct MergeStateMachine<Element>: Sendable {
 
   func next() async -> RegulatedElement<Element> {
     await withTaskCancellationHandler {
-      self.unsuspendAndClearOnCancel()
-    } operation: {
       self.requestNextRegulatedElements()
 
       let regulatedElement = await withUnsafeContinuation { (continuation: UnsafeContinuation<RegulatedElement<Element>, Never>) in
@@ -245,6 +243,8 @@ struct MergeStateMachine<Element>: Sendable {
       }
 
       return regulatedElement
+    } onCancel: {
+      self.unsuspendAndClearOnCancel()
     }
   }
 }

Sources/Combiners/WithLatestFrom/AsyncWithLatestFrom2Sequence.swift

@@ -172,12 +172,7 @@ where Other1: Sendable, Other2: Sendable, Other1.Element: Sendable, Other2.Eleme
       let shouldReturnNil = self.isTerminated.withCriticalRegion { $0 }
       guard !shouldReturnNil else { return nil }
 
-      return try await withTaskCancellationHandler { [isTerminated, othersTask] in
-        isTerminated.withCriticalRegion { isTerminated in
-          isTerminated = true
-        }
-        othersTask?.cancel()
-      } operation: { [othersTask, othersState, onBaseElement] in
+      return try await withTaskCancellationHandler { [othersTask, othersState, onBaseElement] in
         do {
           while true {
             guard let baseElement = try await self.base.next() else {
@@ -219,6 +214,11 @@ where Other1: Sendable, Other2: Sendable, Other1.Element: Sendable, Other2.Eleme
           othersTask?.cancel()
           throw error
         }
+      } onCancel: { [isTerminated, othersTask] in
+        isTerminated.withCriticalRegion { isTerminated in
+          isTerminated = true
+        }
+        othersTask?.cancel()
       }
     }
   }

Sources/Combiners/WithLatestFrom/AsyncWithLatestFromSequence.swift

@@ -121,9 +121,7 @@ where Other: Sendable, Other.Element: Sendable {
     public mutating func next() async rethrows -> Element? {
       guard !self.isTerminated else { return nil }
 
-      return try await withTaskCancellationHandler { [otherTask] in
-        otherTask?.cancel()
-      } operation: { [otherTask, otherState, onBaseElement] in
+      return try await withTaskCancellationHandler { [otherTask, otherState, onBaseElement] in
         do {
           while true {
             guard let baseElement = try await self.base.next() else {
@@ -157,6 +155,8 @@ where Other: Sendable, Other.Element: Sendable {
           otherTask?.cancel()
           throw error
         }
+      } onCancel: { [otherTask] in
+        otherTask?.cancel()
       }
     }
   }

Sources/Combiners/Zip/Zip2Runtime.swift

@@ -148,12 +148,6 @@ where Base1: Sendable, Base1.Element: Sendable, Base2: Sendable, Base2.Element:
 
   func next() async rethrows -> (Base1.Element, Base2.Element)? {
     try await withTaskCancellationHandler {
-      let output = self.stateMachine.withCriticalRegion { stateMachine in
-        stateMachine.rootTaskIsCancelled()
-      }
-
-      self.handle(rootTaskIsCancelledOutput: output)
-    } operation: {
       let results = await withUnsafeContinuation { (continuation: UnsafeContinuation<(Result<Base1.Element, Error>, Result<Base2.Element, Error>)?, Never>) in
         let output = self.stateMachine.withCriticalRegion { stateMachine in
           stateMachine.newDemandFromConsumer(suspendedDemand: continuation)
@@ -173,6 +167,12 @@ where Base1: Sendable, Base1.Element: Sendable, Base2: Sendable, Base2.Element:
       self.handle(demandIsFulfilledOutput: output)
 
       return try (results.0._rethrowGet(), results.1._rethrowGet())
+    } onCancel: {
+      let output = self.stateMachine.withCriticalRegion { stateMachine in
+        stateMachine.rootTaskIsCancelled()
+      }
+
+      self.handle(rootTaskIsCancelledOutput: output)
     }
   }
 

Sources/Combiners/Zip/Zip3Runtime.swift

@@ -186,12 +186,6 @@ where Base1: Sendable, Base1.Element: Sendable, Base2: Sendable, Base2.Element:
 
   func next() async rethrows -> (Base1.Element, Base2.Element, Base3.Element)? {
     try await withTaskCancellationHandler {
-      let output = self.stateMachine.withCriticalRegion { stateMachine in
-        stateMachine.rootTaskIsCancelled()
-      }
-
-      self.handle(rootTaskIsCancelledOutput: output)
-    } operation: {
       let results = await withUnsafeContinuation { (continuation: UnsafeContinuation<(Result<Base1.Element, Error>, Result<Base2.Element, Error>, Result<Base3.Element, Error>)?, Never>) in
         let output = self.stateMachine.withCriticalRegion { stateMachine in
           stateMachine.newDemandFromConsumer(suspendedDemand: continuation)
@@ -211,6 +205,12 @@ where Base1: Sendable, Base1.Element: Sendable, Base2: Sendable, Base2.Element:
       self.handle(demandIsFulfilledOutput: output)
 
       return try (results.0._rethrowGet(), results.1._rethrowGet(), results.2._rethrowGet())
+    } onCancel: {
+      let output = self.stateMachine.withCriticalRegion { stateMachine in
+        stateMachine.rootTaskIsCancelled()
+      }
+
+      self.handle(rootTaskIsCancelledOutput: output)
     }
   }
 

Sources/Combiners/Zip/ZipRuntime.swift

@@ -120,12 +120,6 @@ where Base: Sendable, Base.Element: Sendable {
 
   func next() async rethrows -> [Base.Element]? {
     try await withTaskCancellationHandler {
-      let output = self.stateMachine.withCriticalRegion { stateMachine in
-        stateMachine.rootTaskIsCancelled()
-      }
-
-      self.handle(rootTaskIsCancelledOutput: output)
-    } operation: {
       let results = await withUnsafeContinuation { (continuation: UnsafeContinuation<[Int: Result<Base.Element, Error>]?, Never>) in
         let output = self.stateMachine.withCriticalRegion { stateMachine in
           stateMachine.newDemandFromConsumer(suspendedDemand: continuation)
@@ -145,6 +139,12 @@ where Base: Sendable, Base.Element: Sendable {
       self.handle(demandIsFulfilledOutput: output)
 
       return try results.sorted { $0.key < $1.key }.map { try $0.value._rethrowGet() }
+    } onCancel: {
+      let output = self.stateMachine.withCriticalRegion { stateMachine in
+        stateMachine.rootTaskIsCancelled()
+      }
+
+      self.handle(rootTaskIsCancelledOutput: output)
     }
   }
 

Sources/Creators/AsyncTimerSequence.swift

@@ -78,11 +78,11 @@ public struct AsyncTimerSequence: AsyncSequence {
     }
 
     public mutating func next() async -> Element? {
-      await withTaskCancellationHandler { [task] in
-        task.cancel()
-      } operation: {
+      await withTaskCancellationHandler {
         guard !Task.isCancelled else { return nil }
         return await self.iterator.next()
+      } onCancel: { [task] in
+        task.cancel()
       }
     }
   }

Sources/Operators/AsyncMulticastSequence.swift

@@ -101,47 +101,42 @@ where Base.Element == Subject.Element, Subject.Failure == Error, Base.AsyncItera
 
   /// Allow the `AsyncIterator` to produce elements.
   public func connect() {
+    self.isConnected.apply(criticalState: true)
     self.connectedGate.send(())
   }
 
   func next() async {
-    guard !Task.isCancelled else { return }
-
-    let (canAccessBase, iterator) = self.state.withCriticalRegion { state -> (Bool, Base.AsyncIterator?) in
-      switch state {
-        case .available(let iterator):
-          state = .busy
-          return (true, iterator)
-        case .busy:
-          return (false, nil)
+    await Task {
+      let (canAccessBase, iterator) = self.state.withCriticalRegion { state -> (Bool, Base.AsyncIterator?) in
+        switch state {
+          case .available(let iterator):
+            state = .busy
+            return (true, iterator)
+          case .busy:
+            return (false, nil)
+        }
+      }
+
+      guard canAccessBase, var iterator = iterator else { return }
+
+      let toSend: Result<Element?, Error>
+      do {
+        let element = try await iterator.next()
+        toSend = .success(element)
+      } catch {
+        toSend = .failure(error)
       }
-    }
 
-    guard canAccessBase, var iterator = iterator else { return }
-    defer {
       self.state.withCriticalRegion { state in
         state = .available(iterator)
       }
-    }
-
-    guard !Task.isCancelled else { return }
 
-    let toSend: Result<Element?, Error>
-    do {
-      let element = try await iterator.next()
-      toSend = .success(element)
-    } catch {
-      guard !Task.isCancelled else { return }
-      toSend = .failure(error)
-    }
-
-    guard !Task.isCancelled else { return }
-
-    switch toSend {
-      case .success(.some(let element)): self.subject.send(element)
-      case .success(.none): self.subject.send(.finished)
-      case .failure(let error): self.subject.send(.failure(error))
-    }
+      switch toSend {
+        case .success(.some(let element)): self.subject.send(element)
+        case .success(.none): self.subject.send(.finished)
+        case .failure(let error): self.subject.send(.failure(error))
+      }
+    }.value
   }
 
   public func makeAsyncIterator() -> AsyncIterator {
@@ -163,14 +158,8 @@ where Base.Element == Subject.Element, Subject.Failure == Error, Base.AsyncItera
     public mutating func next() async rethrows -> Element? {
       guard !Task.isCancelled else { return nil }
 
-      let shouldWaitForGate = self.isConnected.withCriticalRegion { isConnected -> Bool in
-        if !isConnected {
-          isConnected = true
-          return true
-        }
-        return false
-      }
-      if shouldWaitForGate {
+      let isConnected = self.isConnected.withCriticalRegion { $0 }
+      if !isConnected {
         await self.connectedGateIterator.next()
       }