Detach atom dependencies while executing custom refresh

Sources/Atoms/Core/StoreContext.swift

@@ -180,8 +180,15 @@ internal struct StoreContext {
         let key = AtomKey(atom, scopeKey: scopeKey)
         let state = getState(of: atom, for: key)
         let context = AtomCurrentContext(store: self, coordinator: state.coordinator)
+
+        // Detach the dependencies once to delay updating the downstream until
+        // this atom's value refresh is complete.
+        let dependencies = detachDependencies(for: key)
         let value = await atom.refresh(context: context)
 
+        // Restore dependencies when the refresh is completed.
+        attachDependencies(dependencies, for: key)
+
         guard let transaction = state.transaction, let cache = lookupCache(of: atom, for: key) else {
             checkAndRelease(for: key)
             return value
@@ -454,6 +461,28 @@ private extension StoreContext {
         release(for: key)
     }
 
+    func detachDependencies(for key: AtomKey) -> Set<AtomKey> {
+        // Remove current dependencies.
+        let dependencies = store.graph.dependencies.removeValue(forKey: key) ?? []
+
+        // Detatch the atom from its children.
+        for dependency in dependencies {
+            store.graph.children[dependency]?.remove(key)
+        }
+
+        return dependencies
+    }
+
+    func attachDependencies(_ dependencies: Set<AtomKey>, for key: AtomKey) {
+        // Set dependencies.
+        store.graph.dependencies[key] = dependencies
+
+        // Attach the atom to its children.
+        for dependency in dependencies {
+            store.graph.children[dependency]?.insert(key)
+        }
+    }
+
     func unsubscribe<Keys: Sequence<AtomKey>>(_ keys: Keys, for subscriberKey: SubscriberKey) {
         for key in keys {
             store.state.subscriptions[key]?.removeValue(forKey: subscriberKey)
@@ -468,13 +497,7 @@ private extension StoreContext {
         for key: AtomKey
     ) -> AtomProducerContext<Node.Produced, Node.Coordinator> {
         let transaction = Transaction(key: key) {
-            // Remove current dependencies.
-            let oldDependencies = store.graph.dependencies.removeValue(forKey: key) ?? []
-
-            // Detatch the atom from its children.
-            for dependency in oldDependencies {
-                store.graph.children[dependency]?.remove(key)
-            }
+            let oldDependencies = detachDependencies(for: key)
 
             return {
                 let dependencies = store.graph.dependencies[key] ?? []

Tests/AtomsTests/Attribute/RefreshableTests.swift

@@ -1,4 +1,3 @@
-import Combine
 import XCTest
 
 @testable import Atoms
@@ -9,10 +8,10 @@ final class RefreshableTests: XCTestCase {
         let store = AtomStore()
         let subscriberState = SubscriberState()
         let subscriber = Subscriber(subscriberState)
-        let atom = TestCustomRefreshableAtom {
-            Just(0)
-        } refresh: {
-            .success(1)
+        let atom = TestCustomRefreshableAtom { _ in
+            0
+        } refresh: { _ in
+            1
         }
         let key = AtomKey(atom)
         var snapshots = [Snapshot]()
@@ -22,33 +21,32 @@ final class RefreshableTests: XCTestCase {
         do {
             // Should call custom refresh behavior
 
-            let phase0 = await context.refresh(atom)
-            XCTAssertEqual(phase0.value, 1)
+            let value0 = await context.refresh(atom)
+            XCTAssertEqual(value0, 1)
             XCTAssertNil(store.state.caches[key])
             XCTAssertNil(store.state.states[key])
             XCTAssertTrue(snapshots.isEmpty)
 
             var updateCount = 0
-            let phase1 = context.watch(
+            let value1 = context.watch(
                 atom,
                 subscriber: subscriber,
                 subscription: Subscription {
                     updateCount += 1
                 }
             )
 
-            XCTAssertTrue(phase1.isSuspending)
+            XCTAssertEqual(value1, 0)
 
             snapshots.removeAll()
-
-            let phase2 = await context.refresh(atom)
-            XCTAssertEqual(phase2.value, 1)
+            let value2 = await context.refresh(atom)
+            XCTAssertEqual(value2, 1)
             XCTAssertNotNil(store.state.states[key])
-            XCTAssertEqual((store.state.caches[key] as? AtomCache<TestCustomRefreshableAtom<Just<Int>>>)?.value, .success(1))
+            XCTAssertEqual((store.state.caches[key] as? AtomCache<TestCustomRefreshableAtom<Int>>)?.value, 1)
             XCTAssertEqual(updateCount, 1)
             XCTAssertEqual(
-                snapshots.map { $0.caches.mapValues { $0.value as? AsyncPhase<Int, Never> } },
-                [[key: .success(1)]]
+                snapshots.map { $0.caches.mapValues { $0.value as? Int } },
+                [[key: 1]]
             )
 
             context.unwatch(atom, subscriber: subscriber)
@@ -64,30 +62,58 @@ final class RefreshableTests: XCTestCase {
                 scopeID: ScopeID(DefaultScopeID()),
                 observers: [],
                 overrides: [
-                    OverrideKey(atom): Override<TestCustomRefreshableAtom<Just<Int>>>(isScoped: true) { _ in .success(2) }
+                    OverrideKey(atom): Override<TestCustomRefreshableAtom<Int>>(isScoped: true) { _ in 2 }
                 ]
             )
 
-            let phase0 = scopedContext.watch(atom, subscriber: subscriber, subscription: Subscription())
-            XCTAssertEqual(phase0.value, 2)
+            let value0 = scopedContext.watch(atom, subscriber: subscriber, subscription: Subscription())
+            XCTAssertEqual(value0, 2)
 
-            let phase1 = await scopedContext.refresh(atom)
-            XCTAssertEqual(phase1.value, 1)
+            let value1 = await scopedContext.refresh(atom)
+            XCTAssertEqual(value1, 1)
             XCTAssertNotNil(store.state.states[overrideAtomKey])
             XCTAssertEqual(
-                (store.state.caches[overrideAtomKey] as? AtomCache<TestCustomRefreshableAtom<Just<Int>>>)?.value,
-                .success(1)
+                (store.state.caches[overrideAtomKey] as? AtomCache<TestCustomRefreshableAtom<Int>>)?.value,
+                1
             )
         }
 
         do {
             // Should not make new state and cache
 
-            let phase = await context.refresh(atom)
+            let value = await context.refresh(atom)
 
-            XCTAssertEqual(phase.value, 1)
+            XCTAssertEqual(value, 1)
             XCTAssertNil(store.state.states[key])
             XCTAssertNil(store.state.caches[key])
         }
     }
+
+    @MainActor
+    func testTransitiveRefresh() async {
+        let parentAtom = TestTaskAtom { 0 }
+        let atom = TestCustomRefreshableAtom { context in
+            context.watch(parentAtom.phase)
+        } refresh: { context in
+            await context.refresh(parentAtom.phase)
+        }
+        let context = AtomTestContext()
+
+        var updateCount = 0
+        context.onUpdate = {
+            updateCount += 1
+        }
+
+        XCTAssertTrue(context.watch(atom).isSuspending)
+
+        await context.waitForUpdate()
+        XCTAssertEqual(context.watch(atom).value, 0)
+        XCTAssertEqual(updateCount, 1)
+
+        let value = await context.refresh(atom).value
+
+        XCTAssertEqual(value, 0)
+        XCTAssertEqual(context.watch(atom).value, 0)
+        XCTAssertEqual(updateCount, 2)
+    }
 }

Tests/AtomsTests/Attribute/ResettableTests.swift

@@ -99,4 +99,28 @@ final class ResettableTests: XCTestCase {
             XCTAssertNil(store.state.caches[key])
         }
     }
+
+    @MainActor
+    func testTransitiveReset() {
+        let parentAtom = TestValueAtom(value: 0)
+        let atom = TestCustomResettableAtom { context in
+            context.watch(parentAtom)
+        } reset: { context in
+            context.reset(parentAtom)
+        }
+        let context = AtomTestContext()
+
+        var updateCount = 0
+        context.onUpdate = {
+            updateCount += 1
+        }
+
+        XCTAssertEqual(context.watch(atom), 0)
+        XCTAssertEqual(updateCount, 0)
+
+        context.reset(atom)
+
+        XCTAssertEqual(context.watch(atom), 0)
+        XCTAssertEqual(updateCount, 1)
+    }
 }

Tests/AtomsTests/Context/AtomCurrentContextTests.swift

@@ -41,15 +41,15 @@ final class AtomCurrentContextTests: XCTestCase {
 
     @MainActor
     func testCustomRefresh() async {
-        let atom = TestCustomRefreshableAtom {
-            Just(100)
-        } refresh: {
-            .success(200)
+        let atom = TestCustomRefreshableAtom { _ in
+            100
+        } refresh: { _ in
+            200
         }
         let store = AtomStore()
         let context = AtomCurrentContext(store: StoreContext(store: store), coordinator: ())
-        let value = await context.refresh(atom).value
 
+        let value = await context.refresh(atom)
         XCTAssertEqual(value, 200)
     }
 

Tests/AtomsTests/Context/AtomTestContextTests.swift

@@ -211,10 +211,10 @@ final class AtomTestContextTests: XCTestCase {
 
     @MainActor
     func testCustomRefresh() async {
-        let atom = TestCustomRefreshableAtom {
-            Just(100)
-        } refresh: {
-            .success(200)
+        let atom = TestCustomRefreshableAtom { _ in
+            100
+        } refresh: { _ in
+            200
         }
         let context = AtomTestContext()
         var updateCount = 0
@@ -223,17 +223,13 @@ final class AtomTestContextTests: XCTestCase {
             updateCount += 1
         }
 
-        XCTAssertTrue(context.watch(atom).isSuspending)
-
-        await context.waitForUpdate()
-
-        XCTAssertEqual(context.watch(atom).value, 100)
+        XCTAssertEqual(context.watch(atom), 100)
 
-        let value = await context.refresh(atom).value
+        let value = await context.refresh(atom)
 
         XCTAssertEqual(value, 200)
-        XCTAssertEqual(context.watch(atom).value, 200)
-        XCTAssertEqual(updateCount, 2)
+        XCTAssertEqual(context.watch(atom), 200)
+        XCTAssertEqual(updateCount, 1)
     }
 
     @MainActor

Tests/AtomsTests/Context/AtomTransactionContextTests.swift

@@ -48,21 +48,20 @@ final class AtomTransactionContextTests: XCTestCase {
     @MainActor
     func testCustomRefresh() async {
         let atom0 = TestValueAtom(value: 0)
-        let atom1 = TestCustomRefreshableAtom {
-            Just(100)
-        } refresh: {
-            .success(200)
+        let atom1 = TestCustomRefreshableAtom { _ in
+            100
+        } refresh: { _ in
+            200
         }
         let store = AtomStore()
         let transaction = Transaction(key: AtomKey(atom0))
         let context = AtomTransactionContext(store: StoreContext(store: store), transaction: transaction, coordinator: ())
 
-        XCTAssertTrue(context.watch(atom1).isSuspending)
-
-        let value = await context.refresh(atom1).value
+        XCTAssertEqual(context.watch(atom1), 100)
 
+        let value = await context.refresh(atom1)
         XCTAssertEqual(value, 200)
-        XCTAssertEqual(context.watch(atom1).value, 200)
+        XCTAssertEqual(context.watch(atom1), 200)
     }
 
     @MainActor

Tests/AtomsTests/Context/AtomViewContextTests.swift

@@ -57,10 +57,10 @@ final class AtomViewContextTests: XCTestCase {
 
     @MainActor
     func testCustomRefresh() async {
-        let atom = TestCustomRefreshableAtom {
-            Just(100)
-        } refresh: {
-            .success(200)
+        let atom = TestCustomRefreshableAtom { _ in
+            100
+        } refresh: { _ in
+            200
         }
         let store = AtomStore()
         let subscriberState = SubscriberState()
@@ -70,12 +70,11 @@ final class AtomViewContextTests: XCTestCase {
             subscription: Subscription()
         )
 
-        XCTAssertTrue(context.watch(atom).isSuspending)
-
-        let value = await context.refresh(atom).value
+        XCTAssertEqual(context.watch(atom), 100)
 
+        let value = await context.refresh(atom)
         XCTAssertEqual(value, 200)
-        XCTAssertEqual(context.watch(atom).value, 200)
+        XCTAssertEqual(context.watch(atom), 200)
     }
 
     @MainActor

Tests/AtomsTests/Utilities/TestAtom.swift

@@ -78,20 +78,20 @@ struct TestThrowingTaskAtom<Success: Sendable>: ThrowingTaskAtom {
     }
 }
 
-struct TestCustomRefreshableAtom<Publisher: Combine.Publisher>: PublisherAtom, Refreshable {
-    var makePublisher: () -> Publisher
-    var refresh: () -> AsyncPhase<Publisher.Output, Publisher.Failure>
+struct TestCustomRefreshableAtom<T>: ValueAtom, Refreshable {
+    var getValue: (Context) -> T
+    var refresh: (CurrentContext) async -> T
 
     var key: UniqueKey {
         UniqueKey()
     }
 
-    func publisher(context: Context) -> Publisher {
-        makePublisher()
+    func value(context: Context) -> T {
+        getValue(context)
     }
 
-    func refresh(context: CurrentContext) async -> AsyncPhase<Publisher.Output, Publisher.Failure> {
-        refresh()
+    func refresh(context: CurrentContext) async -> T {
+        await refresh(context)
     }
 }