Feature: Custom Resettable Attribute

README.md

@@ -805,6 +805,38 @@ struct FetchMoviesPhaseAtom: ValueAtom, Refreshable, Hashable {
 }
 ```
 
+
+#### [Resettable](https://ra1028.github.io/swiftui-atom-properties/documentation/atoms/resettable)
+
+`Resettable` allows you to implement a custom reset behavior to an atom.
+
+<details><summary><code>📖 Expand to see example</code></summary>
+
+It adds custom reset behavior to an Atom that will be executed upon atom reset.
+
+It's useful when need to have arbitrary reset ability or implementing reset when value depends on private atom.
+
+In following example, `RandomIntAtom` generates a random value using generated from private `RandomNumberGeneratorAtom`, and `Resettable` gives ability to replace exposed reset with  `RandomNumberGeneratorAtom` reset.
+
+```swift
+struct RandomIntAtom: ValueAtom, Resettable, Hashable {
+    func value(context: Context) -> Int {
+        var generator = context.watch(RandomNumberGeneratorAtom())
+        return .random(in: 0..<100, using: &generator)
+    }
+
+    func reset(context: ResetContext) {
+        context.reset(RandomNumberGeneratorAtom())
+    }
+}
+
+private struct RandomNumberGeneratorAtom: ValueAtom, Hashable {
+    func value(context: Context) -> CustomRandomNumberGenerator {
+        CustomRandomNumberGenerator()
+    }
+}
+```
+
 </details>
 
 ---

Sources/Atoms/Atoms.docc/Atoms.md

@@ -35,6 +35,7 @@ Building state by compositing atoms automatically optimizes rendering based on i
 
 - ``KeepAlive``
 - ``Refreshable``
+- ``Resettable``
 
 ### Property Wrappers
 

Sources/Atoms/Attribute/Resettable.swift

@@ -0,0 +1,36 @@
+/// An attribute protocol allows an atom to have a custom reset override.
+///
+/// Note that the custom reset will be triggered even when the atom is overridden.
+///
+/// ```swift
+/// struct UserAtom: ValueAtom, Resettable, Hashable {
+///     func value(context: Context) -> User? {
+///          context.watch(FetchUserAtom()).phase.value
+///     }
+///
+///     func reset(context: ResetContext) {
+///          context.reset(FetchUserAtom())
+///     }
+/// }
+///
+/// private struct FetchUserAtom: TaskAtom, Hashable {
+///     func value(context: Context) async -> User? {
+///          await fetchUser()
+///     }
+/// }
+/// ```
+///
+public protocol Resettable where Self: Atom {
+    /// A type of the context structure to read, set, and otherwise interact
+    /// with other atoms.
+    typealias ResetContext = AtomCurrentContext<Loader.Coordinator>
+
+    /// Arbitrary reset method to be executed on atom reset.
+    ///
+    /// This is arbitrary custom reset method that replaces regular atom reset functionality.
+    ///
+    /// - Parameter context: A context structure to read, set, and otherwise interact
+    ///                      with other atoms.
+    @MainActor
+    func reset(context: ResetContext)
+}

Sources/Atoms/Context/AtomContext.swift

