Flatten scoped stores (#65)

* Cache overridden atom states into a single store instead of hierarchical multi stores

* Add reasonably unique suffix to overridden atoms for graphDesdription

* Release atoms that marked as KeepAlive when it is no longer watched

* Do not access StateObject that is not installed on a view

* Drop support to lookup overridden atoms in Snapshot for consistency

* Lint

* Refactoring

Author: ra1028

Sources/Atoms/AtomRoot.swift

@@ -60,14 +60,26 @@ public struct AtomRoot<Content: View>: View {
     public var body: some View {
         content.environment(
             \.store,
-            StoreContext(
+            .scoped(
                 state.store,
-                overrides: overrides,
-                observers: observers
+                scopeKey: ScopeKey(token: state.token),
+                observers: observers,
+                overrides: overrides
             )
         )
     }
 
+    /// For debugging purposes, each time there is a change in the internal state,
+    /// a snapshot is taken that captures the state of the atoms and their dependency graph
+    /// at that point in time.
+    ///
+    /// - Parameter onUpdate: A closure to handle a snapshot of recent updates.
+    ///
+    /// - Returns: The self instance.
+    public func observe(_ onUpdate: @escaping @MainActor (Snapshot) -> Void) -> Self {
+        mutating { $0.observers.append(Observer(onUpdate: onUpdate)) }
+    }
+
     /// Overrides the atom value with the given value.
     ///
     /// When accessing the overridden atom, this context will create and return the given value
@@ -97,23 +109,13 @@ public struct AtomRoot<Content: View>: View {
     public func override<Node: Atom>(_ atomType: Node.Type, with value: @escaping (Node) -> Node.Loader.Value) -> Self {
         mutating { $0.overrides[OverrideKey(atomType)] = AtomOverride(value: value) }
     }
-
-    /// For debugging purposes, each time there is a change in the internal state,
-    /// a snapshot is taken that captures the state of the atoms and their dependency graph
-    /// at that point in time.
-    ///
-    /// - Parameter onUpdate: A closure to handle a snapshot of recent updates.
-    ///
-    /// - Returns: The self instance.
-    public func observe(_ onUpdate: @escaping @MainActor (Snapshot) -> Void) -> Self {
-        mutating { $0.observers.append(Observer(onUpdate: onUpdate)) }
-    }
 }
 
 private extension AtomRoot {
     @MainActor
     final class State: ObservableObject {
         let store = AtomStore()
+        let token = ScopeKey.Token()
     }
 
     func `mutating`(_ mutation: (inout Self) -> Void) -> Self {

Sources/Atoms/AtomScope.swift

@@ -101,13 +101,27 @@ public struct AtomScope<Content: View>: View {
         content.environment(
             \.store,
             (store ?? environmentStore).scoped(
-                store: state.store,
-                overrides: overrides,
-                observers: observers
+                key: ScopeKey(token: state.token),
+                observers: observers,
+                overrides: overrides
             )
         )
     }
 
+    /// For debugging purposes, each time there is a change in the internal state,
+    /// a snapshot is taken that captures the state of the atoms and their dependency graph
+    /// at that point in time.
+    ///
+    /// Note that unlike observed by ``AtomRoot``, this is triggered only by internal state changes
+    /// caused by atoms use in this scope.
+    ///
+    /// - Parameter onUpdate: A closure to handle a snapshot of recent updates.
+    ///
+    /// - Returns: The self instance.
+    public func observe(_ onUpdate: @escaping @MainActor (Snapshot) -> Void) -> Self {
+        mutating { $0.observers.append(Observer(onUpdate: onUpdate)) }
+    }
+
     /// Override the atom value used in this scope with the given value.
     ///
     /// When accessing the overridden atom, this context will create and return the given value
@@ -141,26 +155,12 @@ public struct AtomScope<Content: View>: View {
     public func override<Node: Atom>(_ atomType: Node.Type, with value: @escaping (Node) -> Node.Loader.Value) -> Self {
         mutating { $0.overrides[OverrideKey(atomType)] = AtomOverride(value: value) }
     }
-
-    /// For debugging purposes, each time there is a change in the internal state,
-    /// a snapshot is taken that captures the state of the atoms and their dependency graph
-    /// at that point in time.
-    ///
-    /// Note that unlike observed by ``AtomRoot``, this is triggered only by internal state changes
-    /// caused by atoms use in this scope.
-    ///
-    /// - Parameter onUpdate: A closure to handle a snapshot of recent updates.
-    ///
-    /// - Returns: The self instance.
-    public func observe(_ onUpdate: @escaping @MainActor (Snapshot) -> Void) -> Self {
-        mutating { $0.observers.append(Observer(onUpdate: onUpdate)) }
-    }
 }
 
 private extension AtomScope {
     @MainActor
     final class State: ObservableObject {
-        let store = AtomStore()
+        let token = ScopeKey.Token()
     }
 
     func `mutating`(_ mutation: (inout Self) -> Void) -> Self {

Sources/Atoms/Context/AtomTestContext.swift

@@ -250,9 +250,8 @@ public struct AtomTestContext: AtomWatchableContext {
     /// It simulates cases where other atoms or views no longer watches to the atom.
     ///
     /// - Parameter atom: An atom that associates the value.
-    public func unwatch<Node: Atom>(_ atom: Node) {
-        let key = AtomKey(atom)
-        container.subscriptions.removeValue(forKey: key)?.unsubscribe()
+    public func unwatch(_ atom: some Atom) {
+        store.unwatch(atom, container: container)
     }
 
     /// Overrides the atom value with the given value.
@@ -285,6 +284,7 @@ public struct AtomTestContext: AtomWatchableContext {
 private extension AtomTestContext {
     final class State {
         let store = AtomStore()
+        let token = ScopeKey.Token()
         let container = SubscriptionContainer()
         let notifier = PassthroughSubject<Void, Never>()
         var overrides = [OverrideKey: any AtomOverrideProtocol]()
@@ -297,7 +297,12 @@ private extension AtomTestContext {
     }
 
     var store: StoreContext {
-        StoreContext(state.store, overrides: state.overrides)
+        .scoped(
+            state.store,
+            scopeKey: ScopeKey(token: state.token),
+            observers: [],
+            overrides: state.overrides
+        )
     }
 
     var container: SubscriptionContainer.Wrapper {

Sources/Atoms/Core/AtomKey.swift

@@ -1,24 +1,36 @@
-internal struct AtomKey: Hashable {
+internal struct AtomKey: Hashable, CustomStringConvertible {
     private let key: AnyHashable
     private let type: ObjectIdentifier
+    private let overrideScopeKey: ScopeKey?
     private let getName: () -> String
 
-    var name: String {
-        getName()
+    var isOverridden: Bool {
+        overrideScopeKey != nil
     }
 
-    init<Node: Atom>(_ atom: Node) {
-        key = AnyHashable(atom.key)
-        type = ObjectIdentifier(Node.self)
-        getName = { String(describing: Node.self) }
+    var description: String {
+        if let overrideScopeKey {
+            return getName() + "-override:\(overrideScopeKey.id)"
+        }
+        else {
+            return getName()
+        }
+    }
+
+    init<Node: Atom>(_ atom: Node, overrideScopeKey: ScopeKey?) {
+        self.key = AnyHashable(atom.key)
+        self.type = ObjectIdentifier(Node.self)
+        self.overrideScopeKey = overrideScopeKey
+        self.getName = { String(describing: Node.self) }
     }
 
     func hash(into hasher: inout Hasher) {
         hasher.combine(key)
         hasher.combine(type)
+        hasher.combine(overrideScopeKey)
     }
 
     static func == (lhs: Self, rhs: Self) -> Bool {
-        lhs.key == rhs.key && lhs.type == rhs.type
+        lhs.key == rhs.key && lhs.type == rhs.type && lhs.overrideScopeKey == rhs.overrideScopeKey
     }
 }

Sources/Atoms/Core/AtomOverride.swift

@@ -1,10 +1,30 @@
-@MainActor
 internal protocol AtomOverrideProtocol {
     associatedtype Node: Atom
 
     var value: (Node) -> Node.Loader.Value { get }
+
+    func scoped(key: ScopeKey) -> any AtomScopedOverrideProtocol
 }
 
 internal struct AtomOverride<Node: Atom>: AtomOverrideProtocol {
     let value: (Node) -> Node.Loader.Value
+
+    func scoped(key: ScopeKey) -> any AtomScopedOverrideProtocol {
+        AtomScopedOverride<Node>(scopeKey: key, value: value)
+    }
+}
+
+// As a workaround to the problem of not getting ScopeKey synchronously
+// when the AtomRoot or AtomScope's override modifier is called, those modifiers
+// temporarily register AtomOverride and convert them to AtomScopedOverride when
+// their View body is evaluated. This is not ideal from a performance standpoint,
+// so it will be improved as soon as an alternative way to grant per-scope keys
+// independent of the SwiftUI lifecycle is came up.
+internal protocol AtomScopedOverrideProtocol {
+    var scopeKey: ScopeKey { get }
+}
+
+internal struct AtomScopedOverride<Node: Atom>: AtomScopedOverrideProtocol {
+    let scopeKey: ScopeKey
+    let value: (Node) -> Node.Loader.Value
 }

Sources/Atoms/Core/ScopeKey.swift

@@ -0,0 +1,13 @@
+internal struct ScopeKey: Hashable {
+    final class Token {}
+
+    private let identifier: ObjectIdentifier
+
+    var id: String {
+        String(hashValue, radix: 36, uppercase: false)
+    }
+
+    init(token: Token) {
+        identifier = ObjectIdentifier(token)
+    }
+}