fix: Lifecycle and unsubscription problem

Sources/Atoms/AtomRoot.swift

@@ -42,7 +42,7 @@ import SwiftUI
 public struct AtomRoot<Content: View>: View {
     @StateObject
     private var state: State
-    private var overrides: Overrides
+    private var overrides = Overrides()
     private var observers = [AtomObserver]()
     private let content: Content
 
@@ -51,7 +51,6 @@ public struct AtomRoot<Content: View>: View {
     /// - Parameter content: The content that uses atoms.
     public init(@ViewBuilder content: () -> Content) {
         self._state = StateObject(wrappedValue: State())
-        self.overrides = Overrides()
         self.content = content()
     }
 

Sources/Atoms/Context/AtomTestContext.swift

@@ -275,18 +275,13 @@ private extension AtomTestContext {
     @MainActor
     final class State {
         private let _store = Store()
+        private let _container = SubscriptionContainer()
 
-        let container: SubscriptionContainer
-        var overrides: Overrides
+        var overrides = Overrides()
         var observers = [AtomObserver]()
         let notifier = PassthroughSubject<Void, Never>()
         var onUpdate: (() -> Void)?
 
-        init() {
-            overrides = Overrides()
-            container = SubscriptionContainer()
-        }
-
         var store: StoreContext {
             StoreContext(
                 _store,
@@ -295,6 +290,10 @@ private extension AtomTestContext {
             )
         }
 
+        var container: SubscriptionContainer.Wrapper {
+            _container.wrapper
+        }
+
         func notifyUpdate() {
             onUpdate?()
             notifier.send()

Sources/Atoms/Context/AtomViewContext.swift

@@ -7,13 +7,13 @@ public struct AtomViewContext: AtomWatchableContext {
     @usableFromInline
     internal let _store: StoreContext
     @usableFromInline
-    internal let _container: SubscriptionContainer
+    internal let _container: SubscriptionContainer.Wrapper
     @usableFromInline
     internal let _notifyUpdate: () -> Void
 
     internal init(
         store: StoreContext,
-        container: SubscriptionContainer,
+        container: SubscriptionContainer.Wrapper,
         notifyUpdate: @escaping () -> Void
     ) {
         _store = store

Sources/Atoms/Core/Overrides.swift

@@ -3,6 +3,8 @@ internal struct Overrides {
     private var _entriesForNode = [AtomKey: Override]()
     private var _entriesForType = [AtomTypeKey: Override]()
 
+    nonisolated init() {}
+
     mutating func insert<Node: Atom>(
         _ atom: Node,
         with value: @escaping (Node) -> Node.Loader.Value

Sources/Atoms/Core/StoreContext.swift

@@ -19,7 +19,13 @@ internal struct StoreContext {
 
     @usableFromInline
     func read<Node: Atom>(_ atom: Node) -> Node.Loader.Value {
-        getValue(for: atom)
+        let key = AtomKey(atom)
+
+        // Register if it doesn't exist yet because the atom needs to be maintained if it's marked as `KeepAlive`.
+        registerIfAbsent(atom: atom)
+        defer { checkRelease(for: key) }
+
+        return getValue(for: atom)
     }
 
     @usableFromInline
@@ -53,7 +59,7 @@ internal struct StoreContext {
     func watch<Node: Atom>(_ atom: Node, in transaction: Transaction) -> Node.Loader.Value {
         // Return a new value immediately if the transaction is already terminated.
         guard !transaction.isTerminated else {
-            return getNewValue(for: atom)
+            return read(atom)
         }
 
         let store = getStore()
@@ -72,19 +78,18 @@ internal struct StoreContext {
     @usableFromInline
     func watch<Node: Atom>(
         _ atom: Node,
-        container: SubscriptionContainer,
+        container: SubscriptionContainer.Wrapper,
         notifyUpdate: @escaping () -> Void
     ) -> Node.Loader.Value {
         let store = getStore()
         let key = AtomKey(atom)
-        let subscriptionKey = SubscriptionKey(container)
         let subscription = Subscription(notifyUpdate: notifyUpdate) { [weak store] in
             guard let store = store else {
                 return
             }
 
             // Unsubscribe and release if it's no longer used.
-            store.state.subscriptions[key]?.removeValue(forKey: subscriptionKey)
+            store.state.subscriptions[key]?.removeValue(forKey: container.key)
             checkRelease(for: key)
         }
 
@@ -93,13 +98,19 @@ internal struct StoreContext {
 
         // Register the subscription to both the store and the container.
         container.subscriptions[key] = subscription
-        store.state.subscriptions[key, default: [:]].updateValue(subscription, forKey: subscriptionKey)
+        store.state.subscriptions[key, default: [:]].updateValue(subscription, forKey: container.key)
 
         return getValue(for: atom)
     }
 
     @usableFromInline
     func refresh<Node: Atom>(_ atom: Node) async -> Node.Loader.Value where Node.Loader: RefreshableAtomLoader {
+        let key = AtomKey(atom)
+
+        // Register if it doesn't exist yet because the atom needs to be maintained if it's marked as `KeepAlive`.
+        registerIfAbsent(atom: atom)
+        defer { checkRelease(for: key) }
+
         let context = prepareTransaction(for: atom)
         let value: Node.Loader.Value
 
@@ -117,6 +128,12 @@ internal struct StoreContext {
 
     @usableFromInline
     func reset<Node: Atom>(_ atom: Node) {
+        let key = AtomKey(atom)
+
+        // Register if it doesn't exist yet because the atom needs to be maintained if it's marked as `KeepAlive`.
+        registerIfAbsent(atom: atom)
+        defer { checkRelease(for: key) }
+
         let value = getNewValue(for: atom)
         update(atom: atom, with: value)
     }

Sources/Atoms/Core/SubscriptionContainer.swift

@@ -1,11 +1,37 @@
 @usableFromInline
 @MainActor
 internal final class SubscriptionContainer {
-    var subscriptions = [AtomKey: Subscription]()
+    private var subscriptions = [AtomKey: Subscription]()
+
+    var wrapper: Wrapper {
+        Wrapper(self)
+    }
+
+    nonisolated init() {}
 
     deinit {
         for subscription in ContiguousArray(subscriptions.values) {
             subscription.unsubscribe()
         }
     }
 }
+
+internal extension SubscriptionContainer {
+    @usableFromInline
+    @MainActor
+    struct Wrapper {
+        private weak var container: SubscriptionContainer?
+
+        let key: SubscriptionKey
+
+        var subscriptions: [AtomKey: Subscription] {
+            get { container?.subscriptions ?? [:] }
+            nonmutating set { container?.subscriptions = newValue }
+        }
+
+        init(_ container: SubscriptionContainer) {
+            self.container = container
+            self.key = SubscriptionKey(container)
+        }
+    }
+}

Sources/Atoms/PropertyWrapper/ViewContext.swift

@@ -52,7 +52,7 @@ public struct ViewContext: DynamicProperty {
     public var wrappedValue: AtomViewContext {
         AtomViewContext(
             store: _store,
-            container: state.container,
+            container: state.container.wrapper,
             notifyUpdate: state.objectWillChange.send
         )
     }
@@ -61,10 +61,6 @@ public struct ViewContext: DynamicProperty {
 private extension ViewContext {
     @MainActor
     final class State: ObservableObject {
-        let container: SubscriptionContainer
-
-        init() {
-            container = SubscriptionContainer()
-        }
+        let container = SubscriptionContainer()
     }
 }

Tests/AtomsTests/Context/AtomTestContextTests.swift

@@ -92,15 +92,15 @@ final class AtomTestContextTests: XCTestCase {
 
         for observer in observers {
             XCTAssertEqual(observer.assignedAtomKeys, [key])
-            XCTAssertEqual(observer.unassignedAtomKeys, [])
+            XCTAssertTrue(observer.unassignedAtomKeys.isEmpty)
             XCTAssertEqual(observer.changedAtomKeys, [key])
         }
 
         context[atom] = 200
 
         for observer in observers {
             XCTAssertEqual(observer.assignedAtomKeys, [key])
-            XCTAssertEqual(observer.unassignedAtomKeys, [])
+            XCTAssertTrue(observer.unassignedAtomKeys.isEmpty)
             XCTAssertEqual(observer.changedAtomKeys, [key, key])
         }
 

Tests/AtomsTests/Context/AtomViewContextTests.swift

@@ -10,7 +10,7 @@ final class AtomViewContextTests: XCTestCase {
         let container = SubscriptionContainer()
         let context = AtomViewContext(
             store: StoreContext(store),
-            container: container,
+            container: container.wrapper,
             notifyUpdate: {}
         )
 
@@ -23,7 +23,7 @@ final class AtomViewContextTests: XCTestCase {
         let container = SubscriptionContainer()
         let context = AtomViewContext(
             store: StoreContext(store),
-            container: container,
+            container: container.wrapper,
             notifyUpdate: {}
         )
 
@@ -40,7 +40,7 @@ final class AtomViewContextTests: XCTestCase {
         let container = SubscriptionContainer()
         let context = AtomViewContext(
             store: StoreContext(store),
-            container: container,
+            container: container.wrapper,
             notifyUpdate: {}
         )
 
@@ -57,7 +57,7 @@ final class AtomViewContextTests: XCTestCase {
         let container = SubscriptionContainer()
         let context = AtomViewContext(
             store: StoreContext(store),
-            container: container,
+            container: container.wrapper,
             notifyUpdate: {}
         )
 
@@ -78,7 +78,7 @@ final class AtomViewContextTests: XCTestCase {
         let container = SubscriptionContainer()
         let context = AtomViewContext(
             store: StoreContext(store),
-            container: container,
+            container: container.wrapper,
             notifyUpdate: {}
         )
 
@@ -88,4 +88,22 @@ final class AtomViewContextTests: XCTestCase {
 
         XCTAssertEqual(context.watch(atom), 200)
     }
+
+    func testUnsubscription() {
+        let atom = TestValueAtom(value: 100)
+        let key = AtomKey(atom)
+        let store = Store()
+        var container: SubscriptionContainer? = SubscriptionContainer()
+        let context = AtomViewContext(
+            store: StoreContext(store),
+            container: container!.wrapper,
+            notifyUpdate: {}
+        )
+
+        context.watch(atom)
+        XCTAssertNotNil(store.state.atomStates[key])
+
+        container = nil
+        XCTAssertNil(store.state.atomStates[key])
+    }
 }

Tests/AtomsTests/Core/StoreContextTests.swift

@@ -12,9 +12,9 @@ final class StoreContextTests: XCTestCase {
 
         XCTAssertEqual(context.read(atom), 0)
         XCTAssertNil(store.state.atomStates[AtomKey(atom)])
-        XCTAssertTrue(observer.assignedAtomKeys.isEmpty)
+        XCTAssertEqual(observer.assignedAtomKeys, [AtomKey(atom)])
         XCTAssertEqual(observer.changedAtomKeys, [AtomKey(atom)])
-        XCTAssertTrue(observer.unassignedAtomKeys.isEmpty)
+        XCTAssertEqual(observer.unassignedAtomKeys, [AtomKey(atom)])
 
         store.state.atomStates[AtomKey(atom)] = ConcreteAtomState(atom: atom, value: 1)
 
@@ -95,9 +95,9 @@ final class StoreContextTests: XCTestCase {
         XCTAssertEqual(store.graph.dependencies, [key: [dependency0Key]])
         XCTAssertEqual(store.graph.children, [dependency0Key: [key]])
         XCTAssertNil(store.state.atomStates[dependency1Key])
-        XCTAssertEqual(observer.assignedAtomKeys, [dependency0Key])
-        XCTAssertEqual(observer.changedAtomKeys, [dependency0Key])
-        XCTAssertTrue(observer.unassignedAtomKeys.isEmpty)
+        XCTAssertEqual(observer.assignedAtomKeys, [dependency0Key, dependency1Key])
+        XCTAssertEqual(observer.changedAtomKeys, [dependency0Key, dependency1Key])
+        XCTAssertEqual(observer.unassignedAtomKeys, [dependency1Key])
     }
 
     func testWatchFromView() {
@@ -123,12 +123,12 @@ final class StoreContextTests: XCTestCase {
         let key = AtomKey(atom)
         let dependencyKey = AtomKey(dependency)
         var updateCount = 0
-        let initialValue = context.watch(atom, container: container) {
+        let initialValue = context.watch(atom, container: container.wrapper) {
             updateCount += 1
         }
 
         XCTAssertEqual(initialValue, 0)
-        XCTAssertNotNil(container.subscriptions[key])
+        XCTAssertNotNil(container.wrapper.subscriptions[key])
         XCTAssertNotNil(store.state.subscriptions[key]?[subscriptionKey])
         XCTAssertEqual((store.state.atomStates[key] as? ConcreteAtomState<TestAtom>)?.value, 0)
         XCTAssertEqual((store.state.atomStates[dependencyKey] as? ConcreteAtomState<DependencyAtom>)?.value, 0)
@@ -160,13 +160,13 @@ final class StoreContextTests: XCTestCase {
 
         XCTAssertEqual(value0, 0)
         XCTAssertNil(store.state.atomStates[key])
-        XCTAssertTrue(observer.assignedAtomKeys.isEmpty)
+        XCTAssertEqual(observer.assignedAtomKeys, [key])
         XCTAssertEqual(observer.changedAtomKeys, [key])
-        XCTAssertTrue(observer.unassignedAtomKeys.isEmpty)
+        XCTAssertEqual(observer.unassignedAtomKeys, [key])
 
         var updateCount = 0
 
-        _ = context.watch(atom, container: container) {
+        _ = context.watch(atom, container: container.wrapper) {
             updateCount += 1
         }
 
@@ -186,7 +186,7 @@ final class StoreContextTests: XCTestCase {
         let atom = TestStateAtom(defaultValue: 0)
         let key = AtomKey(atom)
         var updateCount = 0
-        let initialValue = context.watch(atom, container: container) {
+        let initialValue = context.watch(atom, container: container.wrapper) {
             updateCount += 1
         }
 
@@ -218,7 +218,7 @@ final class StoreContextTests: XCTestCase {
         let context = StoreContext(store, observers: [observer0])
         let relayedContext = context.relay(observers: [observer1])
 
-        _ = relayedContext.watch(atom, container: container) {}
+        _ = relayedContext.watch(atom, container: container.wrapper) {}
 
         XCTAssertFalse(observer0.assignedAtomKeys.isEmpty)
         XCTAssertFalse(observer1.assignedAtomKeys.isEmpty)
@@ -324,7 +324,7 @@ final class StoreContextTests: XCTestCase {
         let phase = PhaseAtom()
 
         func watch() async -> Int {
-            await atomStore.watch(atom, container: container, notifyUpdate: {}).value
+            await atomStore.watch(atom, container: container.wrapper, notifyUpdate: {}).value
         }
 
         do {

Tests/AtomsTests/Core/SubscriptionContainerTests.swift

@@ -13,7 +13,7 @@ final class SubscriptionContainerTests: XCTestCase {
         let atom0 = TestValueAtom(value: 0)
         let atom1 = TestValueAtom(value: 1)
 
-        container?.subscriptions = [
+        container?.wrapper.subscriptions = [
             AtomKey(atom0): subscription,
             AtomKey(atom1): subscription,
         ]