Bug fix: Atom effect receives context with unintended scope (#179)

* Add test case

* Call atom effect with the scope at which they were initialized

* Write more tests

* Switch context when mutating operation

* Refactoring

* Disable test case for Swift 5 compiler

Author: ra1028

Sources/Atoms/AtomScope.swift

@@ -80,8 +80,8 @@ public struct AtomScope<Content: View>: View {
     /// Observes the state changes with a snapshot that captures the whole atom states and
     /// their dependency graph at the point in time for debugging purposes.
     ///
-    /// Note that unlike ``AtomRoot/observe(_:)``, this observes only the state changes caused by atoms
-    /// used in this scope.
+    /// Note that unlike ``AtomRoot/observe(_:)``, this observes only the state changes of atoms
+    /// initialized in this scope.
     ///
     /// - Note: It ignores the observers if this scope inherits the parent scope.
     ///

Sources/Atoms/Core/StoreContext.swift

@@ -71,7 +71,8 @@ internal struct StoreContext {
         let (key, _) = lookupAtomKeyAndOverride(of: atom)
 
         if let cache = lookupCache(of: atom, for: key) {
-            update(atom: atom, for: key, cache: cache, newValue: value)
+            switchContext(scope: cache.initializedScope)
+                .update(atom: atom, for: key, cache: cache, newValue: value)
         }
     }
 
@@ -81,7 +82,8 @@ internal struct StoreContext {
 
         if let cache = lookupCache(of: atom, for: key) {
             let newValue = mutating(cache.value, body)
-            update(atom: atom, for: key, cache: cache, newValue: newValue)
+            switchContext(scope: cache.initializedScope)
+                .update(atom: atom, for: key, cache: cache, newValue: newValue)
         }
     }
 
@@ -128,7 +130,10 @@ internal struct StoreContext {
     @_disfavoredOverload
     func refresh<Node: AsyncAtom>(_ atom: Node) async -> Node.Produced {
         let (key, override) = lookupAtomKeyAndOverride(of: atom)
-        let context = prepareForTransaction(of: atom, for: key)
+        let cache = lookupCache(of: atom, for: key)
+        let localContext = switchContext(scope: cache?.initializedScope)
+        let context = localContext.prepareForTransaction(of: atom, for: key)
+
         let value: Node.Produced
 
         if let override {
@@ -140,14 +145,14 @@ internal struct StoreContext {
 
         await atom.refreshProducer.refreshValue(value, context)
 
-        guard let cache = lookupCache(of: atom, for: key) else {
+        guard let cache else {
             checkAndRelease(for: key)
             return value
         }
 
         // Notify update unless it's cancelled or terminated by other operations.
         if !Task.isCancelled && !context.isTerminated {
-            update(atom: atom, for: key, cache: cache, newValue: value)
+            localContext.update(atom: atom, for: key, cache: cache, newValue: value)
         }
 
         return value
@@ -156,8 +161,10 @@ internal struct StoreContext {
     @usableFromInline
     func refresh<Node: Refreshable>(_ atom: Node) async -> Node.Produced {
         let (key, _) = lookupAtomKeyAndOverride(of: atom)
-        let state = getState(of: atom, for: key)
-        let currentContext = AtomCurrentContext(store: self)
+        let cache = lookupCache(of: atom, for: key)
+        let localContext = switchContext(scope: cache?.initializedScope)
+        let state = localContext.getState(of: atom, for: key)
+        let currentContext = AtomCurrentContext(store: localContext)
 
         // Detach the dependencies once to delay updating the downstream until
         // this atom's value refresh is complete.
@@ -167,14 +174,14 @@ internal struct StoreContext {
         // Restore dependencies when the refresh is completed.
         attachDependencies(dependencies, for: key)
 
-        guard let transactionState = state.transactionState, let cache = lookupCache(of: atom, for: key) else {
+        guard let transactionState = state.transactionState, let cache else {
             checkAndRelease(for: key)
             return value
         }
 
         // Notify update unless it's cancelled or terminated by other operations.
         if !Task.isCancelled && !transactionState.isTerminated {
-            update(atom: atom, for: key, cache: cache, newValue: value)
+            localContext.update(atom: atom, for: key, cache: cache, newValue: value)
         }
 
         return value
@@ -186,14 +193,18 @@ internal struct StoreContext {
         let (key, override) = lookupAtomKeyAndOverride(of: atom)
 
         if let cache = lookupCache(of: atom, for: key) {
-            let newValue = getValue(of: atom, for: key, override: override)
-            update(atom: atom, for: key, cache: cache, newValue: newValue)
+            let localContext = switchContext(scope: cache.initializedScope)
+            let newValue = localContext.getValue(of: atom, for: key, override: override)
+            localContext.update(atom: atom, for: key, cache: cache, newValue: newValue)
         }
     }
 
     @usableFromInline
     func reset(_ atom: some Resettable) {
-        let currentContext = AtomCurrentContext(store: self)
+        let (key, _) = lookupAtomKeyAndOverride(of: atom)
+        let cache = lookupCache(of: atom, for: key)
+        let localContext = switchContext(scope: cache?.initializedScope)
+        let currentContext = AtomCurrentContext(store: localContext)
         atom.reset(context: currentContext)
     }
 
@@ -306,22 +317,27 @@ private extension StoreContext {
             updatedScopes[currentScope.key] = currentScope
         }
 
-        func transitiveUpdate(for key: AtomKey, cache: some AtomCacheProtocol) {
+        func updatePropagation(for key: AtomKey, cache: some AtomCacheProtocol) {
             // Dependents must be updated with the scope at which they were initialised.
-            let localContext = StoreContext(
-                store: store,
-                rootScope: rootScope,
-                currentScope: cache.initializedScope
-            )
+            let localContext = switchContext(scope: cache.initializedScope)
 
-            let didUpdate = localContext.transitiveUpdate(for: key, cache: cache)
+            // Overridden atoms don't get updated transitively.
+            let newValue = localContext.getValue(of: cache.atom, for: key, override: nil)
 
-            guard didUpdate else {
+            store.state.caches[key] = cache.updated(value: newValue)
+
+            // Check whether if the dependent atoms should be updated transitively.
+            guard cache.atom.producer.shouldUpdate(cache.value, newValue) else {
                 // Record the atom to avoid downstream from being update.
                 skippedDependencies.insert(key)
                 return
             }
 
+            // Perform side effects before updating downstream.
+            let state = localContext.getState(of: cache.atom, for: key)
+            let currentContext = AtomCurrentContext(store: localContext)
+            state.effect.updated(context: currentContext)
+
             if let scope = cache.initializedScope {
                 updatedScopes[scope.key] = scope
             }
@@ -362,7 +378,7 @@ private extension StoreContext {
 
                 if let cache, let dependencyCache {
                     performUpdate(dependency: dependencyCache.atom) {
-                        transitiveUpdate(for: key, cache: cache)
+                        updatePropagation(for: key, cache: cache)
                     }
                 }
 
@@ -384,30 +400,12 @@ private extension StoreContext {
         notifyUpdateToObservers(scopes: updatedScopes.values)
     }
 
-    func transitiveUpdate(for key: AtomKey, cache: some AtomCacheProtocol) -> Bool {
-        // Overridden atoms don't get updated transitively.
-        let newValue = getValue(of: cache.atom, for: key, override: nil)
-
-        store.state.caches[key] = cache.updated(value: newValue)
-
-        // Check whether if the dependent atoms should be updated transitively.
-        guard cache.atom.producer.shouldUpdate(cache.value, newValue) else {
-            return false
-        }
-
-        // Perform side effects before updating downstream.
-        let state = getState(of: cache.atom, for: key)
-        let currentContext = AtomCurrentContext(store: self)
-        state.effect.updated(context: currentContext)
-        return true
-    }
-
     func release(for key: AtomKey) {
         let dependencies = store.graph.dependencies.removeValue(forKey: key)
         let state = store.state.states.removeValue(forKey: key)
+        let cache = store.state.caches.removeValue(forKey: key)
 
         store.graph.children.removeValue(forKey: key)
-        store.state.caches.removeValue(forKey: key)
         store.state.subscriptions.removeValue(forKey: key)
 
         if let dependencies {
@@ -419,8 +417,12 @@ private extension StoreContext {
 
         state?.transactionState?.terminate()
 
-        let currentContext = AtomCurrentContext(store: self)
-        state?.effect.released(context: currentContext)
+        if let state, let cache {
+            // It must call release effect with the scope at which they were initialised.
+            let localContext = switchContext(scope: cache.initializedScope)
+            let currentContext = AtomCurrentContext(store: localContext)
+            state.effect.released(context: currentContext)
+        }
     }
 
     func checkAndRelease(for key: AtomKey) {
@@ -505,7 +507,8 @@ private extension StoreContext {
 
         return AtomProducerContext(store: self, transactionState: transactionState) { newValue in
             if let cache = lookupCache(of: atom, for: key) {
-                update(atom: atom, for: key, cache: cache, newValue: newValue)
+                switchContext(scope: cache.initializedScope)
+                    .update(atom: atom, for: key, cache: cache, newValue: newValue)
             }
         }
     }
@@ -640,4 +643,12 @@ private extension StoreContext {
             observer.onUpdate(snapshot)
         }
     }
+
+    func switchContext(scope: Scope?) -> StoreContext {
+        StoreContext(
+            store: store,
+            rootScope: rootScope,
+            currentScope: scope
+        )
+    }
 }

Tests/AtomsTests/Core/StoreContextTests.swift

@@ -1199,6 +1199,145 @@ final class StoreContextTests: XCTestCase {
         XCTAssertEqual(effect.releasedCount, 1)
     }
 
+    #if compiler(>=6)
+        @MainActor
+        func testEffectCrossScopeBoundary() async {
+            struct TestDependencyAtom: ValueAtom, Hashable {
+                func value(context: Context) -> Int {
+                    1
+                }
+            }
+
+            struct TestAtom: StateAtom, Refreshable, Resettable {
+                let testEffect: TestEffect
+
+                var key: UniqueKey {
+                    UniqueKey()
+                }
+
+                func defaultValue(context: Context) -> Int {
+                    context.read(TestDependencyAtom())
+                }
+
+                func refresh(context: CurrentContext) async -> Int {
+                    context.read(TestDependencyAtom())
+                }
+
+                func reset(context: CurrentContext) {
+                    context.reset(TestDependencyAtom())
+                }
+
+                func effect(context: CurrentContext) -> some AtomEffect {
+                    testEffect.initContext = context
+                    return testEffect
+                }
+            }
+
+            struct TestAsyncAtom: AsyncPhaseAtom {
+                let testEffect: TestEffect
+
+                var key: UniqueKey {
+                    UniqueKey()
+                }
+
+                func value(context: Context) async -> Int {
+                    context.read(TestDependencyAtom())
+                }
+
+                func effect(context: CurrentContext) -> some AtomEffect {
+                    testEffect.initContext = context
+                    return testEffect
+                }
+            }
+
+            final class TestEffect: AtomEffect {
+                var initContext: AtomCurrentContext?
+                var value1: Int?
+                var value2: Int?
+
+                func updateValues(context: Context) {
+                    value1 = context.read(TestDependencyAtom())
+                    value2 = initContext?.read(TestDependencyAtom())
+                }
+
+                func initialized(context: Context) {
+                    updateValues(context: context)
+                }
+
+                func updated(context: Context) {
+                    updateValues(context: context)
+                }
+
+                func released(context: Context) {
+                    updateValues(context: context)
+                }
+            }
+
+            let subscriberState = SubscriberState()
+            let subscriber = Subscriber(subscriberState)
+            let store = AtomStore()
+            let rootScopeToken = ScopeKey.Token()
+            let scopeToken = ScopeKey.Token()
+            let context = StoreContext.root(store: store, scopeKey: rootScopeToken.key)
+            let scopedContext = context.scoped(
+                scopeID: ScopeID(DefaultScopeID()),
+                scopeKey: scopeToken.key,
+                observers: [],
+                overrideContainer: OverrideContainer()
+                    .addingOverride(for: TestDependencyAtom()) { _ in
+                        2
+                    }
+            )
+            let testEffect = TestEffect()
+            let testAtom = TestAtom(testEffect: testEffect)
+            let testAsyncAtom = TestAsyncAtom(testEffect: testEffect)
+
+            func assert<Node: Atom>(
+                _ atom: Node,
+                _ expected: Node.Produced,
+                file: StaticString = #filePath,
+                line: UInt = #line
+            ) where Node.Produced: Equatable {
+                XCTAssertEqual(testEffect.value1, 1, file: file, line: line)
+                XCTAssertEqual(testEffect.value2, 1, file: file, line: line)
+                XCTAssertEqual(context.read(atom), expected, file: file, line: line)
+            }
+
+            _ = context.read(testAtom)
+            assert(testAtom, 1)
+
+            _ = context.watch(testAtom, subscriber: subscriber, subscription: Subscription())
+            assert(testAtom, 1)
+
+            _ = context.watch(testAsyncAtom, subscriber: subscriber, subscription: Subscription())
+            assert(testAsyncAtom, .suspending)
+
+            context.modify(testAtom) { $0 = 1 }
+            assert(testAtom, 1)
+
+            scopedContext.set(context.read(testAtom), for: testAtom)
+            assert(testAtom, 1)
+
+            scopedContext.modify(testAtom) { $0 = 1 }
+            assert(testAtom, 1)
+
+            _ = await scopedContext.refresh(testAtom)
+            assert(testAtom, 1)
+
+            _ = await scopedContext.refresh(testAsyncAtom)
+            assert(testAsyncAtom, .success(1))
+
+            scopedContext.reset(testAtom)
+            assert(testAtom, 1)
+
+            scopedContext.reset(testAsyncAtom)
+            assert(testAsyncAtom, .suspending)
+
+            scopedContext.unwatch(testAtom, subscriber: subscriber)
+            assert(testAtom, 1)
+        }
+    #endif
+
     @MainActor
     func testUpdateInTopologicalOrder() {
         struct TestAtom1: StateAtom, Hashable {
@@ -1254,7 +1393,7 @@ final class StoreContextTests: XCTestCase {
     }
 
     @MainActor
-    func testTransitiveUpdate() {
+    func testUpdatePropagation() {
         struct TestAtom1: StateAtom, Hashable {
             func defaultValue(context: Context) -> Int {
                 0