Move restore function to ViewContext (#66)

* Remove restore function from Snapshot itself and move it to AtomViewContext

* Fix example app

* Update README

Author: ra1028

Examples/Packages/iOS/Sources/ExampleTimeTravel/ExampleTimeTravel.swift

@@ -78,6 +78,9 @@ struct TimeTravelDebug<Content: View>: View {
     @State
     var position = 0
 
+    @ViewContext
+    var context
+
     var body: some View {
         AtomScope {
             ZStack(alignment: .bottom) {
@@ -105,7 +108,7 @@ struct TimeTravelDebug<Content: View>: View {
                     set: { value in
                         Task { @MainActor in
                             position = Int(value)
-                            snapshots[position].restore()
+                            context.restore(snapshots[position])
                         }
                     }
                 ),

README.md

@@ -1302,9 +1302,30 @@ AtomRoot {
 }
 ```
 
-Calling the [restore()](https://ra1028.github.io/swiftui-atom-properties/documentation/atoms/snapshot/restore()) method of the obtained `Snapshot` will roll back to the states and dependency graph at that point in time to see what happened.  
+`@ViewContext` also supports restoring the values of atoms and the dependency graph captured at a point in time in a retrieved snapshot and its dependency graph so that you can investigate what happend.  
 The debugging technique is called [time travel debugging](https://en.wikipedia.org/wiki/Time_travel_debugging), and the example application [here](Examples/Packages/iOS/Sources/ExampleTimeTravel) demonstrates how it works.  
 
+```swift
+@ViewContext
+var context
+
+@State
+var snapshot: Snapshot?
+
+var body: some View {
+    VStack {
+        Button("Capture") {
+            snapshot = context.snapshot()
+        }
+        Button("Restore") {
+            if let snapshot {
+                context.restore(snapshot)
+            }
+        }
+    }
+}
+```
+
 In addition, [graphDescription()](https://ra1028.github.io/swiftui-atom-properties/documentation/atoms/snapshot/graphdescription()) method returns a string, that represents the dependencies graph and where they are used, as a String in [graph description language DOT](https://graphviz.org/doc/info/lang.html).  
 This can be converted to an image using [Graphviz](https://graphviz.org), a graph visualization tool, to visually analyze information about the state of the application, as shown below.  
 

Sources/Atoms/Context/AtomViewContext.swift

@@ -170,4 +170,14 @@ public struct AtomViewContext: AtomWatchableContext {
     public func snapshot() -> Snapshot {
         _store.snapshot()
     }
+
+    /// For debugging, restore atom values and the dependency graph captured at a point in time in the given snapshot.
+    ///
+    /// Atoms and their dependencies that are no longer subscribed to from anywhere are then released.
+    ///
+    /// - Parameter snapshot: A snapshot that captures specific set of values of atoms.
+    @inlinable
+    public func restore(_ snapshot: Snapshot) {
+        _store.restore(snapshot)
+    }
 }

Sources/Atoms/Core/StoreContext.swift

@@ -191,44 +191,44 @@ internal struct StoreContext {
     @usableFromInline
     func snapshot() -> Snapshot {
         let store = getStore()
-        let graph = store.graph
-        let caches = store.state.caches
-        let subscriptions = store.state.subscriptions
 
         return Snapshot(
-            graph: graph,
-            caches: caches,
-            subscriptions: subscriptions
-        ) {
-            let store = getStore()
-            let keys = ContiguousArray(caches.keys)
-            var obsoletedDependencies = [AtomKey: Set<AtomKey>]()
-
-            for key in keys {
-                let oldDependencies = store.graph.dependencies[key]
-                let newDependencies = graph.dependencies[key]
-
-                // Update atom values and the graph.
-                store.state.caches[key] = caches[key]
-                store.graph.dependencies[key] = newDependencies
-                store.graph.children[key] = graph.children[key]
-                obsoletedDependencies[key] = oldDependencies?.subtracting(newDependencies ?? [])
-            }
+            graph: store.graph,
+            caches: store.state.caches,
+            subscriptions: store.state.subscriptions
+        )
+    }
 
-            for key in keys {
-                // Release if the atom is no longer used.
-                checkRelease(for: key)
+    @usableFromInline
+    func restore(_ snapshot: Snapshot) {
+        let store = getStore()
+        let keys = ContiguousArray(snapshot.caches.keys)
+        var obsoletedDependencies = [AtomKey: Set<AtomKey>]()
 
-                // Release dependencies that are no longer dependent.
-                if let dependencies = obsoletedDependencies[key] {
-                    checkReleaseDependencies(dependencies, for: key)
-                }
+        for key in keys {
+            let oldDependencies = store.graph.dependencies[key]
+            let newDependencies = snapshot.graph.dependencies[key]
 
-                // Notify updates only for the subscriptions of restored atoms.
-                if let subscriptions = store.state.subscriptions[key] {
-                    for subscription in ContiguousArray(subscriptions.values) {
-                        subscription.notifyUpdate()
-                    }
+            // Update atom values and the graph.
+            store.state.caches[key] = snapshot.caches[key]
+            store.graph.dependencies[key] = newDependencies
+            store.graph.children[key] = snapshot.graph.children[key]
+            obsoletedDependencies[key] = oldDependencies?.subtracting(newDependencies ?? [])
+        }
+
+        for key in keys {
+            // Release if the atom is no longer used.
+            checkRelease(for: key)
+
+            // Release dependencies that are no longer dependent.
+            if let dependencies = obsoletedDependencies[key] {
+                checkReleaseDependencies(dependencies, for: key)
+            }
+
+            // Notify updates only for the subscriptions of restored atoms.
+            if let subscriptions = store.state.subscriptions[key] {
+                for subscription in ContiguousArray(subscriptions.values) {
+                    subscription.notifyUpdate()
                 }
             }
         }

Sources/Atoms/Snapshot.swift

@@ -3,18 +3,15 @@ public struct Snapshot: CustomStringConvertible {
     internal let graph: Graph
     internal let caches: [AtomKey: any AtomCacheProtocol]
     internal let subscriptions: [AtomKey: [SubscriptionKey: Subscription]]
-    private let _restore: @MainActor () -> Void
 
     internal init(
         graph: Graph,
         caches: [AtomKey: any AtomCacheProtocol],
-        subscriptions: [AtomKey: [SubscriptionKey: Subscription]],
-        restore: @MainActor @escaping () -> Void
+        subscriptions: [AtomKey: [SubscriptionKey: Subscription]]
     ) {
         self.graph = graph
         self.caches = caches
         self.subscriptions = subscriptions
-        self._restore = restore
     }
 
     /// A textual representation of this snapshot.
@@ -26,12 +23,6 @@ public struct Snapshot: CustomStringConvertible {
         """
     }
 
-    /// Restores the atom values that are captured in this snapshot.
-    @MainActor
-    public func restore() {
-        _restore()
-    }
-
     /// Lookup a value associated with the given atom from the set captured in this snapshot..
     ///
     /// Note that this does not look up a overridden atom.

Tests/AtomsTests/Core/StoreContextTests.swift

@@ -284,19 +284,15 @@ final class StoreContextTests: XCTestCase {
         XCTAssertEqual(context.read(atom), 0)
     }
 
-    func testSnapshot() {
+    func testSnapshotAndRestore() {
         let store = AtomStore()
         let token = SubscriptionKey.Token()
         let subscriptionKey = SubscriptionKey(token: token)
         let context = StoreContext(store)
         let atom0 = TestAtom(value: 0)
         let atom1 = TestAtom(value: 1)
-        let atom2 = TestAtom(value: 2)
-        let atom3 = TestAtom(value: 3)
         let key0 = AtomKey(atom0)
         let key1 = AtomKey(atom1)
-        let key2 = AtomKey(atom2)
-        let key3 = AtomKey(atom3)
         let graph = Graph(
             dependencies: [key0: [key1]],
             children: [key1: [key0]]
@@ -319,34 +315,6 @@ final class StoreContextTests: XCTestCase {
             snapshot.caches.mapValues { $0 as? AtomCache<TestAtom<Int>> },
             caches
         )
-
-        // Modify graph and caches to a form that atom2 & atom3 should be released.
-        store.graph = Graph(
-            dependencies: [key0: [key1, key2], key2: [key3]],
-            children: [key1: [key0], key2: [key0], key3: [key2]]
-        )
-        store.state.caches = [
-            key0: AtomCache(atom: atom0, value: 0),
-            key1: AtomCache(atom: atom1, value: 1),
-            key2: AtomCache(atom: atom2, value: 2),
-            key3: AtomCache(atom: atom3, value: 3),
-        ]
-
-        snapshot.restore()
-
-        XCTAssertEqual(store.graph, snapshot.graph)
-        XCTAssertEqual(
-            store.state.caches.mapValues { $0 as? AtomCache<TestAtom<Int>> },
-            snapshot.caches.mapValues { $0 as? AtomCache<TestAtom<Int>> }
-        )
-
-        // Remove all subscriptions so all atoms should be released.
-        store.state.subscriptions = [:]
-
-        snapshot.restore()
-
-        XCTAssertEqual(store.graph, Graph())
-        XCTAssertTrue(store.state.caches.isEmpty)
     }
 
     func testScopedOverride() async {
@@ -898,7 +866,7 @@ final class StoreContextTests: XCTestCase {
             AtomKey(atom2): [subscriptionKey: subscription2],
         ]
 
-        snapshot.restore()
+        context.restore(snapshot)
 
         // Notifies updated only for the subscriptions of the atoms that are restored.
         XCTAssertEqual(updated, [AtomKey(atom0), AtomKey(atom1)])
@@ -921,7 +889,7 @@ final class StoreContextTests: XCTestCase {
 
         // Restore with no subscriptions.
         store.state.subscriptions.removeAll()
-        snapshot.restore()
+        context.restore(snapshot)
 
         XCTAssertEqual(store.graph, Graph())
 

Tests/AtomsTests/SnapshotTests.swift

@@ -4,21 +4,6 @@ import XCTest
 
 @MainActor
 final class SnapshotTests: XCTestCase {
-    func testRestore() {
-        var isRestoreCalled = false
-        let snapshot = Snapshot(
-            graph: Graph(),
-            caches: [:],
-            subscriptions: [:]
-        ) {
-            isRestoreCalled = true
-        }
-
-        snapshot.restore()
-
-        XCTAssertTrue(isRestoreCalled)
-    }
-
     func testLookup() {
         let atom0 = TestAtom(value: 0)
         let atom1 = TestAtom(value: 1)
@@ -29,7 +14,7 @@ final class SnapshotTests: XCTestCase {
             graph: Graph(),
             caches: atomCache,
             subscriptions: [:]
-        ) {}
+        )
 
         XCTAssertEqual(snapshot.lookup(atom0), 0)
         XCTAssertNil(snapshot.lookup(atom1))
@@ -40,7 +25,7 @@ final class SnapshotTests: XCTestCase {
             graph: Graph(),
             caches: [:],
             subscriptions: [:]
-        ) {}
+        )
 
         XCTAssertEqual(snapshot.graphDescription(), "digraph {}")
     }
@@ -93,7 +78,7 @@ final class SnapshotTests: XCTestCase {
                 key3: [subscriber: subscription],
                 key4: [subscriber: subscription],
             ]
-        ) {}
+        )
 
         XCTAssertEqual(
             snapshot.graphDescription(),