From b919c7c732c91ca00b9b1984d7d16a3b835fa422 Mon Sep 17 00:00:00 2001
From: pjechris <pjechris@users.noreply.github.com>
Date: Wed, 27 Mar 2024 14:33:04 +0100
Subject: [PATCH] [store] post nil when alias is removed

---
 Sources/CohesionKit/EntityStore.swift         | 43 ++++++++++++++++---
 .../CohesionKit/Storage/AliasContainer.swift  | 16 ++++++-
 .../CohesionKit/Storage/AliasStorage.swift    |  2 +-
 Sources/CohesionKit/Storage/EntityNode.swift  |  8 ++++
 Tests/CohesionKitTests/EntityStoreTests.swift | 31 +++++++++++++
 5 files changed, 91 insertions(+), 9 deletions(-)

diff --git a/Sources/CohesionKit/EntityStore.swift b/Sources/CohesionKit/EntityStore.swift
index ef44b40..4342961 100644
--- a/Sources/CohesionKit/EntityStore.swift
+++ b/Sources/CohesionKit/EntityStore.swift
@@ -349,25 +349,54 @@ extension EntityStore {
 extension EntityStore {
     /// Removes an alias from the storage
     public func removeAlias<T>(named: AliasKey<T>) {
-        refAliases[named] = nil
-        logger?.didUnregisterAlias(named)
+        transaction {
+            if let alias = refAliases[named] {
+                do {
+                    try alias.updateEntity(AliasContainer(key: named, content: nil), modifiedAt: nil)
+                    logger?.didUnregisterAlias(named)
+                }
+                catch {
+
+                }
+            }
+        }
     }
 
     /// Removes an alias from the storage
     public func removeAlias<C: Collection>(named: AliasKey<C>) {
-        refAliases[named] = nil
-        logger?.didUnregisterAlias(named)
+        transaction {
+            if let alias = refAliases[named] {
+                do {
+                    try alias.updateEntity(AliasContainer(key: named, content: nil), modifiedAt: nil)
+                    logger?.didUnregisterAlias(named)
+                }
+                catch {
+
+                }
+            }
+        }
+
     }
 
     /// Removes all alias from identity map
     public func removeAllAlias() {
-        refAliases.removeAll()
+        transaction {
+            removeAliases()
+        }
     }
 
     /// Removes all alias AND all objects stored weakly. You should not need this method and rather use `removeAlias`.
     /// But this can be useful if you fear retain cycles
     public func removeAll() {
-        refAliases.removeAll()
-        storage.removeAll()
+        transaction {
+            removeAliases()
+            storage.removeAll()
+        }
+    }
+
+    private func removeAliases() {
+        for (_, node) in refAliases {
+                node.nullify()
+        }
     }
 }
diff --git a/Sources/CohesionKit/Storage/AliasContainer.swift b/Sources/CohesionKit/Storage/AliasContainer.swift
index 37e0585..dfa6fef 100644
--- a/Sources/CohesionKit/Storage/AliasContainer.swift
+++ b/Sources/CohesionKit/Storage/AliasContainer.swift
@@ -68,4 +68,18 @@ extension AliasContainer: CollectionIdentifiableKeyPathsEraser where T: MutableC
     var erasedEntitiesKeyPaths: [Any] {
         [PartialIdentifiableKeyPath<Self>(\.content)]
     }
-}
\ No newline at end of file
+}
+
+extension AliasContainer: Equatable where T: Equatable {
+
+}
+
+protocol Nullable {
+  func nullified() -> Self
+}
+
+extension AliasContainer: Nullable {
+  func nullified() -> AliasContainer<T> {
+    AliasContainer(key: key, content: nil)
+  }
+}
diff --git a/Sources/CohesionKit/Storage/AliasStorage.swift b/Sources/CohesionKit/Storage/AliasStorage.swift
index 58cc9ec..a8a341e 100644
--- a/Sources/CohesionKit/Storage/AliasStorage.swift
+++ b/Sources/CohesionKit/Storage/AliasStorage.swift
@@ -1,5 +1,5 @@
 /// Keep a strong reference on each aliased node
-typealias AliasStorage = [String: Any]
+typealias AliasStorage = [String: AnyEntityNode]
 
 extension AliasStorage {
     subscript<T>(_ aliasKey: AliasKey<T>) -> EntityNode<AliasContainer<T>>? {
diff --git a/Sources/CohesionKit/Storage/EntityNode.swift b/Sources/CohesionKit/Storage/EntityNode.swift
index 9623342..ac77322 100644
--- a/Sources/CohesionKit/Storage/EntityNode.swift
+++ b/Sources/CohesionKit/Storage/EntityNode.swift
@@ -4,6 +4,8 @@ import Combine
 /// Typed erased protocol
 protocol AnyEntityNode: AnyObject {
     var value: Any { get }
+
+  func nullify()
 }
 
 /// A graph node representing a entity of type `T` and its children. Anytime one of its children is updated the node
@@ -53,6 +55,12 @@ class EntityNode<T>: AnyEntityNode {
         onChange?(self)
     }
 
+    func nullify() {
+        if let value = ref.value as? Nullable {
+            try? updateEntity(value.nullified() as! T, modifiedAt: nil)
+        }
+    }
+
     func removeAllChildren() {
         children = [:]
     }
diff --git a/Tests/CohesionKitTests/EntityStoreTests.swift b/Tests/CohesionKitTests/EntityStoreTests.swift
index f404243..cdf00c5 100644
--- a/Tests/CohesionKitTests/EntityStoreTests.swift
+++ b/Tests/CohesionKitTests/EntityStoreTests.swift
@@ -424,6 +424,37 @@ extension EntityStoreTests {
     }
 }
 
+// MARK: Remove
+extension EntityStoreTests {
+    func test_removeAlias_itEnqueuesNilInRegistry() {
+        let registry = ObserverRegistryStub()
+        let store = EntityStore(registry: registry)
+
+        _ = store.store(entity: SingleNodeFixture(id: 1), named: .test)
+
+        registry.clearPendingChangesStub()
+
+        store.removeAlias(named: .test)
+
+        XCTAssertTrue(registry.hasPendingChange(for: AliasContainer(key: .test, content: nil)))
+    }
+
+    func test_removeAllAlias_itEnqueuesNilForEachAlias() {
+        let registry = ObserverRegistryStub()
+        let store = EntityStore(registry: registry)
+
+        _ = store.store(entity: SingleNodeFixture(id: 1), named: .test)
+        _ = store.store(entities: [SingleNodeFixture(id: 2)], named: .listOfNodes)
+
+        registry.clearPendingChangesStub()
+
+        store.removeAllAlias()
+
+        XCTAssertTrue(registry.hasPendingChange(for: AliasContainer(key: .test, content: nil)))
+        XCTAssertTrue(registry.hasPendingChange(for: AliasContainer(key: .listOfNodes, content: nil)))
+    }
+}
+
 private extension AliasKey where T == SingleNodeFixture {
     static let test = AliasKey(named: "test")
 }