diff --git a/Sources/Atoms/Core/StoreContext.swift b/Sources/Atoms/Core/StoreContext.swift index 227b5e6b..1797e500 100644 --- 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: 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 index 194b983f..d143f941 100644 --- 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() + private var unsubscribe: ((Set) -> 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 { + get { container?.subscribingKeys ?? [] } + nonmutating set { container?.subscribingKeys = newValue } } - var unsubscribe: (([AtomKey]) -> Void)? { + var unsubscribe: ((Set) -> Void)? { get { container?.unsubscribe } nonmutating set { container?.unsubscribe = newValue } } diff --git a/Tests/AtomsTests/Core/StoreContextTests.swift b/Tests/AtomsTests/Core/StoreContextTests.swift index 1f6a118a..d07b879c 100644 --- 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)?.value, 0) XCTAssertEqual((store.state.caches[dependencyKey] as? AtomCache)?.value, 0)