Skip to content

Commit df7d6cc

Browse files
ra1028ra1028
authored
Merge pull request #23 from ra1028/feat/fix-atom-lifecycle
fix: Lifecycle and unsubscription problem
2 parents a8c210a + 3bf14ea commit df7d6cc

File tree

11 files changed

+102
-45
lines changed

11 files changed

+102
-45
lines changed

‎Sources/Atoms/AtomRoot.swift‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ import SwiftUI
4242
public struct AtomRoot<Content: View>: View {
4343
@StateObject
4444
private var state: State
45-
private var overrides: Overrides
45+
private var overrides = Overrides()
4646
private var observers = [AtomObserver]()
4747
private let content: Content
4848

@@ -51,7 +51,6 @@ public struct AtomRoot<Content: View>: View {
5151
/// - Parameter content: The content that uses atoms.
5252
public init(@ViewBuilder content: () -> Content) {
5353
self._state = StateObject(wrappedValue: State())
54-
self.overrides = Overrides()
5554
self.content = content()
5655
}
5756

‎Sources/Atoms/Context/AtomTestContext.swift‎

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -275,18 +275,13 @@ private extension AtomTestContext {
275275
@MainActor
276276
final class State {
277277
private let _store = Store()
278+
private let _container = SubscriptionContainer()
278279

279-
let container: SubscriptionContainer
280-
var overrides: Overrides
280+
var overrides = Overrides()
281281
var observers = [AtomObserver]()
282282
let notifier = PassthroughSubject<Void, Never>()
283283
var onUpdate: (() -> Void)?
284284

285-
init() {
286-
overrides = Overrides()
287-
container = SubscriptionContainer()
288-
}
289-
290285
var store: StoreContext {
291286
StoreContext(
292287
_store,
@@ -295,6 +290,10 @@ private extension AtomTestContext {
295290
)
296291
}
297292

293+
var container: SubscriptionContainer.Wrapper {
294+
_container.wrapper
295+
}
296+
298297
func notifyUpdate() {
299298
onUpdate?()
300299
notifier.send()

‎Sources/Atoms/Context/AtomViewContext.swift‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,13 +7,13 @@ public struct AtomViewContext: AtomWatchableContext {
77
@usableFromInline
88
internal let _store: StoreContext
99
@usableFromInline
10-
internal let _container: SubscriptionContainer
10+
internal let _container: SubscriptionContainer.Wrapper
1111
@usableFromInline
1212
internal let _notifyUpdate: () -> Void
1313

1414
internal init(
1515
store: StoreContext,
16-
container: SubscriptionContainer,
16+
container: SubscriptionContainer.Wrapper,
1717
notifyUpdate: @escaping () -> Void
1818
) {
1919
_store = store

‎Sources/Atoms/Core/Overrides.swift‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@ internal struct Overrides {
33
private var _entriesForNode = [AtomKey: Override]()
44
private var _entriesForType = [AtomTypeKey: Override]()
55

6+
nonisolated init() {}
7+
68
mutating func insert<Node: Atom>(
79
_ atom: Node,
810
with value: @escaping (Node) -> Node.Loader.Value

‎Sources/Atoms/Core/StoreContext.swift‎

Lines changed: 23 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,13 @@ internal struct StoreContext {
1919

2020
@usableFromInline
2121
func read<Node: Atom>(_ atom: Node) -> Node.Loader.Value {
22-
getValue(for: atom)
22+
let key = AtomKey(atom)
23+
24+
// Register if it doesn't exist yet because the atom needs to be maintained if it's marked as `KeepAlive`.
25+
registerIfAbsent(atom: atom)
26+
defer { checkRelease(for: key) }
27+
28+
return getValue(for: atom)
2329
}
2430

2531
@usableFromInline
@@ -53,7 +59,7 @@ internal struct StoreContext {
5359
func watch<Node: Atom>(_ atom: Node, in transaction: Transaction) -> Node.Loader.Value {
5460
// Return a new value immediately if the transaction is already terminated.
5561
guard !transaction.isTerminated else {
56-
return getNewValue(for: atom)
62+
return read(atom)
5763
}
5864

5965
let store = getStore()
@@ -72,19 +78,18 @@ internal struct StoreContext {
7278
@usableFromInline
7379
func watch<Node: Atom>(
7480
_ atom: Node,
75-
container: SubscriptionContainer,
81+
container: SubscriptionContainer.Wrapper,
7682
notifyUpdate: @escaping () -> Void
7783
) -> Node.Loader.Value {
7884
let store = getStore()
7985
let key = AtomKey(atom)
80-
let subscriptionKey = SubscriptionKey(container)
8186
let subscription = Subscription(notifyUpdate: notifyUpdate) { [weak store] in
8287
guard let store = store else {
8388
return
8489
}
8590

8691
// Unsubscribe and release if it's no longer used.
87-
store.state.subscriptions[key]?.removeValue(forKey: subscriptionKey)
92+
store.state.subscriptions[key]?.removeValue(forKey: container.key)
8893
checkRelease(for: key)
8994
}
9095

@@ -93,13 +98,19 @@ internal struct StoreContext {
9398

9499
// Register the subscription to both the store and the container.
95100
container.subscriptions[key] = subscription
96-
store.state.subscriptions[key, default: [:]].updateValue(subscription, forKey: subscriptionKey)
101+
store.state.subscriptions[key, default: [:]].updateValue(subscription, forKey: container.key)
97102

98103
return getValue(for: atom)
99104
}
100105

101106
@usableFromInline
102107
func refresh<Node: Atom>(_ atom: Node) async -> Node.Loader.Value where Node.Loader: RefreshableAtomLoader {
108+
let key = AtomKey(atom)
109+
110+
// Register if it doesn't exist yet because the atom needs to be maintained if it's marked as `KeepAlive`.
111+
registerIfAbsent(atom: atom)
112+
defer { checkRelease(for: key) }
113+
103114
let context = prepareTransaction(for: atom)
104115
let value: Node.Loader.Value
105116

@@ -117,6 +128,12 @@ internal struct StoreContext {
117128

118129
@usableFromInline
119130
func reset<Node: Atom>(_ atom: Node) {
131+
let key = AtomKey(atom)
132+
133+
// Register if it doesn't exist yet because the atom needs to be maintained if it's marked as `KeepAlive`.
134+
registerIfAbsent(atom: atom)
135+
defer { checkRelease(for: key) }
136+
120137
let value = getNewValue(for: atom)
121138
update(atom: atom, with: value)
122139
}

‎Sources/Atoms/Core/SubscriptionContainer.swift‎

Lines changed: 27 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,37 @@
11
@usableFromInline
22
@MainActor
33
internal final class SubscriptionContainer {
4-
var subscriptions = [AtomKey: Subscription]()
4+
private var subscriptions = [AtomKey: Subscription]()
5+
6+
var wrapper: Wrapper {
7+
Wrapper(self)
8+
}
9+
10+
nonisolated init() {}
511

612
deinit {
713
for subscription in ContiguousArray(subscriptions.values) {
814
subscription.unsubscribe()
915
}
1016
}
1117
}
18+
19+
internal extension SubscriptionContainer {
20+
@usableFromInline
21+
@MainActor
22+
struct Wrapper {
23+
private weak var container: SubscriptionContainer?
24+
25+
let key: SubscriptionKey
26+
27+
var subscriptions: [AtomKey: Subscription] {
28+
get { container?.subscriptions ?? [:] }
29+
nonmutating set { container?.subscriptions = newValue }
30+
}
31+
32+
init(_ container: SubscriptionContainer) {
33+
self.container = container
34+
self.key = SubscriptionKey(container)
35+
}
36+
}
37+
}

‎Sources/Atoms/PropertyWrapper/ViewContext.swift‎

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ public struct ViewContext: DynamicProperty {
5252
public var wrappedValue: AtomViewContext {
5353
AtomViewContext(
5454
store: _store,
55-
container: state.container,
55+
container: state.container.wrapper,
5656
notifyUpdate: state.objectWillChange.send
5757
)
5858
}
@@ -61,10 +61,6 @@ public struct ViewContext: DynamicProperty {
6161
private extension ViewContext {
6262
@MainActor
6363
final class State: ObservableObject {
64-
let container: SubscriptionContainer
65-
66-
init() {
67-
container = SubscriptionContainer()
68-
}
64+
let container = SubscriptionContainer()
6965
}
7066
}

‎Tests/AtomsTests/Context/AtomTestContextTests.swift‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -92,15 +92,15 @@ final class AtomTestContextTests: XCTestCase {
9292

9393
for observer in observers {
9494
XCTAssertEqual(observer.assignedAtomKeys, [key])
95-
XCTAssertEqual(observer.unassignedAtomKeys, [])
95+
XCTAssertTrue(observer.unassignedAtomKeys.isEmpty)
9696
XCTAssertEqual(observer.changedAtomKeys, [key])
9797
}
9898

9999
context[atom] = 200
100100

101101
for observer in observers {
102102
XCTAssertEqual(observer.assignedAtomKeys, [key])
103-
XCTAssertEqual(observer.unassignedAtomKeys, [])
103+
XCTAssertTrue(observer.unassignedAtomKeys.isEmpty)
104104
XCTAssertEqual(observer.changedAtomKeys, [key, key])
105105
}
106106

‎Tests/AtomsTests/Context/AtomViewContextTests.swift‎

Lines changed: 23 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ final class AtomViewContextTests: XCTestCase {
1010
let container = SubscriptionContainer()
1111
let context = AtomViewContext(
1212
store: StoreContext(store),
13-
container: container,
13+
container: container.wrapper,
1414
notifyUpdate: {}
1515
)
1616

@@ -23,7 +23,7 @@ final class AtomViewContextTests: XCTestCase {
2323
let container = SubscriptionContainer()
2424
let context = AtomViewContext(
2525
store: StoreContext(store),
26-
container: container,
26+
container: container.wrapper,
2727
notifyUpdate: {}
2828
)
2929

@@ -40,7 +40,7 @@ final class AtomViewContextTests: XCTestCase {
4040
let container = SubscriptionContainer()
4141
let context = AtomViewContext(
4242
store: StoreContext(store),
43-
container: container,
43+
container: container.wrapper,
4444
notifyUpdate: {}
4545
)
4646

@@ -57,7 +57,7 @@ final class AtomViewContextTests: XCTestCase {
5757
let container = SubscriptionContainer()
5858
let context = AtomViewContext(
5959
store: StoreContext(store),
60-
container: container,
60+
container: container.wrapper,
6161
notifyUpdate: {}
6262
)
6363

@@ -78,7 +78,7 @@ final class AtomViewContextTests: XCTestCase {
7878
let container = SubscriptionContainer()
7979
let context = AtomViewContext(
8080
store: StoreContext(store),
81-
container: container,
81+
container: container.wrapper,
8282
notifyUpdate: {}
8383
)
8484

@@ -88,4 +88,22 @@ final class AtomViewContextTests: XCTestCase {
8888

8989
XCTAssertEqual(context.watch(atom), 200)
9090
}
91+
92+
func testUnsubscription() {
93+
let atom = TestValueAtom(value: 100)
94+
let key = AtomKey(atom)
95+
let store = Store()
96+
var container: SubscriptionContainer? = SubscriptionContainer()
97+
let context = AtomViewContext(
98+
store: StoreContext(store),
99+
container: container!.wrapper,
100+
notifyUpdate: {}
101+
)
102+
103+
context.watch(atom)
104+
XCTAssertNotNil(store.state.atomStates[key])
105+
106+
container = nil
107+
XCTAssertNil(store.state.atomStates[key])
108+
}
91109
}

‎Tests/AtomsTests/Core/StoreContextTests.swift‎

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,9 @@ final class StoreContextTests: XCTestCase {
1212

1313
XCTAssertEqual(context.read(atom), 0)
1414
XCTAssertNil(store.state.atomStates[AtomKey(atom)])
15-
XCTAssertTrue(observer.assignedAtomKeys.isEmpty)
15+
XCTAssertEqual(observer.assignedAtomKeys, [AtomKey(atom)])
1616
XCTAssertEqual(observer.changedAtomKeys, [AtomKey(atom)])
17-
XCTAssertTrue(observer.unassignedAtomKeys.isEmpty)
17+
XCTAssertEqual(observer.unassignedAtomKeys, [AtomKey(atom)])
1818

1919
store.state.atomStates[AtomKey(atom)] = ConcreteAtomState(atom: atom, value: 1)
2020

@@ -95,9 +95,9 @@ final class StoreContextTests: XCTestCase {
9595
XCTAssertEqual(store.graph.dependencies, [key: [dependency0Key]])
9696
XCTAssertEqual(store.graph.children, [dependency0Key: [key]])
9797
XCTAssertNil(store.state.atomStates[dependency1Key])
98-
XCTAssertEqual(observer.assignedAtomKeys, [dependency0Key])
99-
XCTAssertEqual(observer.changedAtomKeys, [dependency0Key])
100-
XCTAssertTrue(observer.unassignedAtomKeys.isEmpty)
98+
XCTAssertEqual(observer.assignedAtomKeys, [dependency0Key, dependency1Key])
99+
XCTAssertEqual(observer.changedAtomKeys, [dependency0Key, dependency1Key])
100+
XCTAssertEqual(observer.unassignedAtomKeys, [dependency1Key])
101101
}
102102

103103
func testWatchFromView() {
@@ -123,12 +123,12 @@ final class StoreContextTests: XCTestCase {
123123
let key = AtomKey(atom)
124124
let dependencyKey = AtomKey(dependency)
125125
var updateCount = 0
126-
let initialValue = context.watch(atom, container: container) {
126+
let initialValue = context.watch(atom, container: container.wrapper) {
127127
updateCount += 1
128128
}
129129

130130
XCTAssertEqual(initialValue, 0)
131-
XCTAssertNotNil(container.subscriptions[key])
131+
XCTAssertNotNil(container.wrapper.subscriptions[key])
132132
XCTAssertNotNil(store.state.subscriptions[key]?[subscriptionKey])
133133
XCTAssertEqual((store.state.atomStates[key] as? ConcreteAtomState<TestAtom>)?.value, 0)
134134
XCTAssertEqual((store.state.atomStates[dependencyKey] as? ConcreteAtomState<DependencyAtom>)?.value, 0)
@@ -160,13 +160,13 @@ final class StoreContextTests: XCTestCase {
160160

161161
XCTAssertEqual(value0, 0)
162162
XCTAssertNil(store.state.atomStates[key])
163-
XCTAssertTrue(observer.assignedAtomKeys.isEmpty)
163+
XCTAssertEqual(observer.assignedAtomKeys, [key])
164164
XCTAssertEqual(observer.changedAtomKeys, [key])
165-
XCTAssertTrue(observer.unassignedAtomKeys.isEmpty)
165+
XCTAssertEqual(observer.unassignedAtomKeys, [key])
166166

167167
var updateCount = 0
168168

169-
_ = context.watch(atom, container: container) {
169+
_ = context.watch(atom, container: container.wrapper) {
170170
updateCount += 1
171171
}
172172

@@ -186,7 +186,7 @@ final class StoreContextTests: XCTestCase {
186186
let atom = TestStateAtom(defaultValue: 0)
187187
let key = AtomKey(atom)
188188
var updateCount = 0
189-
let initialValue = context.watch(atom, container: container) {
189+
let initialValue = context.watch(atom, container: container.wrapper) {
190190
updateCount += 1
191191
}
192192

@@ -218,7 +218,7 @@ final class StoreContextTests: XCTestCase {
218218
let context = StoreContext(store, observers: [observer0])
219219
let relayedContext = context.relay(observers: [observer1])
220220

221-
_ = relayedContext.watch(atom, container: container) {}
221+
_ = relayedContext.watch(atom, container: container.wrapper) {}
222222

223223
XCTAssertFalse(observer0.assignedAtomKeys.isEmpty)
224224
XCTAssertFalse(observer1.assignedAtomKeys.isEmpty)
@@ -324,7 +324,7 @@ final class StoreContextTests: XCTestCase {
324324
let phase = PhaseAtom()
325325

326326
func watch() async -> Int {
327-
await atomStore.watch(atom, container: container, notifyUpdate: {}).value
327+
await atomStore.watch(atom, container: container.wrapper, notifyUpdate: {}).value
328328
}
329329

330330
do {

0 commit comments

Comments
 (0)