Reduce process redundancy

Sources/Atoms/Core/StoreContext.swift

@@ -61,7 +61,11 @@ internal struct StoreContext {
         else {
             let cache = makeNewCache(of: atom, for: key, override: override)
             notifyUpdateToObservers()
-            checkRelease(for: key)
+
+            if checkRelease(for: key) {
+                notifyUpdateToObservers()
+            }
+
             return cache.value
         }
     }
@@ -73,7 +77,6 @@ internal struct StoreContext {
 
         if let cache = lookupCache(of: atom, for: key) {
             update(atom: atom, for: key, value: value, cache: cache, order: .newValue)
-            checkRelease(for: key)
         }
     }
 
@@ -86,7 +89,6 @@ internal struct StoreContext {
             var value = cache.value
             body(&value)
             update(atom: atom, for: key, value: value, cache: cache, order: .newValue)
-            checkRelease(for: key)
         }
     }
 
@@ -129,17 +131,14 @@ internal struct StoreContext {
             location: container.location,
             requiresObjectUpdate: requiresObjectUpdate,
             notifyUpdate: notifyUpdate
-        ) {
-            let store = getStore()
-            // Unsubscribe and release if it's no longer used.
-            store.state.subscriptions[key]?.removeValue(forKey: container.key)
-            notifyUpdateToObservers()
-            checkRelease(for: key)
-        }
+        )
         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
+        container.unsubscribe = { keys in
+            unsubscribe(keys, for: container.key)
+        }
 
         if isInserted || cache == nil {
             notifyUpdateToObservers()
@@ -165,8 +164,12 @@ internal struct StoreContext {
         if let cache = lookupCache(of: atom, for: key) {
             update(atom: atom, for: key, value: value, cache: cache, order: .newValue)
         }
+        else {
+            // Release the temporarily created state.
+            // Do not notify update to observers here because refresh doesn't create a new cache.
+            release(for: key)
+        }
 
-        checkRelease(for: key)
         return value
     }
 
@@ -178,15 +181,16 @@ internal struct StoreContext {
         if let cache = lookupCache(of: atom, for: key) {
             let newCache = makeNewCache(of: atom, for: key, override: override)
             update(atom: atom, for: key, value: newCache.value, cache: cache, order: .newValue)
-            checkRelease(for: key)
         }
     }
 
     @usableFromInline
     func unwatch(_ atom: some Atom, container: SubscriptionContainer.Wrapper) {
         let override = lookupOverride(of: atom)
         let key = AtomKey(atom, overrideScopeKey: override?.scopeKey)
-        container.subscriptions.removeValue(forKey: key)?.unsubscribe()
+
+        container.subscriptions.removeValue(forKey: key)
+        unsubscribe([key], for: container.key)
     }
 
     @usableFromInline
@@ -223,7 +227,7 @@ internal struct StoreContext {
 
             // Release dependencies that are no longer dependent.
             if let dependencies = obsoletedDependencies[key] {
-                checkReleaseDependencies(dependencies, for: key)
+                detach(key, fromDependencies: dependencies)
             }
 
             // Notify updates only for the subscriptions of restored atoms.
@@ -233,6 +237,8 @@ internal struct StoreContext {
                 }
             }
         }
+
+        notifyUpdateToObservers()
     }
 }
 
@@ -267,8 +273,10 @@ private extension StoreContext {
             let dependencies = store.graph.dependencies[key] ?? []
             let obsoletedDependencies = oldDependencies.subtracting(dependencies)
 
-            // Check if the dependencies that are no longer used and release them if possible.
-            checkReleaseDependencies(obsoletedDependencies, for: key)
+            if !obsoletedDependencies.isEmpty {
+                detach(key, fromDependencies: obsoletedDependencies)
+                notifyUpdateToObservers()
+            }
         }
 
         // Register the transaction state so it can be terminated from anywhere.
@@ -358,6 +366,17 @@ private extension StoreContext {
         }
     }
 
+    func unsubscribe(_ keys: [AtomKey], for subscriptionKey: SubscriptionKey) {
+        let store = getStore()
+
+        for key in keys {
+            store.state.subscriptions[key]?.removeValue(forKey: subscriptionKey)
+            checkRelease(for: key)
+        }
+
+        notifyUpdateToObservers()
+    }
+
     func invalidate(for key: AtomKey) -> Set<AtomKey> {
         let store = getStore()
 
@@ -378,14 +397,11 @@ private extension StoreContext {
         store.state.states.removeValue(forKey: key)
         store.state.subscriptions.removeValue(forKey: key)
 
-        // Check if the dependencies are releasable.
-        checkReleaseDependencies(dependencies, for: key)
-
-        // Notify release.
-        notifyUpdateToObservers()
+        detach(key, fromDependencies: dependencies)
     }
 
-    func checkRelease(for key: AtomKey) {
+    @discardableResult
+    func checkRelease(for key: AtomKey) -> Bool {
         let store = getStore()
 
         // The condition under which an atom may be released are as follows:
@@ -398,16 +414,16 @@ private extension StoreContext {
         lazy var shouldRelease = !shouldKeepAlive && isChildrenEmpty && isSubscriptionEmpty
 
         guard shouldRelease else {
-            return
+            return false
         }
 
         release(for: key)
+        return true
     }
 
-    func checkReleaseDependencies(_ dependencies: Set<AtomKey>, for key: AtomKey) {
+    func detach(_ key: AtomKey, fromDependencies dependencies: Set<AtomKey>) {
         let store = getStore()
 
-        // Recursively release dependencies while unlinking the dependent.
         for dependency in dependencies {
             store.graph.children[dependency]?.remove(key)
             checkRelease(for: dependency)
@@ -480,6 +496,7 @@ private extension StoreContext {
 
             // Release the invalid registration as a fallback.
             release(for: key)
+            notifyUpdateToObservers()
             return makeState()
         }
 
@@ -509,6 +526,7 @@ private extension StoreContext {
 
             // Release the invalid registration as a fallback.
             release(for: key)
+            notifyUpdateToObservers()
             return nil
         }
 

Sources/Atoms/Core/Subscription.swift

@@ -3,5 +3,4 @@ internal struct Subscription {
     let location: SourceLocation
     let requiresObjectUpdate: Bool
     let notifyUpdate: () -> Void
-    let unsubscribe: () -> Void
 }

Sources/Atoms/Core/SubscriptionContainer.swift

@@ -2,14 +2,13 @@
 @MainActor
 internal final class SubscriptionContainer {
     private var subscriptions = [AtomKey: Subscription]()
+    private var unsubscribe: (([AtomKey]) -> Void)?
     private let token = SubscriptionKey.Token()
 
     nonisolated init() {}
 
     deinit {
-        for subscription in ContiguousArray(subscriptions.values) {
-            subscription.unsubscribe()
-        }
+        unsubscribe?(Array(subscriptions.keys))
     }
 
     func wrapper(location: SourceLocation) -> Wrapper {
@@ -31,6 +30,11 @@ internal extension SubscriptionContainer {
             nonmutating set { container?.subscriptions = newValue }
         }
 
+        var unsubscribe: (([AtomKey]) -> Void)? {
+            get { container?.unsubscribe }
+            nonmutating set { container?.unsubscribe = newValue }
+        }
+
         init(
             _ container: SubscriptionContainer,
             token: SubscriptionKey.Token,