Fix bug that KeepAlive is not effective

Sources/Atoms/Core/StoreContext.swift

@@ -424,7 +424,7 @@ private extension StoreContext {
         //     1. It's not marked as `KeepAlive`.
         //     2. It has no downstream atoms.
         //     3. It has no subscriptions from views.
-        let shouldKeepAlive = store.state.caches[key].map { $0 is any KeepAlive } ?? false
+        let shouldKeepAlive = store.state.caches[key].map { $0.atom is any KeepAlive } ?? false
         let isChildrenEmpty = store.graph.children[key]?.isEmpty ?? true
         let isSubscriptionEmpty = store.state.subscriptions[key]?.isEmpty ?? true
         let shouldRelease = !shouldKeepAlive && isChildrenEmpty && isSubscriptionEmpty

Tests/AtomsTests/Core/StoreContextTests.swift

@@ -286,6 +286,44 @@ final class StoreContextTests: XCTestCase {
         XCTAssertTrue(state?.coordinator === newState?.coordinator)
     }
 
+    func testRelease() {
+        struct KeepAliveAtom<T: Hashable>: ValueAtom, KeepAlive, Hashable {
+            var value: T
+
+            func value(context: Context) -> T {
+                value
+            }
+        }
+
+        let store = AtomStore()
+        let context = StoreContext(store)
+
+        XCTContext.runActivity(named: "Normal atoms should be released.") { _ in
+            let atom = TestAtom(value: 0)
+            let key = AtomKey(atom)
+            let transaction = Transaction(key: key) {}
+
+            _ = context.watch(atom, in: transaction)
+            XCTAssertNotNil(store.state.caches[key])
+
+            context.reset(atom)
+            XCTAssertNil(store.state.caches[key])
+        }
+
+        XCTContext.runActivity(named: "KeepAlive atoms should not be released.") { _ in
+            let atom = KeepAliveAtom(value: 0)
+            let key = AtomKey(atom)
+            let transaction = Transaction(key: key) {}
+
+            _ = context.watch(atom, in: transaction)
+
+            XCTAssertNotNil(store.state.caches[key])
+
+            context.reset(atom)
+            XCTAssertNotNil(store.state.caches[key])
+        }
+    }
+
     func testObservers() {
         let store = AtomStore()
         let container = SubscriptionContainer()