@@ -27,7 +27,7 @@ public protocol AtomContext {
     /// and assigns a new value for the atom.
     /// When you assign a new value, it immediately notifies downstream atoms and views.
     ///
-    /// - SeeAlso: ``AtomContext/subscript``
+    /// - SeeAlso: ``AtomContext/subscript(_:)``
     ///
     /// ```swift
     /// let context = ...
@@ -102,7 +102,25 @@ public protocol AtomContext {
     /// ```
     ///
     /// - Parameter atom: An atom to reset.
-    func reset(_ atom: some Atom)
+    @_disfavoredOverload
+    func reset<Node: Atom>(_ atom: Node)
+
+    /// Calls arbitrary reset function of the given atom.
+    ///
+    /// This method only accepts atoms that conform to ``Resettable`` protocol.
+    /// Calls custom reset function of the given atom. Hence, it does not generate any new cache value or notify subscribers.
+    ///
+    /// ```swift
+    /// let context = ...
+    /// print(context.watch(ResettableTextAtom()) // Prints "Text"
+    /// context[ResettableTextAtom()] = "New text"
+    /// print(context.read(ResettableTextAtom())) // Prints "New text"
+    /// context.reset(ResettableTextAtom()) // Calls the custom reset function
+    /// print(context.read(ResettableTextAtom())) // Prints "New text"
+    /// ```
+    ///
+    /// - Parameter atom: An atom to reset.
+    func reset<Node: Resettable>(_ atom: Node)
 }
 
 public extension AtomContext {

Sources/Atoms/Context/AtomCurrentContext.swift

@@ -40,7 +40,7 @@ public struct AtomCurrentContext<Coordinator>: AtomContext {
     /// and assigns a new value for the atom.
     /// When you assign a new value, it immediately notifies downstream atoms and views.
     ///
-    /// - SeeAlso: ``AtomViewContext/subscript``
+    /// - SeeAlso: ``AtomViewContext/subscript(_:)``
     ///
     /// ```swift
     /// let context = ...
@@ -142,7 +142,28 @@ public struct AtomCurrentContext<Coordinator>: AtomContext {
     ///
     /// - Parameter atom: An atom to reset.
     @inlinable
-    public func reset(_ atom: some Atom) {
+    @_disfavoredOverload
+    public func reset<Node: Atom>(_ atom: Node) {
+        _store.reset(atom)
+    }
+
+    /// Calls arbitrary reset function of the given atom.
+    ///
+    /// This method only accepts atoms that conform to ``Resettable`` protocol.
+    /// Calls custom reset function of the given atom. Hence, it does not generate any new cache value or notify subscribers.
+    ///
+    /// ```swift
+    /// let context = ...
+    /// print(context.watch(ResettableTextAtom()) // Prints "Text"
+    /// context[ResettableTextAtom()] = "New text"
+    /// print(context.read(ResettableTextAtom())) // Prints "New text"
+    /// context.reset(ResettableTextAtom()) // Calls the custom reset function
+    /// print(context.read(ResettableTextAtom())) // Prints "New text"
+    /// ```
+    ///
+    /// - Parameter atom: An atom to reset.
+    @inlinable
+    public func reset<Node: Resettable>(_ atom: Node) {
         _store.reset(atom)
     }
 }

Sources/Atoms/Context/AtomTestContext.swift

@@ -173,7 +173,7 @@ public struct AtomTestContext: AtomWatchableContext {
     /// and assigns a new value for the atom.
     /// When you assign a new value, it immediately notifies downstream atoms and views.
     ///
-    /// - SeeAlso: ``AtomTestContext/subscript``
+    /// - SeeAlso: ``AtomTestContext/subscript(_:)``
     ///
     /// ```swift
     /// let context = AtomTestContext()
@@ -277,7 +277,28 @@ public struct AtomTestContext: AtomWatchableContext {
     ///
     /// - Parameter atom: An atom to reset.
     @inlinable
-    public func reset(_ atom: some Atom) {
+    @_disfavoredOverload
+    public func reset<Node: Atom>(_ atom: Node) {
+        _store.reset(atom)
+    }
+
+    /// Calls arbitrary reset function of the given atom.
+    ///
+    /// This method only accepts atoms that conform to ``Resettable`` protocol.
+    /// Calls custom reset function of the given atom. Hence, it does not generate any new cache value or notify subscribers.
+    ///
+    /// ```swift
+    /// let context = ...
+    /// print(context.watch(ResettableTextAtom()) // Prints "Text"
+    /// context[ResettableTextAtom()] = "New text"
+    /// print(context.read(ResettableTextAtom())) // Prints "New text"
+    /// context.reset(ResettableTextAtom()) // Calls the custom reset function
+    /// print(context.read(ResettableTextAtom())) // Prints "New text"
+    /// ```
+    ///
+    /// - Parameter atom: An atom to reset.
+    @inlinable
+    public func reset<Node: Resettable>(_ atom: Node) {
         _store.reset(atom)
     }
 

Sources/Atoms/Context/AtomTransactionContext.swift

@@ -47,7 +47,7 @@ public struct AtomTransactionContext<Coordinator>: AtomWatchableContext {
     /// and assigns a new value for the atom.
     /// When you assign a new value, it immediately notifies downstream atoms and views.
     ///
-    /// - SeeAlso: ``AtomTransactionContext/subscript``
+    /// - SeeAlso: ``AtomTransactionContext/subscript(_:)``
     ///
     /// ```swift
     /// let context = ...
@@ -142,16 +142,37 @@ public struct AtomTransactionContext<Coordinator>: AtomWatchableContext {
     ///
     /// ```swift
     /// let context = ...
-    /// print(context.watch(TextAtom())) // Prints "Text"
-    /// context[TextAtom()] = "New text"
-    /// print(context.read(TextAtom())) // Prints "New text"
-    /// context.reset(TextAtom())
-    /// print(context.read(TextAtom())) // Prints "Text"
+    /// print(context.watch(ResettableTextAtom())) // Prints "Text"
+    /// context[ResettableTextAtom()] = "New text"
+    /// print(context.read(ResettableTextAtom())) // Prints "New text"
+    /// context.reset(ResettableTextAtom())
+    /// print(context.read(ResettableTextAtom())) // Prints "Text"
+    /// ```
+    ///
+    /// - Parameter atom: An atom to reset.
+    @inlinable
+    @_disfavoredOverload
+    public func reset<Node: Atom>(_ atom: Node) {
+        _store.reset(atom)
+    }
+
+    /// Calls arbitrary reset function of the given atom.
+    ///
+    /// This method only accepts atoms that conform to ``Resettable`` protocol.
+    /// Calls custom reset function of the given atom. Hence, it does not generate any new cache value or notify subscribers.
+    ///
+    /// ```swift
+    /// let context = ...
+    /// print(context.watch(ResettableTextAtom()) // Prints "Text"
+    /// context[ResettableTextAtom()] = "New text"
+    /// print(context.read(ResettableTextAtom())) // Prints "New text"
+    /// context.reset(ResettableTextAtom()) // Calls the custom reset function
+    /// print(context.read(ResettableTextAtom())) // Prints "New text"
     /// ```
     ///
     /// - Parameter atom: An atom to reset.
     @inlinable
-    public func reset(_ atom: some Atom) {
+    public func reset<Node: Resettable>(_ atom: Node) {
         _store.reset(atom)
     }
 

Sources/Atoms/Context/AtomViewContext.swift

@@ -46,7 +46,7 @@ public struct AtomViewContext: AtomWatchableContext {
     /// and assigns a new value for the atom.
     /// When you assign a new value, it immediately notifies downstream atoms and views.
     ///
-    /// - SeeAlso: ``AtomViewContext/subscript``
+    /// - SeeAlso: ``AtomViewContext/subscript(_:)``
     ///
     /// ```swift
     /// let context = ...
@@ -150,7 +150,28 @@ public struct AtomViewContext: AtomWatchableContext {
     ///
     /// - Parameter atom: An atom to reset.
     @inlinable
-    public func reset(_ atom: some Atom) {
+    @_disfavoredOverload
+    public func reset<Node: Atom>(_ atom: Node) {
+        _store.reset(atom)
+    }
+
+    /// Calls arbitrary reset function of the given atom.
+    ///
+    /// This method only accepts atoms that conform to ``Resettable`` protocol.
+    /// Calls custom reset function of the given atom. Hence, it does not generate any new cache value or notify subscribers.
+    ///
+    /// ```swift
+    /// let context = ...
+    /// print(context.watch(ResettableTextAtom()) // Prints "Text"
+    /// context[ResettableTextAtom()] = "New text"
+    /// print(context.read(ResettableTextAtom())) // Prints "New text"
+    /// context.reset(ResettableTextAtom()) // Calls the custom reset function
+    /// print(context.read(ResettableTextAtom())) // Prints "New text"
+    /// ```
+    ///
+    /// - Parameter atom: An atom to reset.
+    @inlinable
+    public func reset<Node: Resettable>(_ atom: Node) {
         _store.reset(atom)
     }
 

Sources/Atoms/Core/StoreContext.swift

@@ -201,7 +201,8 @@ internal struct StoreContext {
     }
 
     @usableFromInline
-    func reset(_ atom: some Atom) {
+    @_disfavoredOverload
+    func reset<Node: Atom>(_ atom: Node) {
         let override = lookupOverride(of: atom)
         let key = AtomKey(atom, overrideScopeKey: override?.scopeKey)
 
@@ -211,6 +212,15 @@ internal struct StoreContext {
         }
     }
 
+    @usableFromInline
+    func reset<Node: Resettable>(_ atom: Node) {
+        let override = lookupOverride(of: atom)
+        let key = AtomKey(atom, overrideScopeKey: override?.scopeKey)
+        let state = getState(of: atom, for: key)
+        let context = AtomCurrentContext(store: self, coordinator: state.coordinator)
+        atom.reset(context: context)
+    }
+
     @usableFromInline
     func lookup<Node: Atom>(_ atom: Node) -> Node.Loader.Value? {
         let override = lookupOverride(of: atom)

Tests/AtomsTests/Atom/TaskAtomTests.swift

@@ -66,9 +66,6 @@ final class TaskAtomTests: XCTestCase {
 
         do {
             // Cancellation
-            var updateCount = 0
-            context.onUpdate = { updateCount += 1 }
-
             let refreshTask0 = Task {
                 await context.refresh(atom)
             }

Tests/AtomsTests/Atom/ThrowingTaskAtomTests.swift

@@ -81,9 +81,6 @@ final class ThrowingTaskAtomTests: XCTestCase {
 
         do {
             // Cancellation
-            var updateCount = 0
-            context.onUpdate = { updateCount += 1 }
-
             let refreshTask = Task {
                 await context.refresh(atom)
             }

Tests/AtomsTests/Context/AtomCurrentContextTests.swift

@@ -72,4 +72,36 @@ final class AtomCurrentContextTests: XCTestCase {
 
         XCTAssertEqual(context.read(dependency), 0)
     }
+
+    @MainActor
+    func testCustomReset() {
+        let store = AtomStore()
+        let context = AtomCurrentContext(store: StoreContext(store), coordinator: ())
+        let storeContext = StoreContext(store)
+        let transactionAtom = TestValueAtom(value: 0)
+        let atom = TestStateAtom(defaultValue: 0)
+        let transaction = Transaction(key: AtomKey(transactionAtom)) {}
+
+        let resettableAtom = TestCustomResettableAtom(
+            defaultValue: { context in
+                context.watch(atom)
+            },
+            reset: { context in
+                context[atom] = 300
+            }
+        )
+
+        XCTAssertEqual(storeContext.watch(atom, in: transaction), 0)
+        XCTAssertEqual(storeContext.watch(resettableAtom, in: transaction), 0)
+
+        context[atom] = 100
+
+        XCTAssertEqual(storeContext.watch(atom, in: transaction), 100)
+        XCTAssertEqual(storeContext.watch(resettableAtom, in: transaction), 100)
+
+        context.reset(resettableAtom)
+
+        XCTAssertEqual(storeContext.watch(atom, in: transaction), 300)
+        XCTAssertEqual(storeContext.watch(resettableAtom, in: transaction), 300)
+    }
 }