Fix bug where scoped overrides not used when calling refresh if cache is missing (#186)

* Refresh now uses a store with the current scope if cache is missing

* Fix compilation in Swift 5

Author: ra1028

Sources/Atoms/Core/StoreContext.swift

@@ -71,7 +71,7 @@ internal struct StoreContext {
         let (key, _) = lookupAtomKeyAndOverride(of: atom)
 
         if let cache = lookupCache(of: atom, for: key) {
-            switchContext(scope: cache.initializedScope)
+            switchContext(with: cache)
                 .update(atom: atom, for: key, cache: cache, newValue: value)
         }
     }
@@ -82,7 +82,7 @@ internal struct StoreContext {
 
         if let cache = lookupCache(of: atom, for: key) {
             let newValue = mutating(cache.value, body)
-            switchContext(scope: cache.initializedScope)
+            switchContext(with: cache)
                 .update(atom: atom, for: key, cache: cache, newValue: newValue)
         }
     }
@@ -131,7 +131,7 @@ internal struct StoreContext {
     func refresh<Node: AsyncAtom>(_ atom: Node) async -> Node.Produced {
         let (key, override) = lookupAtomKeyAndOverride(of: atom)
         let cache = lookupCache(of: atom, for: key)
-        let localContext = switchContext(scope: cache?.initializedScope)
+        let localContext = cache.map(switchContext) ?? self
         let context = localContext.prepareForTransaction(of: atom, for: key)
 
         let value: Node.Produced
@@ -162,7 +162,7 @@ internal struct StoreContext {
     func refresh<Node: Refreshable>(_ atom: Node) async -> Node.Produced {
         let (key, _) = lookupAtomKeyAndOverride(of: atom)
         let cache = lookupCache(of: atom, for: key)
-        let localContext = switchContext(scope: cache?.initializedScope)
+        let localContext = cache.map(switchContext) ?? self
         let state = localContext.getState(of: atom, for: key)
         let currentContext = AtomCurrentContext(store: localContext)
 
@@ -193,7 +193,7 @@ internal struct StoreContext {
         let (key, override) = lookupAtomKeyAndOverride(of: atom)
 
         if let cache = lookupCache(of: atom, for: key) {
-            let localContext = switchContext(scope: cache.initializedScope)
+            let localContext = switchContext(with: cache)
             let newValue = localContext.getValue(of: atom, for: key, override: override)
             localContext.update(atom: atom, for: key, cache: cache, newValue: newValue)
         }
@@ -202,10 +202,12 @@ internal struct StoreContext {
     @usableFromInline
     func reset(_ atom: some Resettable) {
         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)
+
+        if let cache = lookupCache(of: atom, for: key) {
+            let localContext = switchContext(with: cache)
+            let currentContext = AtomCurrentContext(store: localContext)
+            atom.reset(context: currentContext)
+        }
     }
 
     @usableFromInline
@@ -319,7 +321,7 @@ private extension StoreContext {
 
         func updatePropagation(for key: AtomKey, cache: some AtomCacheProtocol) {
             // Dependents must be updated with the scope at which they were initialised.
-            let localContext = switchContext(scope: cache.initializedScope)
+            let localContext = switchContext(with: cache)
 
             // Overridden atoms don't get updated transitively.
             let newValue = localContext.getValue(of: cache.atom, for: key, override: nil)
@@ -419,7 +421,7 @@ private extension StoreContext {
 
         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 localContext = switchContext(with: cache)
             let currentContext = AtomCurrentContext(store: localContext)
             state.effect.released(context: currentContext)
         }
@@ -507,7 +509,7 @@ private extension StoreContext {
 
         return AtomProducerContext(store: self, transactionState: transactionState) { newValue in
             if let cache = lookupCache(of: atom, for: key) {
-                switchContext(scope: cache.initializedScope)
+                switchContext(with: cache)
                     .update(atom: atom, for: key, cache: cache, newValue: newValue)
             }
         }
@@ -644,11 +646,11 @@ private extension StoreContext {
         }
     }
 
-    func switchContext(scope: Scope?) -> StoreContext {
+    func switchContext(with cache: some AtomCacheProtocol) -> StoreContext {
         StoreContext(
             store: store,
             rootScope: rootScope,
-            currentScope: scope
+            currentScope: cache.initializedScope
         )
     }
 }

Tests/AtomsTests/Attribute/RefreshableTests.swift

@@ -96,6 +96,34 @@ final class RefreshableTests: XCTestCase {
         }
     }
 
+    @MainActor
+    func testCustomRefreshNotCached() async {
+        let store = AtomStore()
+        let dependencyAtom = TestValueAtom(value: 1)
+        let atom = TestCustomRefreshableAtom { _ in
+            0
+        } refresh: { context in
+            context.read(dependencyAtom)
+        }
+        let rootScopeToken = ScopeKey.Token()
+        let scopeToken = ScopeKey.Token()
+        let scopedContext =
+            StoreContext
+            .root(store: store, scopeKey: rootScopeToken.key)
+            .scoped(
+                scopeID: ScopeID(DefaultScopeID()),
+                scopeKey: scopeToken.key,
+                observers: [],
+                overrideContainer: OverrideContainer()
+                    .addingOverride(for: dependencyAtom) { _ in
+                        2
+                    }
+            )
+
+        let value = await scopedContext.refresh(atom)
+        XCTAssertEqual(value, 2)
+    }
+
     @MainActor
     func testTransitiveRefresh() async {
         let parentAtom = TestTaskAtom { 0 }

Tests/AtomsTests/Core/StoreContextTests.swift

@@ -321,6 +321,29 @@ final class StoreContextTests: XCTestCase {
         )
     }
 
+    @MainActor
+    func testRefreshNotCached() async {
+        let store = AtomStore()
+        let atom = TestAsyncPhaseAtom<Int, Never> { .success(0) }
+        let rootScopeToken = ScopeKey.Token()
+        let scopeToken = ScopeKey.Token()
+        let scopedContext =
+            StoreContext
+            .root(store: store, scopeKey: rootScopeToken.key)
+            .scoped(
+                scopeID: ScopeID(DefaultScopeID()),
+                scopeKey: scopeToken.key,
+                observers: [],
+                overrideContainer: OverrideContainer()
+                    .addingOverride(for: atom) { _ in
+                        .success(1)
+                    }
+            )
+
+        let phase = await scopedContext.refresh(atom)
+        XCTAssertEqual(phase.value, 1)
+    }
+
     @MainActor
     func testReset() {
         let store = AtomStore()