feat: Ensure that ObservableObjectAtom notifies updates after a new @Published value is set

Sources/Atoms/Context/AtomContext.swift

@@ -130,6 +130,26 @@ public protocol AtomWatchableContext: AtomContext {
     /// - Returns: The value associated with the given atom.
     @discardableResult
     func watch<Node: Atom>(_ atom: Node) -> Node.Hook.Value
+
+    /// Accesses the observable object associated with the given atom for reading and initialing watch to
+    /// receive its updates.
+    ///
+    /// This method returns an observable object for the given atom and initiate watching the atom so that
+    /// the current context to get updated when the atom notifies updates.
+    /// The observable object associated with the atom is cached until it is no longer watched to or until
+    /// it is updated.
+    ///
+    /// ```swift
+    /// let context = ...
+    /// let store = context.watch(AccountStoreAtom())
+    /// print(store.currentUser) // Prints the user value after update.
+    /// ```
+    ///
+    /// - Parameter atom: An atom that associates the observable object.
+    ///
+    /// - Returns: The observable object associated with the given atom.
+    @discardableResult
+    func watch<Node: Atom>(_ atom: Node) -> Node.Hook.Value where Node.Hook: AtomObservableObjectHook
 }
 
 public extension AtomWatchableContext {

Sources/Atoms/Context/AtomRelationContext.swift

@@ -117,7 +117,30 @@ public struct AtomRelationContext: AtomWatchableContext {
     @inlinable
     @discardableResult
     public func watch<Node: Atom>(_ atom: Node) -> Node.Hook.Value {
-        _box.watch(atom)
+        _box.watch(atom, shouldNotifyAfterUpdates: false)
+    }
+
+    /// Accesses the observable object associated with the given atom for reading and initialing watch to
+    /// receive its updates.
+    ///
+    /// This method returns an observable object for the given atom and initiate watching the atom so that
+    /// the current context to get updated when the atom notifies updates.
+    /// The observable object associated with the atom is cached until it is no longer watched to or until
+    /// it is updated.
+    ///
+    /// ```swift
+    /// let context = ...
+    /// let store = context.watch(AccountStoreAtom())
+    /// print(store.currentUser) // Prints the user value after update.
+    /// ```
+    ///
+    /// - Parameter atom: An atom that associates the observable object.
+    ///
+    /// - Returns: The observable object associated with the given atom.
+    @inlinable
+    @discardableResult
+    public func watch<Node: Atom>(_ atom: Node) -> Node.Hook.Value where Node.Hook: AtomObservableObjectHook {
+        _box.watch(atom, shouldNotifyAfterUpdates: true)
     }
 
     /// Add the termination action that will be performed when the atom will no longer be watched to
@@ -175,7 +198,7 @@ public struct AtomRelationContext: AtomWatchableContext {
 internal protocol _AnyAtomRelationContextBox {
     var store: AtomStore { get }
 
-    func watch<Node: Atom>(_ atom: Node) -> Node.Hook.Value
+    func watch<Node: Atom>(_ atom: Node, shouldNotifyAfterUpdates: Bool) -> Node.Hook.Value
     func addTermination(_ termination: @MainActor @escaping () -> Void)
     func keepUntilTermination<Object: AnyObject>(_ object: Object)
 }
@@ -200,8 +223,12 @@ internal struct _AtomRelationContextBox<Caller: Atom>: _AnyAtomRelationContextBo
     let store: AtomStore
 
     @usableFromInline
-    func watch<Node: Atom>(_ atom: Node) -> Node.Hook.Value {
-        store.watch(atom, belongTo: caller)
+    func watch<Node: Atom>(_ atom: Node, shouldNotifyAfterUpdates: Bool) -> Node.Hook.Value {
+        store.watch(
+            atom,
+            belongTo: caller,
+            shouldNotifyAfterUpdates: shouldNotifyAfterUpdates
+        )
     }
 
     @usableFromInline

Sources/Atoms/Context/AtomTestContext.swift

@@ -120,9 +120,29 @@ public struct AtomTestContext: AtomWatchableContext {
     /// - Returns: The value associated with the given atom.
     @discardableResult
     public func watch<Node: Atom>(_ atom: Node) -> Node.Hook.Value {
-        container.store.watch(atom, relationship: container.relationship) {
-            container.onUpdate?()
-        }
+        container.watch(atom, shouldNotifyAfterUpdates: false)
+    }
+
+    /// Accesses the observable object associated with the given atom for reading and initialing watch to
+    /// receive its updates.
+    ///
+    /// This method returns an observable object for the given atom and initiate watching the atom so that
+    /// the current context to get updated when the atom notifies updates.
+    /// The observable object associated with the atom is cached until it is no longer watched to or until
+    /// it is updated.
+    ///
+    /// ```swift
+    /// let context = ...
+    /// let store = context.watch(AccountStoreAtom())
+    /// print(store.currentUser) // Prints the user value after update.
+    /// ```
+    ///
+    /// - Parameter atom: An atom that associates the observable object.
+    ///
+    /// - Returns: The observable object associated with the given atom.
+    @discardableResult
+    public func watch<Node: Atom>(_ atom: Node) -> Node.Hook.Value where Node.Hook: AtomObservableObjectHook {
+        container.watch(atom, shouldNotifyAfterUpdates: true)
     }
 
     /// Unwatches the given atom and do not receive any more updates of it.
@@ -199,5 +219,15 @@ private extension AtomTestContext {
         var relationship: Relationship {
             Relationship(container: relationshipContainer)
         }
+
+        func watch<Node: Atom>(_ atom: Node, shouldNotifyAfterUpdates: Bool) -> Node.Hook.Value {
+            store.watch(
+                atom,
+                relationship: relationship,
+                shouldNotifyAfterUpdates: shouldNotifyAfterUpdates
+            ) { [weak self] in
+                self?.onUpdate?()
+            }
+        }
     }
 }

Sources/Atoms/Context/AtomViewContext.swift

@@ -127,6 +127,11 @@ public struct AtomViewContext: AtomWatchableContext {
     @discardableResult
     @inlinable
     public func watch<Node: Atom>(_ atom: Node) -> Node.Hook.Value {
-        _store.watch(atom, relationship: _relationship, notifyUpdate: _notifyUpdate)
+        _store.watch(
+            atom,
+            relationship: _relationship,
+            shouldNotifyAfterUpdates: false,
+            notifyUpdate: _notifyUpdate
+        )
     }
 }

Sources/Atoms/Core/Hook/AtomHook.swift

@@ -1,3 +1,5 @@
+import Combine
+
 /// Internal use, a hook type that determines behavioral details of atoms.
 @MainActor
 public protocol AtomHook {
@@ -20,6 +22,10 @@ public protocol AtomHook {
     func updateOverride(context: Context, with value: Value)
 }
 
+/// Internal use, a hook type that determines behavioral details of atoms which provide `ObservableObject`.
+@MainActor
+public protocol AtomObservableObjectHook: AtomHook where Value: ObservableObject {}
+
 /// Internal use, a hook type that determines behavioral details of read-write atoms.
 @MainActor
 public protocol AtomStateHook: AtomHook {

Sources/Atoms/Core/Hook/ObservableObjectHook.swift

@@ -2,7 +2,7 @@ import Combine
 
 /// Internal use, a hook type that determines behavioral details of corresponding atoms.
 @MainActor
-public struct ObservableObjectHook<ObjectType: ObservableObject>: AtomHook {
+public struct ObservableObjectHook<ObjectType: ObservableObject>: AtomObservableObjectHook {
     /// A reference type object to manage internal state.
     public final class Coordinator {
         internal var object: ObjectType?

Sources/Atoms/Core/Internal/AtomStore.swift

@@ -25,11 +25,16 @@ internal protocol AtomStore {
     func watch<Node: Atom>(
         _ atom: Node,
         relationship: Relationship,
+        shouldNotifyAfterUpdates: Bool,
         notifyUpdate: @MainActor @escaping () -> Void
     ) -> Node.Hook.Value
 
     @MainActor
-    func watch<Node: Atom, Caller: Atom>(_ atom: Node, belongTo caller: Caller) -> Node.Hook.Value
+    func watch<Node: Atom, Caller: Atom>(
+        _ atom: Node,
+        belongTo caller: Caller,
+        shouldNotifyAfterUpdates: Bool
+    ) -> Node.Hook.Value
 
     @MainActor
     func notifyUpdate<Node: Atom>(_ atom: Node)

Sources/Atoms/Core/Internal/DefaultStore.swift

@@ -36,15 +36,29 @@ internal struct DefaultStore: AtomStore {
     func watch<Node: Atom>(
         _ atom: Node,
         relationship: Relationship,
+        shouldNotifyAfterUpdates: Bool,
         notifyUpdate: @escaping @MainActor () -> Void
     ) -> Node.Hook.Value {
         assertionFailureStoreNotProvided()
-        return fallbackStore.watch(atom, relationship: relationship, notifyUpdate: notifyUpdate)
+        return fallbackStore.watch(
+            atom,
+            relationship: relationship,
+            shouldNotifyAfterUpdates: shouldNotifyAfterUpdates,
+            notifyUpdate: notifyUpdate
+        )
     }
 
-    func watch<Node: Atom, Caller: Atom>(_ atom: Node, belongTo caller: Caller) -> Node.Hook.Value {
+    func watch<Node: Atom, Caller: Atom>(
+        _ atom: Node,
+        belongTo caller: Caller,
+        shouldNotifyAfterUpdates: Bool
+    ) -> Node.Hook.Value {
         assertionFailureStoreNotProvided()
-        return fallbackStore.watch(atom, belongTo: caller)
+        return fallbackStore.watch(
+            atom,
+            belongTo: caller,
+            shouldNotifyAfterUpdates: shouldNotifyAfterUpdates
+        )
     }
 
     func notifyUpdate<Node: Atom>(_ atom: Node) {

Sources/Atoms/Core/Internal/Store.swift

@@ -1,3 +1,5 @@
+import Foundation
+
 @MainActor
 internal struct Store: AtomStore {
     private(set) weak var container: StoreContainer?
@@ -85,15 +87,29 @@ internal struct Store: AtomStore {
     func watch<Node: Atom>(
         _ atom: Node,
         relationship: Relationship,
+        shouldNotifyAfterUpdates: Bool = false,
         notifyUpdate: @MainActor @escaping () -> Void
     ) -> Node.Hook.Value {
         // Assign the observation to the given relationship.
-        relationship[atom] = host(of: atom).observe(notifyUpdate)
+        relationship[atom] = host(of: atom).observe {
+            if shouldNotifyAfterUpdates {
+                RunLoop.current.perform {
+                    notifyUpdate()
+                }
+            }
+            else {
+                notifyUpdate()
+            }
+        }
         return read(atom)
     }
 
-    func watch<Node: Atom, Caller: Atom>(_ atom: Node, belongTo caller: Caller) -> Node.Hook.Value {
-        watch(atom, relationship: host(of: caller).relationship) {
+    func watch<Node: Atom, Caller: Atom>(
+        _ atom: Node,
+        belongTo caller: Caller,
+        shouldNotifyAfterUpdates: Bool = false
+    ) -> Node.Hook.Value {
+        watch(atom, relationship: host(of: caller).relationship, shouldNotifyAfterUpdates: shouldNotifyAfterUpdates) {
             let oldValue = read(caller)
 
             // Terminate the value & the ongoing task, but keep assignment until finishing notify update.

Tests/AtomsTests/Atom/ObservableObjectAtomTests.swift

@@ -25,15 +25,17 @@ final class ObservableObjectAtomTests: XCTestCase {
         let atom = TestAtom(value: 100)
         let context = AtomTestContext()
         let object = context.watch(atom)
-        var isUpdated = false
+        let expectation = expectation(description: "test")
+        var updatedValue: Int?
 
-        XCTAssertEqual(object.value, 100)
-        XCTAssertFalse(isUpdated)
+        context.onUpdate = {
+            updatedValue = object.value
+            expectation.fulfill()
+        }
 
-        context.onUpdate = { isUpdated = true }
         object.value = 200
 
-        XCTAssertEqual(object.value, 200)
-        XCTAssertTrue(isUpdated)
+        wait(for: [expectation], timeout: 1)
+        XCTAssertEqual(updatedValue, 200)
     }
 }

Tests/AtomsTests/Core/Hook/ObservableObjectHookTests.swift

@@ -4,7 +4,15 @@ import XCTest
 
 @MainActor
 final class ObservableObjectHookTests: XCTestCase {
-    final class TestObject: ObservableObject {}
+    @MainActor
+    final class TestObject: ObservableObject {
+        @Published
+        private(set) var updatedCount = 0
+
+        func update() {
+            updatedCount += 1
+        }
+    }
 
     func testMakeCoordinator() {
         let object = TestObject()
@@ -33,43 +41,77 @@ final class ObservableObjectHookTests: XCTestCase {
         let hook = ObservableObjectHook { _ in TestObject() }
         let atom = TestAtom(key: 0, hook: hook)
         let context = AtomTestContext()
-        var updateCount = 0
 
-        context.onUpdate = { updateCount += 1 }
+        XCTContext.runActivity(named: "Update") { _ in
+            let object = context.watch(atom)
+            var updateCount = 0
+            let expectation = expectation(description: "Update")
 
-        let object0 = context.watch(atom)
+            context.onUpdate = {
+                updateCount = object.updatedCount
+                expectation.fulfill()
+            }
+            object.update()
 
-        XCTContext.runActivity(named: "Update") { _ in
-            object0.objectWillChange.send()
+            wait(for: [expectation], timeout: 1)
             XCTAssertEqual(updateCount, 1)
         }
 
         XCTContext.runActivity(named: "Termination") { _ in
             context.unwatch(atom)
 
-            object0.objectWillChange.send()
-            XCTAssertEqual(updateCount, 1)
-        }
+            let object = context.watch(atom)
+            var updateCount = 0
+            let expectation = expectation(description: "Termination")
 
-        let overrideObject = TestObject()
-        context.override(atom) { _ in overrideObject }
+            context.onUpdate = {
+                updateCount = object.updatedCount
+                expectation.fulfill()
+            }
+            object.update()
 
-        let object1 = context.watch(atom)
+            wait(for: [expectation], timeout: 1)
+            XCTAssertEqual(updateCount, 1)
+        }
 
         XCTContext.runActivity(named: "Override") { _ in
-            XCTAssertTrue(object1 === overrideObject)
+            let overrideObject = TestObject()
+            context.unwatch(atom)
+            context.override(atom) { _ in overrideObject }
+
+            let object = context.watch(atom)
+            var updateCount = 0
+            let expectation = expectation(description: "Override")
 
-            object1.objectWillChange.send()
+            context.onUpdate = {
+                updateCount = object.updatedCount
+                expectation.fulfill()
+            }
+            object.update()
 
-            XCTAssertEqual(updateCount, 2)
+            XCTAssertTrue(object === overrideObject)
+            wait(for: [expectation], timeout: 1)
+            XCTAssertEqual(updateCount, 1)
         }
 
         XCTContext.runActivity(named: "Override termination") { _ in
+            let overrideObject = TestObject()
             context.unwatch(atom)
+            context.override(atom) { _ in overrideObject }
+
+            let object = context.watch(atom)
+            var updateCount = 0
+            let expectation = expectation(description: "Override termination")
 
-            object1.objectWillChange.send()
+            context.onUpdate = {
+                updateCount = object.updatedCount
+                expectation.fulfill()
+            }
+            object.update()
 
-            XCTAssertEqual(updateCount, 2)
+            XCTAssertTrue(object === overrideObject)
+            wait(for: [expectation], timeout: 1)
+            XCTAssertEqual(updateCount, 1)
         }
     }
 }