From 24903df85dd19422de693151b0c17abad69df07d Mon Sep 17 00:00:00 2001
From: pjechris <pjechris@users.noreply.github.com>
Date: Wed, 11 Oct 2023 16:41:54 +0200
Subject: [PATCH] fix(store): Fixed `storeAlias` not storing relationships
 (#59)

---
 .../CohesionKit/Identity/IdentityStore.swift  | 20 ++++------
 .../CohesionKit/Storage/AliasContainer.swift  | 38 ++++++++++---------
 Tests/CohesionKitTests/IdentityMapTests.swift | 19 +++++++++-
 3 files changed, 47 insertions(+), 30 deletions(-)

diff --git a/Sources/CohesionKit/Identity/IdentityStore.swift b/Sources/CohesionKit/Identity/IdentityStore.swift
index ca82f8b..a1772cf 100644
--- a/Sources/CohesionKit/Identity/IdentityStore.swift
+++ b/Sources/CohesionKit/Identity/IdentityStore.swift
@@ -147,8 +147,8 @@ public class IdentityMap {
         }
     }
 
-    func nodeStore<T: Identifiable>(entity: T, modifiedAt: Stamp?) -> EntityNode<T> {
-        let node = storage[entity, new: EntityNode(entity, modifiedAt: nil) { [registry] in
+    func nodeStore<T: Identifiable>(in node: EntityNode<T>? = nil, entity: T, modifiedAt: Stamp?) -> EntityNode<T> {
+        let node = node ?? storage[entity, new: EntityNode(entity, modifiedAt: nil) { [registry] in
             registry.enqueueChange(for: $0)
         }]
 
@@ -167,8 +167,8 @@ public class IdentityMap {
         return node
     }
 
-    func nodeStore<T: Aggregate>(entity: T, modifiedAt: Stamp?) -> EntityNode<T> {
-        let node = storage[entity, new: EntityNode(entity, modifiedAt: nil) { [registry] in
+    func nodeStore<T: Aggregate>(in node: EntityNode<T>? = nil, entity: T, modifiedAt: Stamp?) -> EntityNode<T> {
+        let node = node ?? storage[entity, new: EntityNode(entity, modifiedAt: nil) { [registry] in
             registry.enqueueChange(for: $0)
         }]
 
@@ -196,16 +196,12 @@ public class IdentityMap {
 
     private func storeAlias<T>(content: T, key: AliasKey<T>, modifiedAt: Stamp?) {
         let aliasNode = refAliases[safe: key]
+        let aliasContainer = AliasContainer(key: key, content: content)
 
-        do {
-            try aliasNode.updateEntity(AliasContainer(key: key, content: content), modifiedAt: modifiedAt)
+        _ = nodeStore(in: aliasNode, entity: aliasContainer, modifiedAt: modifiedAt)
 
-            registry.enqueueChange(for: aliasNode)
-            logger?.didRegisterAlias(key)
-        }
-        catch {
-
-        }
+        registry.enqueueChange(for: aliasNode)
+        logger?.didRegisterAlias(key)
     }
 
     private func transaction<T>(_ body: () -> T) -> T {
diff --git a/Sources/CohesionKit/Storage/AliasContainer.swift b/Sources/CohesionKit/Storage/AliasContainer.swift
index ace0957..9761503 100644
--- a/Sources/CohesionKit/Storage/AliasContainer.swift
+++ b/Sources/CohesionKit/Storage/AliasContainer.swift
@@ -8,32 +8,36 @@ struct AliasContainer<T>: Identifiable, Aggregate {
     var content: T?
 }
 
-extension AliasContainer where T: Aggregate {
-    var nestedEntitiesKeyPaths: [PartialIdentifiableKeyPath<AliasContainer<T>>] {
-     [.init(\.content)]
+extension AliasContainer {
+    var nestedEntitiesKeyPaths: [PartialIdentifiableKeyPath<Self>] {
+        if let self = self as? IdentifiableKeyPathsEraser {
+            return self.erasedEntitiesKeyPaths as! [PartialIdentifiableKeyPath<Self>]
+        }
+
+        if let self = self as? CollectionIdentifiableKeyPathsEraser {
+            return self.erasedEntitiesKeyPaths as! [PartialIdentifiableKeyPath<Self>]
+        }
+
+        return []
     }
 }
 
-extension AliasContainer where T: Identifiable {
-    var nestedEntitiesKeyPaths: [PartialIdentifiableKeyPath<AliasContainer<T>>] {
-     [.init(\.content)]
-    }
+private protocol IdentifiableKeyPathsEraser {
+    var erasedEntitiesKeyPaths: [Any] { get }
 }
 
-extension AliasContainer where T: MutableCollection, T.Element: Aggregate, T.Index: Hashable {
-    var nestedEntitiesKeyPaths: [PartialIdentifiableKeyPath<AliasContainer<T>>] {
-     [.init(\.content)]
+extension AliasContainer: IdentifiableKeyPathsEraser where T: Identifiable {
+    var erasedEntitiesKeyPaths: [Any] {
+        [PartialIdentifiableKeyPath<Self>(\.content)]
     }
 }
 
-extension AliasContainer where T: MutableCollection, T.Element: Identifiable, T.Index: Hashable {
-    var nestedEntitiesKeyPaths: [PartialIdentifiableKeyPath<AliasContainer<T>>] {
-     [.init(\.content)]
-    }
+private protocol CollectionIdentifiableKeyPathsEraser {
+    var erasedEntitiesKeyPaths: [Any] { get }
 }
 
-extension AliasContainer {
-    var nestedEntitiesKeyPaths: [PartialIdentifiableKeyPath<AliasContainer<T>>] {
-     []
+extension AliasContainer: CollectionIdentifiableKeyPathsEraser where T: MutableCollection, T.Element: Identifiable, T.Index: Hashable {
+    var erasedEntitiesKeyPaths: [Any] {
+        [PartialIdentifiableKeyPath<Self>(\.content)]
     }
 }
\ No newline at end of file
diff --git a/Tests/CohesionKitTests/IdentityMapTests.swift b/Tests/CohesionKitTests/IdentityMapTests.swift
index 7c4afec..f8e0a64 100644
--- a/Tests/CohesionKitTests/IdentityMapTests.swift
+++ b/Tests/CohesionKitTests/IdentityMapTests.swift
@@ -141,7 +141,7 @@ class IdentityMapTests: XCTestCase {
         wait(for: [expectation], timeout: 0)
     }
 
-    func test_storeAlias_itEnqueuesAliasInRegistry() {
+    func test_storeAggregate_named_itEnqueuesAliasInRegistry() {
         let root = SingleNodeFixture(id: 1)
         let registry = ObserverRegistryStub()
         let identityMap = IdentityMap(registry: registry)
@@ -218,6 +218,23 @@ extension IdentityMapTests {
         }
     }
 
+    func test_find_storedByAliasCollection_itReturnsEntity() {
+        let identityMap = IdentityMap()
+
+        _ = identityMap.store(entities: [SingleNodeFixture(id: 1)], named: .listOfNodes)
+
+        XCTAssertNotNil(identityMap.find(SingleNodeFixture.self, id: 1))
+    }
+
+    func test_find_storedByAliasAggregate_itReturnsEntity() {
+        let identityMap = IdentityMap()
+        let aggregate = RootFixture(id: 1, primitive: "", singleNode: SingleNodeFixture(id: 1), listNodes: [])
+
+        _ = identityMap.store(entity: aggregate, named: .root)
+
+        XCTAssertNotNil(identityMap.find(SingleNodeFixture.self, id: 1))
+    }
+
     func test_findNamed_entityStored_noObserver_returnValue() {
         let identityMap = IdentityMap()
         let entity = SingleNodeFixture(id: 1)