Minor refactoring (#88)

ra1028 · Nov 28, 2023 · c76edf2 · c76edf2
1 parent e4d1404
commit c76edf2
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 17 deletions.
diff --git a/Sources/Atoms/Core/StoreContext.swift b/Sources/Atoms/Core/StoreContext.swift
@@ -103,12 +103,12 @@ internal struct StoreContext {
         let key = AtomKey(atom, overrideScopeKey: override?.scopeKey)
         let cache = lookupCache(of: atom, for: key)
         let newCache = cache ?? makeNewCache(of: atom, for: key, override: override)
-        let isInserted = store.graph.children[key, default: []].insert(transaction.key).inserted
+        let isNew = store.graph.children[key, default: []].insert(transaction.key).inserted
 
         // Add an `Edge` from the upstream to downstream.
         store.graph.dependencies[transaction.key, default: []].insert(key)
 
-        if isInserted || cache == nil {
+        if isNew || cache == nil {
             notifyUpdateToObservers()
         }
 
@@ -132,15 +132,14 @@ internal struct StoreContext {
             requiresObjectUpdate: requiresObjectUpdate,
             notifyUpdate: notifyUpdate
         )
-        let isInserted = store.state.subscriptions[key, default: [:]].updateValue(subscription, forKey: container.key) == nil
 
-        // Register the subscription to both the store and the container.
-        container.subscriptions[key] = subscription
+        store.state.subscriptions[key, default: [:]].updateValue(subscription, forKey: container.key)
         container.unsubscribe = { keys in
             unsubscribe(keys, for: container.key)
         }
+        let isNew = container.subscribingKeys.insert(key).inserted
 
-        if isInserted || cache == nil {
+        if isNew || cache == nil {
             notifyUpdateToObservers()
         }
 
@@ -201,7 +200,7 @@ internal struct StoreContext {
         let override = lookupOverride(of: atom)
         let key = AtomKey(atom, overrideScopeKey: override?.scopeKey)
 
-        container.subscriptions.removeValue(forKey: key)
+        container.subscribingKeys.remove(key)
         unsubscribe([key], for: container.key)
     }
 
@@ -380,10 +379,10 @@ private extension StoreContext {
         }
     }
 
-    func unsubscribe(_ keys: [AtomKey], for subscriptionKey: SubscriptionKey) {
+    func unsubscribe<Keys: Sequence<AtomKey>>(_ keys: Keys, for subscriptionKey: SubscriptionKey) {
         let store = getStore()
 
-        for key in keys {
+        for key in ContiguousArray(keys) {
             store.state.subscriptions[key]?.removeValue(forKey: subscriptionKey)
             checkRelease(for: key)
         }

diff --git a/Sources/Atoms/Core/SubscriptionContainer.swift b/Sources/Atoms/Core/SubscriptionContainer.swift
@@ -1,14 +1,14 @@
 @usableFromInline
 @MainActor
 internal final class SubscriptionContainer {
-    private var subscriptions = [AtomKey: Subscription]()
-    private var unsubscribe: (([AtomKey]) -> Void)?
+    private var subscribingKeys = Set<AtomKey>()
+    private var unsubscribe: ((Set<AtomKey>) -> Void)?
     private let token = SubscriptionKey.Token()
 
     nonisolated init() {}
 
     deinit {
-        unsubscribe?(Array(subscriptions.keys))
+        unsubscribe?(subscribingKeys)
     }
 
     func wrapper(location: SourceLocation) -> Wrapper {
@@ -25,12 +25,12 @@ internal extension SubscriptionContainer {
         let key: SubscriptionKey
         let location: SourceLocation
 
-        var subscriptions: [AtomKey: Subscription] {
-            get { container?.subscriptions ?? [:] }
-            nonmutating set { container?.subscriptions = newValue }
+        var subscribingKeys: Set<AtomKey> {
+            get { container?.subscribingKeys ?? [] }
+            nonmutating set { container?.subscribingKeys = newValue }
         }
 
-        var unsubscribe: (([AtomKey]) -> Void)? {
+        var unsubscribe: ((Set<AtomKey>) -> Void)? {
             get { container?.unsubscribe }
             nonmutating set { container?.unsubscribe = newValue }
         }

diff --git a/Tests/AtomsTests/Core/StoreContextTests.swift b/Tests/AtomsTests/Core/StoreContextTests.swift
@@ -270,7 +270,7 @@ final class StoreContextTests: XCTestCase {
         }
 
         XCTAssertEqual(initialValue, 0)
-        XCTAssertNotNil(container!.wrapper.subscriptions[key])
+        XCTAssertTrue(container!.wrapper.subscribingKeys.contains(key))
         XCTAssertNotNil(store.state.subscriptions[key]?[container!.wrapper.key])
         XCTAssertEqual((store.state.caches[key] as? AtomCache<TestAtom>)?.value, 0)
         XCTAssertEqual((store.state.caches[dependencyKey] as? AtomCache<DependencyAtom>)?.value, 0)