Allow keeping Atoms alive within a scope until it's deallocation (#188)

* Keep an atom alive until its scope is unregistered if it conforms both KeepAlive and Scoped

* Refactoring

* Add test case

* Update documentation

* Refactoring

Author: ra1028

README.md

@@ -849,6 +849,7 @@ VStack {
 #### [KeepAlive](https://ra1028.github.io/swiftui-atom-properties/documentation/atoms/keepalive)
 
 `KeepAlive` allows the atom to preserve its data even if it's no longer watched from anywhere.  
+When used with the [Scoped](#scoped) attribute, the atom's data is preserved until the scope ([AtomScope](#atomscope)) it belongs to is dismantled from the view tree. This is useful when you want to prevent the atom from being released prematurely within a specific scope.  
 
 <details><summary><code>📖 Example</code></summary>
 

Sources/Atoms/AtomScope.swift

@@ -150,6 +150,11 @@ public struct AtomScope<Content: View>: View {
 }
 
 private extension AtomScope {
+    @MainActor
+    final class State: ObservableObject {
+        let scopeState = ScopeState()
+    }
+
     enum Inheritance {
         case environment(scopeID: ScopeID)
         case context(AtomViewContext)
@@ -161,21 +166,23 @@ private extension AtomScope {
         let overrideContainer: OverrideContainer
         let content: Content
 
-        @State
-        private var scopeToken = ScopeKey.Token()
         @Environment(\.store)
-        private var environmentStore
+        private var store
+
+        @StateObject
+        private var state = State()
 
         var body: some View {
-            content.environment(
-                \.store,
-                environmentStore?.scoped(
-                    scopeID: scopeID,
-                    scopeKey: scopeToken.key,
-                    observers: observers,
-                    overrideContainer: overrideContainer
-                )
+            let scopedStore = store?.scoped(
+                scopeID: scopeID,
+                scopeKey: state.scopeState.token.key,
+                observers: observers,
+                overrideContainer: overrideContainer
             )
+
+            scopedStore?.registerScope(state: state.scopeState)
+
+            return content.environment(\.store, scopedStore)
         }
     }
 }

Sources/Atoms/Attribute/KeepAlive.swift

@@ -1,7 +1,11 @@
 /// An attribute protocol to allow the value of an atom to continue being retained
 /// even after they are no longer watched.
 ///
-/// Note that overridden or scoped atoms are not retained even with this attribute.
+/// ## Note
+///
+/// Atoms that conform to this attribute and are either scoped using the ``Scoped`` attribute
+/// or overridden via ``AtomScope/scopedOverride(_:with:)-5jen3`` are retained until their scope
+/// is dismantled from the view tree, after which they are released.
 ///
 /// ## Example
 ///

Sources/Atoms/Core/AtomCache.swift

@@ -3,21 +3,26 @@ internal protocol AtomCacheProtocol {
 
     var atom: Node { get }
     var value: Node.Produced { get }
-    var initializedScope: Scope? { get }
+    var scopeValues: ScopeValues? { get }
+    var shouldKeepAlive: Bool { get }
 
     func updated(value: Node.Produced) -> Self
 }
 
 internal struct AtomCache<Node: Atom>: AtomCacheProtocol, CustomStringConvertible {
     let atom: Node
     let value: Node.Produced
-    let initializedScope: Scope?
+    let scopeValues: ScopeValues?
 
     var description: String {
         "\(value)"
     }
 
+    var shouldKeepAlive: Bool {
+        atom is any KeepAlive
+    }
+
     func updated(value: Node.Produced) -> Self {
-        AtomCache(atom: atom, value: value, initializedScope: initializedScope)
+        AtomCache(atom: atom, value: value, scopeValues: scopeValues)
     }
 }

Sources/Atoms/Core/AtomKey.swift

@@ -1,9 +1,10 @@
 internal struct AtomKey: Hashable, Sendable, CustomStringConvertible {
     private let key: UnsafeUncheckedSendable<AnyHashable>
     private let type: ObjectIdentifier
-    private let scopeKey: ScopeKey?
     private let anyAtomType: Any.Type
 
+    let scopeKey: ScopeKey?
+
     var description: String {
         let atomLabel = String(describing: anyAtomType)
 
@@ -15,10 +16,6 @@ internal struct AtomKey: Hashable, Sendable, CustomStringConvertible {
         }
     }
 
-    var isScoped: Bool {
-        scopeKey != nil
-    }
-
     init<Node: Atom>(_ atom: Node, scopeKey: ScopeKey?) {
         self.key = UnsafeUncheckedSendable(atom.key)
         self.type = ObjectIdentifier(Node.self)

Sources/Atoms/Core/Graph.swift

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

Sources/Atoms/Core/Scope.swift

@@ -1,6 +1,4 @@
+@MainActor
 internal struct Scope {
-    let key: ScopeKey
-    let observers: [Observer]
-    let overrideContainer: OverrideContainer
-    let ancestorScopeKeys: [ScopeID: ScopeKey]
+    var atoms = Set<AtomKey>()
 }

Sources/Atoms/Core/ScopeKey.swift

@@ -1,5 +1,6 @@
 @usableFromInline
 internal struct ScopeKey: Hashable, Sendable, CustomStringConvertible {
+    @MainActor
     final class Token {
         private(set) lazy var key = ScopeKey(token: self)
     }

Sources/Atoms/Core/ScopeState.swift

@@ -0,0 +1,19 @@
+@usableFromInline
+@MainActor
+internal final class ScopeState {
+    let token = ScopeKey.Token()
+
+    #if !hasFeature(IsolatedDefaultValues)
+        nonisolated init() {}
+    #endif
+
+    nonisolated(unsafe) var unregister: (@MainActor () -> Void)?
+
+    // TODO: Use isolated synchronous deinit once it's available.
+    // 0371-isolated-synchronous-deinit
+    deinit {
+        MainActor.performIsolated { [unregister] in
+            unregister?()
+        }
+    }
+}

Sources/Atoms/Core/ScopeValues.swift

@@ -0,0 +1,6 @@
+internal struct ScopeValues {
+    let key: ScopeKey
+    let observers: [Observer]
+    let overrideContainer: OverrideContainer
+    let ancestorScopeKeys: [ScopeID: ScopeKey]
+}