Refactoring: Manage subscribed atom keys in a graph instance

Sources/Atoms/AtomStore.swift

@@ -5,5 +5,5 @@ public final class AtomStore {
     internal var state = StoreState()
 
     /// Creates a new store.
-    nonisolated public init() {}
+    public nonisolated init() {}
 }

Sources/Atoms/Core/Graph.swift

@@ -1,4 +1,5 @@
 internal struct Graph: Equatable {
     var dependencies = [AtomKey: Set<AtomKey>]()
     var children = [AtomKey: Set<AtomKey>]()
+    var subscribed = [SubscriberKey: Set<AtomKey>]()
 }

Sources/Atoms/Core/StoreContext.swift

@@ -111,12 +111,12 @@ internal struct StoreContext {
         let (key, override) = lookupAtomKeyAndOverride(of: atom)
         let cache = lookupCache(of: atom, for: key)
         let value = cache?.value ?? initialize(of: atom, for: key, override: override)
-        let isNewSubscription = subscriber.subscribing.insert(key).inserted
+        let isNewSubscription = store.graph.subscribed[subscriber.key, default: []].insert(key).inserted
 
         if isNewSubscription {
             store.state.subscriptions[key, default: [:]][subscriber.key] = subscription
-            subscriber.unsubscribe = { keys in
-                unsubscribe(keys, for: subscriber.key)
+            subscriber.unsubscribe = {
+                unsubscribeAll(for: subscriber.key)
             }
             notifyUpdateToObservers()
         }
@@ -209,7 +209,7 @@ internal struct StoreContext {
     func unwatch(_ atom: some Atom, subscriber: Subscriber) {
         let (key, _) = lookupAtomKeyAndOverride(of: atom)
 
-        subscriber.subscribing.remove(key)
+        store.graph.subscribed[subscriber.key]?.remove(key)
         unsubscribe([key], for: subscriber.key)
     }
 
@@ -462,6 +462,14 @@ private extension StoreContext {
         }
     }
 
+    func unsubscribeAll(for subscriberKey: SubscriberKey) {
+        let keys = store.graph.subscribed.removeValue(forKey: subscriberKey)
+
+        if let keys {
+            unsubscribe(keys, for: subscriberKey)
+        }
+    }
+
     func unsubscribe(_ keys: some Sequence<AtomKey>, for subscriberKey: SubscriberKey) {
         for key in keys {
             store.state.subscriptions[key]?.removeValue(forKey: subscriberKey)

Sources/Atoms/Core/Subscriber.swift

@@ -10,12 +10,7 @@ internal struct Subscriber {
         self.key = state.token.key
     }
 
-    var subscribing: Set<AtomKey> {
-        get { state?.subscribing ?? [] }
-        nonmutating set { state?.subscribing = newValue }
-    }
-
-    var unsubscribe: (@MainActor (Set<AtomKey>) -> Void)? {
+    var unsubscribe: (@MainActor () -> Void)? {
         get { state?.unsubscribe }
         nonmutating set { state?.unsubscribe = newValue }
     }

Sources/Atoms/Core/SubscriberState.swift

@@ -7,33 +7,26 @@ internal final class SubscriberState {
     #endif
 
     #if compiler(>=6)
-        nonisolated(unsafe) var subscribing = Set<AtomKey>()
-        nonisolated(unsafe) var unsubscribe: (@MainActor (Set<AtomKey>) -> Void)?
+        nonisolated(unsafe) var unsubscribe: (@MainActor () -> Void)?
 
         // TODO: Use isolated synchronous deinit once it's available.
         // 0371-isolated-synchronous-deinit
         deinit {
-            MainActor.performIsolated { [unsubscribe, subscribing] in
-                unsubscribe?(subscribing)
+            MainActor.performIsolated { [unsubscribe] in
+                unsubscribe?()
             }
         }
     #else
-        private var _subscribing = UnsafeUncheckedSendable(Set<AtomKey>())
-        private var _unsubscribe = UnsafeUncheckedSendable<(@MainActor (Set<AtomKey>) -> Void)?>(nil)
+        private var _unsubscribe = UnsafeUncheckedSendable<(@MainActor () -> Void)?>(nil)
 
-        var subscribing: Set<AtomKey> {
-            _read { yield _subscribing.value }
-            _modify { yield &_subscribing.value }
-        }
-
-        var unsubscribe: (@MainActor (Set<AtomKey>) -> Void)? {
+        var unsubscribe: (@MainActor () -> Void)? {
             _read { yield _unsubscribe.value }
             _modify { yield &_unsubscribe.value }
         }
 
         deinit {
-            MainActor.performIsolated { [_unsubscribe, _subscribing] in
-                _unsubscribe.value?(_subscribing.value)
+            MainActor.performIsolated { [_unsubscribe] in
+                _unsubscribe.value?()
             }
         }
     #endif

Tests/AtomsTests/Core/StoreContextTests.swift

@@ -157,17 +157,27 @@ final class StoreContextTests: XCTestCase {
         )
 
         XCTAssertEqual(context.watch(dependency0, in: transactionState), 0)
-        XCTAssertEqual(store.graph.dependencies, [key: [dependency0Key]])
-        XCTAssertEqual(store.graph.children, [dependency0Key: [key]])
+        XCTAssertEqual(
+            store.graph,
+            Graph(
+                dependencies: [key: [dependency0Key]],
+                children: [dependency0Key: [key]]
+            )
+        )
         XCTAssertEqual((store.state.caches[dependency0Key] as? AtomCache<TestStateAtom<Int>>)?.value, 0)
         XCTAssertNotNil(store.state.states[dependency0Key])
         XCTAssertTrue(snapshots.flatMap(\.caches).isEmpty)
 
         transactionState.terminate()
 
         XCTAssertEqual(context.watch(dependency1, in: transactionState), 1)
-        XCTAssertEqual(store.graph.dependencies, [key: [dependency0Key]])
-        XCTAssertEqual(store.graph.children, [dependency0Key: [key]])
+        XCTAssertEqual(
+            store.graph,
+            Graph(
+                dependencies: [key: [dependency0Key]],
+                children: [dependency0Key: [key]]
+            )
+        )
         XCTAssertNil(store.state.caches[dependency1Key])
         XCTAssertNil(store.state.states[dependency1Key])
         XCTAssertTrue(snapshots.isEmpty)
@@ -214,7 +224,7 @@ final class StoreContextTests: XCTestCase {
         )
 
         XCTAssertEqual(initialValue, 0)
-        XCTAssertTrue(subscriber.subscribing.contains(key))
+        XCTAssertTrue(store.graph.subscribed[subscriber.key]?.contains(key) ?? false)
         XCTAssertNotNil(store.state.subscriptions[key]?[subscriber.key])
         XCTAssertEqual((store.state.caches[key] as? AtomCache<TestAtom>)?.value, 0)
         XCTAssertEqual((store.state.caches[dependencyKey] as? AtomCache<DependencyAtom>)?.value, 0)
@@ -226,11 +236,15 @@ final class StoreContextTests: XCTestCase {
         snapshots.removeAll()
         store.state.subscriptions[key]?[subscriber.key]?.update()
         subscriberState = nil
+
         XCTAssertEqual(updateCount, 1)
+        XCTAssertNil(store.graph.subscribed[subscriber.key])
         XCTAssertNil(store.state.caches[key])
         XCTAssertNil(store.state.states[key])
+        XCTAssertNil(store.state.subscriptions[key])
         XCTAssertNil(store.state.caches[dependencyKey])
         XCTAssertNil(store.state.states[dependencyKey])
+        XCTAssertNil(store.state.subscriptions[dependencyKey])
         XCTAssertEqual(
             snapshots.map { $0.caches.mapValues { $0.value as? Int } },
             [[:]]
@@ -372,11 +386,21 @@ final class StoreContextTests: XCTestCase {
         )
 
         _ = context.watch(atom, subscriber: subscriber, subscription: Subscription())
-        snapshots.removeAll()
         context.unwatch(atom, subscriber: subscriber)
+
         XCTAssertEqual(
             snapshots.map { $0.caches.mapValues { $0.value as? Int } },
-            [[:]]
+            [
+                [AtomKey(atom): 0],
+                [:],
+            ]
+        )
+        XCTAssertEqual(
+            snapshots.map(\.graph),
+            [
+                Graph(subscribed: [subscriber.key: [AtomKey(atom)]]),
+                Graph(subscribed: [subscriber.key: []]),
+            ]
         )
     }
 
@@ -568,7 +592,7 @@ final class StoreContextTests: XCTestCase {
         context.unwatch(dependency1Atom, subscriber: subscriber)
         context.unwatch(dependency2Atom, subscriber: subscriber)
         scoped1Context.unwatch(dependency1Atom, subscriber: subscriber)
-        scoped2Context.unwatch(dependency1Atom, subscriber: subscriber)
+        scoped2Context.unwatch(dependency2Atom, subscriber: subscriber)
 
         // Override for `scoped1Context` shouldn't inherited to `scoped2Context`.
         XCTAssertEqual(scoped2Context.watch(atom, subscriber: subscriber, subscription: Subscription()), 21)
@@ -644,6 +668,12 @@ final class StoreContextTests: XCTestCase {
                         AtomKey(atom),
                         AtomKey(publisherAtom),
                     ],
+                ],
+                subscribed: [
+                    subscriber.key: [
+                        AtomKey((atom)),
+                        AtomKey((publisherAtom)),
+                    ]
                 ]
             )
         )
@@ -1014,6 +1044,11 @@ final class StoreContextTests: XCTestCase {
                 AtomKey(TestDependency1Atom()): [AtomKey(TestAtom())],
                 AtomKey(TestDependency2Atom()): [AtomKey(TestAtom())],
                 AtomKey(TestDependency3Atom(), scopeKey: scopeToken.key): [AtomKey(TestAtom())],
+            ],
+            subscribed: [
+                subscriber.key: [
+                    AtomKey(atom)
+                ]
             ]
         )
 

Tests/AtomsTests/Core/SubscriberStateTests.swift

@@ -8,7 +8,7 @@ final class SubscriberStateTests: XCTestCase {
         var subscriberState: SubscriberState? = SubscriberState()
         var unsubscribedCount = 0
 
-        subscriberState!.unsubscribe = { _ in
+        subscriberState!.unsubscribe = {
             unsubscribedCount += 1
         }
 

Tests/AtomsTests/Core/SubscriberTestsTests.swift → Tests/AtomsTests/Core/SubscriberTests.swift

@@ -23,42 +23,23 @@ final class SubscriberTests: XCTestCase {
         )
     }
 
-    @MainActor
-    func testSubscribing() {
-        var state: SubscriberState? = SubscriberState()
-        let subscriber = Subscriber(state!)
-        let expected: Set = [
-            AtomKey(TestAtom(value: 0)),
-            AtomKey(TestAtom(value: 1)),
-            AtomKey(TestAtom(value: 2)),
-        ]
-
-        subscriber.subscribing = expected
-
-        XCTAssertEqual(state?.subscribing, expected)
-
-        state = nil
-
-        XCTAssertTrue(subscriber.subscribing.isEmpty)
-    }
-
     @MainActor
     func testUnsubscribe() {
         var state: SubscriberState? = SubscriberState()
         let subscriber = Subscriber(state!)
         var isUnsubscribed = false
 
-        subscriber.unsubscribe = { _ in
+        subscriber.unsubscribe = {
             isUnsubscribed = true
         }
 
-        state?.unsubscribe?([])
+        state?.unsubscribe?()
 
         XCTAssertTrue(isUnsubscribed)
 
         state = nil
         isUnsubscribed = false
-        subscriber.unsubscribe?([])
+        subscriber.unsubscribe?()
 
         XCTAssertFalse(isUnsubscribed)
     }