Skip to content

Commit 74ed60c

Browse files
ra1028ra1028
authored
Refactoring: Manage subscribed atom keys in a graph instance (#177)
* Manage subscribed atom keys in a graph instance * Typo * Fix for compilation failure with Swift 5
1 parent e6b9815 commit 74ed60c

File tree

8 files changed

+69
-56
lines changed

8 files changed

+69
-56
lines changed

‎Sources/Atoms/AtomStore.swift‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,5 +5,5 @@ public final class AtomStore {
55
internal var state = StoreState()
66

77
/// Creates a new store.
8-
nonisolated public init() {}
8+
public nonisolated init() {}
99
}

‎Sources/Atoms/Core/Graph.swift‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
internal struct Graph: Equatable {
22
var dependencies = [AtomKey: Set<AtomKey>]()
33
var children = [AtomKey: Set<AtomKey>]()
4+
var subscribed = [SubscriberKey: Set<AtomKey>]()
45
}

‎Sources/Atoms/Core/StoreContext.swift‎

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -111,12 +111,12 @@ internal struct StoreContext {
111111
let (key, override) = lookupAtomKeyAndOverride(of: atom)
112112
let cache = lookupCache(of: atom, for: key)
113113
let value = cache?.value ?? initialize(of: atom, for: key, override: override)
114-
let isNewSubscription = subscriber.subscribing.insert(key).inserted
114+
let isNewSubscription = store.graph.subscribed[subscriber.key, default: []].insert(key).inserted
115115

116116
if isNewSubscription {
117117
store.state.subscriptions[key, default: [:]][subscriber.key] = subscription
118-
subscriber.unsubscribe = { keys in
119-
unsubscribe(keys, for: subscriber.key)
118+
subscriber.unsubscribe = {
119+
unsubscribeAll(for: subscriber.key)
120120
}
121121
notifyUpdateToObservers()
122122
}
@@ -209,7 +209,7 @@ internal struct StoreContext {
209209
func unwatch(_ atom: some Atom, subscriber: Subscriber) {
210210
let (key, _) = lookupAtomKeyAndOverride(of: atom)
211211

212-
subscriber.subscribing.remove(key)
212+
store.graph.subscribed[subscriber.key]?.remove(key)
213213
unsubscribe([key], for: subscriber.key)
214214
}
215215

@@ -462,6 +462,14 @@ private extension StoreContext {
462462
}
463463
}
464464

465+
func unsubscribeAll(for subscriberKey: SubscriberKey) {
466+
let keys = store.graph.subscribed.removeValue(forKey: subscriberKey)
467+
468+
if let keys {
469+
unsubscribe(keys, for: subscriberKey)
470+
}
471+
}
472+
465473
func unsubscribe(_ keys: some Sequence<AtomKey>, for subscriberKey: SubscriberKey) {
466474
for key in keys {
467475
store.state.subscriptions[key]?.removeValue(forKey: subscriberKey)

‎Sources/Atoms/Core/Subscriber.swift‎

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -10,12 +10,7 @@ internal struct Subscriber {
1010
self.key = state.token.key
1111
}
1212

13-
var subscribing: Set<AtomKey> {
14-
get { state?.subscribing ?? [] }
15-
nonmutating set { state?.subscribing = newValue }
16-
}
17-
18-
var unsubscribe: (@MainActor (Set<AtomKey>) -> Void)? {
13+
var unsubscribe: (@MainActor () -> Void)? {
1914
get { state?.unsubscribe }
2015
nonmutating set { state?.unsubscribe = newValue }
2116
}

‎Sources/Atoms/Core/SubscriberState.swift‎

Lines changed: 7 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -7,33 +7,26 @@ internal final class SubscriberState {
77
#endif
88

99
#if compiler(>=6)
10-
nonisolated(unsafe) var subscribing = Set<AtomKey>()
11-
nonisolated(unsafe) var unsubscribe: (@MainActor (Set<AtomKey>) -> Void)?
10+
nonisolated(unsafe) var unsubscribe: (@MainActor () -> Void)?
1211

1312
// TODO: Use isolated synchronous deinit once it's available.
1413
// 0371-isolated-synchronous-deinit
1514
deinit {
16-
MainActor.performIsolated { [unsubscribe, subscribing] in
17-
unsubscribe?(subscribing)
15+
MainActor.performIsolated { [unsubscribe] in
16+
unsubscribe?()
1817
}
1918
}
2019
#else
21-
private var _subscribing = UnsafeUncheckedSendable(Set<AtomKey>())
22-
private var _unsubscribe = UnsafeUncheckedSendable<(@MainActor (Set<AtomKey>) -> Void)?>(nil)
20+
private var _unsubscribe = UnsafeUncheckedSendable<(@MainActor () -> Void)?>(nil)
2321

24-
var subscribing: Set<AtomKey> {
25-
_read { yield _subscribing.value }
26-
_modify { yield &_subscribing.value }
27-
}
28-
29-
var unsubscribe: (@MainActor (Set<AtomKey>) -> Void)? {
22+
var unsubscribe: (@MainActor () -> Void)? {
3023
_read { yield _unsubscribe.value }
3124
_modify { yield &_unsubscribe.value }
3225
}
3326

3427
deinit {
35-
MainActor.performIsolated { [_unsubscribe, _subscribing] in
36-
_unsubscribe.value?(_subscribing.value)
28+
MainActor.performIsolated { [_unsubscribe] in
29+
_unsubscribe.value?()
3730
}
3831
}
3932
#endif

‎Tests/AtomsTests/Core/StoreContextTests.swift‎

Lines changed: 43 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -157,17 +157,27 @@ final class StoreContextTests: XCTestCase {
157157
)
158158

159159
XCTAssertEqual(context.watch(dependency0, in: transactionState), 0)
160-
XCTAssertEqual(store.graph.dependencies, [key: [dependency0Key]])
161-
XCTAssertEqual(store.graph.children, [dependency0Key: [key]])
160+
XCTAssertEqual(
161+
store.graph,
162+
Graph(
163+
dependencies: [key: [dependency0Key]],
164+
children: [dependency0Key: [key]]
165+
)
166+
)
162167
XCTAssertEqual((store.state.caches[dependency0Key] as? AtomCache<TestStateAtom<Int>>)?.value, 0)
163168
XCTAssertNotNil(store.state.states[dependency0Key])
164169
XCTAssertTrue(snapshots.flatMap(\.caches).isEmpty)
165170

166171
transactionState.terminate()
167172

168173
XCTAssertEqual(context.watch(dependency1, in: transactionState), 1)
169-
XCTAssertEqual(store.graph.dependencies, [key: [dependency0Key]])
170-
XCTAssertEqual(store.graph.children, [dependency0Key: [key]])
174+
XCTAssertEqual(
175+
store.graph,
176+
Graph(
177+
dependencies: [key: [dependency0Key]],
178+
children: [dependency0Key: [key]]
179+
)
180+
)
171181
XCTAssertNil(store.state.caches[dependency1Key])
172182
XCTAssertNil(store.state.states[dependency1Key])
173183
XCTAssertTrue(snapshots.isEmpty)
@@ -214,7 +224,7 @@ final class StoreContextTests: XCTestCase {
214224
)
215225

216226
XCTAssertEqual(initialValue, 0)
217-
XCTAssertTrue(subscriber.subscribing.contains(key))
227+
XCTAssertTrue(store.graph.subscribed[subscriber.key]?.contains(key) ?? false)
218228
XCTAssertNotNil(store.state.subscriptions[key]?[subscriber.key])
219229
XCTAssertEqual((store.state.caches[key] as? AtomCache<TestAtom>)?.value, 0)
220230
XCTAssertEqual((store.state.caches[dependencyKey] as? AtomCache<DependencyAtom>)?.value, 0)
@@ -226,11 +236,15 @@ final class StoreContextTests: XCTestCase {
226236
snapshots.removeAll()
227237
store.state.subscriptions[key]?[subscriber.key]?.update()
228238
subscriberState = nil
239+
229240
XCTAssertEqual(updateCount, 1)
241+
XCTAssertNil(store.graph.subscribed[subscriber.key])
230242
XCTAssertNil(store.state.caches[key])
231243
XCTAssertNil(store.state.states[key])
244+
XCTAssertNil(store.state.subscriptions[key])
232245
XCTAssertNil(store.state.caches[dependencyKey])
233246
XCTAssertNil(store.state.states[dependencyKey])
247+
XCTAssertNil(store.state.subscriptions[dependencyKey])
234248
XCTAssertEqual(
235249
snapshots.map { $0.caches.mapValues { $0.value as? Int } },
236250
[[:]]
@@ -372,11 +386,21 @@ final class StoreContextTests: XCTestCase {
372386
)
373387

374388
_ = context.watch(atom, subscriber: subscriber, subscription: Subscription())
375-
snapshots.removeAll()
376389
context.unwatch(atom, subscriber: subscriber)
390+
377391
XCTAssertEqual(
378392
snapshots.map { $0.caches.mapValues { $0.value as? Int } },
379-
[[:]]
393+
[
394+
[AtomKey(atom): 0],
395+
[:],
396+
]
397+
)
398+
XCTAssertEqual(
399+
snapshots.map(\.graph),
400+
[
401+
Graph(subscribed: [subscriber.key: [AtomKey(atom)]]),
402+
Graph(subscribed: [subscriber.key: []]),
403+
]
380404
)
381405
}
382406

@@ -568,7 +592,7 @@ final class StoreContextTests: XCTestCase {
568592
context.unwatch(dependency1Atom, subscriber: subscriber)
569593
context.unwatch(dependency2Atom, subscriber: subscriber)
570594
scoped1Context.unwatch(dependency1Atom, subscriber: subscriber)
571-
scoped2Context.unwatch(dependency1Atom, subscriber: subscriber)
595+
scoped2Context.unwatch(dependency2Atom, subscriber: subscriber)
572596

573597
// Override for `scoped1Context` shouldn't inherited to `scoped2Context`.
574598
XCTAssertEqual(scoped2Context.watch(atom, subscriber: subscriber, subscription: Subscription()), 21)
@@ -644,6 +668,12 @@ final class StoreContextTests: XCTestCase {
644668
AtomKey(atom),
645669
AtomKey(publisherAtom),
646670
],
671+
],
672+
subscribed: [
673+
subscriber.key: [
674+
AtomKey((atom)),
675+
AtomKey((publisherAtom)),
676+
]
647677
]
648678
)
649679
)
@@ -1014,6 +1044,11 @@ final class StoreContextTests: XCTestCase {
10141044
AtomKey(TestDependency1Atom()): [AtomKey(TestAtom())],
10151045
AtomKey(TestDependency2Atom()): [AtomKey(TestAtom())],
10161046
AtomKey(TestDependency3Atom(), scopeKey: scopeToken.key): [AtomKey(TestAtom())],
1047+
],
1048+
subscribed: [
1049+
subscriber.key: [
1050+
AtomKey(atom)
1051+
]
10171052
]
10181053
)
10191054

‎Tests/AtomsTests/Core/SubscriberStateTests.swift‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ final class SubscriberStateTests: XCTestCase {
88
var subscriberState: SubscriberState? = SubscriberState()
99
var unsubscribedCount = 0
1010

11-
subscriberState!.unsubscribe = { _ in
11+
subscriberState!.unsubscribe = {
1212
unsubscribedCount += 1
1313
}
1414

‎Tests/AtomsTests/Core/SubscriberTestsTests.swift‎ renamed to ‎Tests/AtomsTests/Core/SubscriberTests.swift‎

Lines changed: 3 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -23,42 +23,23 @@ final class SubscriberTests: XCTestCase {
2323
)
2424
}
2525

26-
@MainActor
27-
func testSubscribing() {
28-
var state: SubscriberState? = SubscriberState()
29-
let subscriber = Subscriber(state!)
30-
let expected: Set = [
31-
AtomKey(TestAtom(value: 0)),
32-
AtomKey(TestAtom(value: 1)),
33-
AtomKey(TestAtom(value: 2)),
34-
]
35-
36-
subscriber.subscribing = expected
37-
38-
XCTAssertEqual(state?.subscribing, expected)
39-
40-
state = nil
41-
42-
XCTAssertTrue(subscriber.subscribing.isEmpty)
43-
}
44-
4526
@MainActor
4627
func testUnsubscribe() {
4728
var state: SubscriberState? = SubscriberState()
4829
let subscriber = Subscriber(state!)
4930
var isUnsubscribed = false
5031

51-
subscriber.unsubscribe = { _ in
32+
subscriber.unsubscribe = {
5233
isUnsubscribed = true
5334
}
5435

55-
state?.unsubscribe?([])
36+
state?.unsubscribe?()
5637

5738
XCTAssertTrue(isUnsubscribed)
5839

5940
state = nil
6041
isUnsubscribed = false
61-
subscriber.unsubscribe?([])
42+
subscriber.unsubscribe?()
6243

6344
XCTAssertFalse(isUnsubscribed)
6445
}

0 commit comments

Comments
 (0)