Skip to content

Commit 68efe53

Browse files
swhittyswhitty
committed
Refactor
1 parent 281ed7a commit 68efe53

File tree

1 file changed

+83
-76
lines changed

1 file changed

+83
-76
lines changed

‎Sources/IdentifiableContinuation.swift‎

Lines changed: 83 additions & 76 deletions
Original file line numberDiff line numberDiff line change
@@ -76,25 +76,15 @@ public func withIdentifiableContinuation<T>(
7676
onCancel: (IdentifiableContinuation<T, Never>.ID) -> Void
7777
) async -> T {
7878
let id = IdentifiableContinuation<T, Never>.ID()
79-
let state = LockedState(state: (isStarted: false, isCancelled: false))
80-
return await withTaskCancellationHandler {
81-
await withCheckedContinuation(function: function) {
82-
body(IdentifiableContinuation(id: id, storage: .checked($0)))
83-
let isCancelled = state.withCriticalRegion {
84-
$0.isStarted = true
85-
return $0.isCancelled
79+
return await withoutActuallyEscaping(body, onCancel, result: T.self) {
80+
let state = LockedState(body: $0, onCancel: $1)
81+
return await withTaskCancellationHandler {
82+
await withCheckedContinuation(function: function) {
83+
let continuation = IdentifiableContinuation(id: id, storage: .checked($0))
84+
state.start(with: continuation)
8685
}
87-
if isCancelled {
88-
onCancel(id)
89-
}
90-
}
91-
} onCancel: {
92-
let isStarted = state.withCriticalRegion {
93-
$0.isCancelled = true
94-
return $0.isStarted
95-
}
96-
if isStarted {
97-
onCancel(id)
86+
} onCancel: {
87+
state.cancel(withID: id)
9888
}
9989
}
10090
}
@@ -106,25 +96,15 @@ public func withThrowingIdentifiableContinuation<T>(
10696
onCancel: (IdentifiableContinuation<T, Error>.ID) -> Void
10797
) async throws -> T {
10898
let id = IdentifiableContinuation<T, Error>.ID()
109-
let state = LockedState(state: (isStarted: false, isCancelled: false))
110-
return try await withTaskCancellationHandler {
111-
try await withCheckedThrowingContinuation(function: function) {
112-
body(IdentifiableContinuation(id: id, storage: .checked($0)))
113-
let isCancelled = state.withCriticalRegion {
114-
$0.isStarted = true
115-
return $0.isCancelled
99+
return try await withoutActuallyEscaping(body, onCancel, result: T.self) {
100+
let state = LockedState(body: $0, onCancel: $1)
101+
return try await withTaskCancellationHandler {
102+
try await withCheckedThrowingContinuation(function: function) {
103+
let continuation = IdentifiableContinuation(id: id, storage: .checked($0))
104+
state.start(with: continuation)
116105
}
117-
if isCancelled {
118-
onCancel(id)
119-
}
120-
}
121-
} onCancel: {
122-
let isStarted = state.withCriticalRegion {
123-
$0.isCancelled = true
124-
return $0.isStarted
125-
}
126-
if isStarted {
127-
onCancel(id)
106+
} onCancel: {
107+
state.cancel(withID: id)
128108
}
129109
}
130110
}
@@ -135,25 +115,15 @@ public func withIdentifiableUnsafeContinuation<T>(
135115
onCancel: (IdentifiableContinuation<T, Never>.ID) -> Void
136116
) async -> T {
137117
let id = IdentifiableContinuation<T, Never>.ID()
138-
let state = LockedState(state: (isStarted: false, isCancelled: false))
139-
return await withTaskCancellationHandler {
140-
await withUnsafeContinuation {
141-
body(IdentifiableContinuation(id: id, storage: .unsafe($0)))
142-
let isCancelled = state.withCriticalRegion {
143-
$0.isStarted = true
144-
return $0.isCancelled
145-
}
146-
if isCancelled {
147-
onCancel(id)
118+
return await withoutActuallyEscaping(body, onCancel, result: T.self) {
119+
let state = LockedState(body: $0, onCancel: $1)
120+
return await withTaskCancellationHandler {
121+
await withUnsafeContinuation {
122+
let continuation = IdentifiableContinuation(id: id, storage: .unsafe($0))
123+
state.start(with: continuation)
148124
}
149-
}
150-
} onCancel: {
151-
let isStarted = state.withCriticalRegion {
152-
$0.isCancelled = true
153-
return $0.isStarted
154-
}
155-
if isStarted {
156-
onCancel(id)
125+
} onCancel: {
126+
state.cancel(withID: id)
157127
}
158128
}
159129
}
@@ -164,25 +134,15 @@ public func withThrowingIdentifiableUnsafeContinuation<T>(
164134
onCancel: (IdentifiableContinuation<T, Error>.ID) -> Void
165135
) async throws -> T {
166136
let id = IdentifiableContinuation<T, Error>.ID()
167-
let state = LockedState(state: (isStarted: false, isCancelled: false))
168-
return try await withTaskCancellationHandler {
169-
try await withUnsafeThrowingContinuation {
170-
body(IdentifiableContinuation(id: id, storage: .unsafe($0)))
171-
let isCancelled = state.withCriticalRegion {
172-
$0.isStarted = true
173-
return $0.isCancelled
174-
}
175-
if isCancelled {
176-
onCancel(id)
137+
return try await withoutActuallyEscaping(body, onCancel, result: T.self) {
138+
let state = LockedState(body: $0, onCancel: $1)
139+
return try await withTaskCancellationHandler {
140+
try await withUnsafeThrowingContinuation {
141+
let continuation = IdentifiableContinuation(id: id, storage: .unsafe($0))
142+
state.start(with: continuation)
177143
}
178-
}
179-
} onCancel: {
180-
let isStarted = state.withCriticalRegion {
181-
$0.isCancelled = true
182-
return $0.isStarted
183-
}
184-
if isStarted {
185-
onCancel(id)
144+
} onCancel: {
145+
state.cancel(withID: id)
186146
}
187147
}
188148
}
@@ -248,13 +208,49 @@ public struct IdentifiableContinuation<T, E>: Sendable, Identifiable where E : E
248208
}
249209

250210
@usableFromInline
251-
final class LockedState<State> {
211+
final class LockedState<T, Failure: Error>: @unchecked Sendable {
252212
private let lock = NSLock()
253213
private var state: State
214+
private let body: (IdentifiableContinuation<T, Failure>) -> Void
215+
private let onCancel: (IdentifiableContinuation<T, Failure>.ID) -> Void
216+
217+
@usableFromInline
218+
struct State {
219+
@usableFromInline
220+
var isStarted: Bool = false
221+
@usableFromInline
222+
var isCancelled: Bool = false
223+
}
254224

255225
@usableFromInline
256-
init(state: State) {
257-
self.state = state
226+
init(body: @escaping (IdentifiableContinuation<T, Failure>) -> Void,
227+
onCancel: @escaping (IdentifiableContinuation<T, Failure>.ID) -> Void) {
228+
self.state = State()
229+
self.body = body
230+
self.onCancel = onCancel
231+
}
232+
233+
@usableFromInline
234+
func start(with continuation: IdentifiableContinuation<T, Failure>) {
235+
body(continuation)
236+
let isCancelled = withCriticalRegion {
237+
$0.isStarted = true
238+
return $0.isCancelled
239+
}
240+
if isCancelled {
241+
onCancel(continuation.id)
242+
}
243+
}
244+
245+
@usableFromInline
246+
func cancel(withID id: IdentifiableContinuation<T, Failure>.ID) {
247+
let isStarted = withCriticalRegion {
248+
$0.isCancelled = true
249+
return $0.isStarted
250+
}
251+
if isStarted {
252+
onCancel(id)
253+
}
258254
}
259255

260256
@usableFromInline
@@ -265,4 +261,15 @@ final class LockedState<State> {
265261
}
266262
}
267263

268-
extension LockedState: @unchecked Sendable where State: Sendable { }
264+
@usableFromInline
265+
func withoutActuallyEscaping<T, Failure: Error, U>(
266+
_ c1: (IdentifiableContinuation<T, Failure>) -> Void,
267+
_ c2: (IdentifiableContinuation<T, Failure>.ID) -> Void,
268+
result: U.Type,
269+
do body: (@escaping (IdentifiableContinuation<T, Failure>) -> Void, @escaping (IdentifiableContinuation<T, Failure>.ID) -> Void) async throws -> U) async rethrows -> U {
270+
try await withoutActuallyEscaping(c1) { (escapingC1) -> U in
271+
try await withoutActuallyEscaping(c2) { (escapingC2) -> U in
272+
try await body(escapingC1, escapingC2)
273+
}
274+
}
275+
}

0 commit comments

Comments
 (0)