tech(entity): Introduce metadata system (#72)

Author: pjechris

Sources/CohesionKit/EntityStore.swift

@@ -162,12 +162,15 @@ public class EntityStore {
 
         do {
             try node.updateEntity(entity, modifiedAt: modifiedAt)
+            registry.enqueueChange(for: node)
             logger?.didStore(T.self, id: entity.id)
         }
         catch {
             logger?.didFailedToStore(T.self, id: entity.id, error: error)
         }
 
+        updateParents(of: node)
+
         return node
     }
 
@@ -180,6 +183,14 @@ public class EntityStore {
             return node
         }
 
+        for (childRef, _) in node.metadata.childrenRefs {
+            guard let childNode = storage[childRef]?.unwrap() as? any AnyEntityNode else {
+                continue
+            }
+
+            childNode.removeParent(node)
+        }
+
         // clear all children to avoid a removed child to be kept as child
         node.removeAllChildren()
 
@@ -191,15 +202,30 @@ public class EntityStore {
 
         do {
             try node.updateEntity(entity, modifiedAt: modifiedAt)
+            registry.enqueueChange(for: node)
             logger?.didStore(T.self, id: entity.id)
         }
         catch {
             logger?.didFailedToStore(T.self, id: entity.id, error: error)
         }
 
+        updateParents(of: node)
+
         return node
     }
 
+    func updateParents(of node: some AnyEntityNode) {
+        for parentRef in node.metadata.parentsRefs {
+            guard let parentNode = storage[parentRef]?.unwrap() as? any AnyEntityNode ?? refAliases[parentRef] else {
+                continue
+            }
+
+            parentNode.updateEntityRelationship(node)
+            parentNode.enqueue(in: registry)
+            updateParents(of: parentNode)
+        }
+    }
+
     private func storeAlias<T>(content: T?, key: AliasKey<T>, modifiedAt: Stamp?) {
         let aliasNode = refAliases[safe: key, onChange: registry.enqueueChange(for:)]
         let aliasContainer = AliasContainer(key: key, content: content)
@@ -386,7 +412,9 @@ extension EntityStore {
 
     private func removeAliases() {
         for (_, node) in refAliases {
-                node.nullify()
+            if node.nullify() {
+                node.enqueue(in: registry)
+            }
         }
     }
 }

Sources/CohesionKit/Storage/AliasStorage.swift

@@ -1,5 +1,5 @@
 /// Keep a strong reference on each aliased node
-typealias AliasStorage = [String: AnyEntityNode]
+typealias AliasStorage = [String: any AnyEntityNode]
 
 extension AliasStorage {
     subscript<T>(_ aliasKey: AliasKey<T>) -> EntityNode<AliasContainer<T>>? {
@@ -9,7 +9,8 @@ extension AliasStorage {
 
     subscript<T>(safe key: AliasKey<T>, onChange onChange: ((EntityNode<AliasContainer<T>>) -> Void)? = nil) -> EntityNode<AliasContainer<T>> {
         mutating get {
-            self[key: key, default: EntityNode(AliasContainer(key: key), modifiedAt: nil, onChange: onChange)]
+            let storeKey = buildKey(for: T.self, key: key)
+            return self[key: key, default: EntityNode(AliasContainer(key: key), key: storeKey, modifiedAt: nil, onChange: onChange)]
         }
     }
 

Sources/CohesionKit/Storage/EntitiesStorage.swift

@@ -19,6 +19,11 @@ struct EntitiesStorage {
         set { indexes[key(for: T.self, id: id)] = Weak(value: newValue) }
     }
 
+    subscript(_ key: String) -> AnyWeak? {
+        get { indexes[key] }
+        set { indexes[key] = newValue }
+    }
+
     private func key<T>(for type: T.Type, id: Any) -> String {
         "\(type)-\(id)"
     }

Sources/CohesionKit/Storage/EntityNode.swift

@@ -1,45 +1,85 @@
 import Foundation
 import Combine
 
+struct EntityMetadata {
+    /// children this entity is referencing/using
+    // TODO: change key to a ObjectKey
+    var childrenRefs: [String: AnyKeyPath] = [:]
+
+    /// parents referencing this entity. This means this entity should be listed inside its parents `EntityMetadata.childrenRefs` attribute
+    // TODO: Change value to ObjectKey
+    var parentsRefs: Set<String> = []
+    /// alias referencing this entity
+    var aliasesRefs: Set<String> = []
+
+    /// number of observers
+    var observersCount: Int = 0
+
+    var isActivelyUsed: Bool {
+        observersCount > 0 || !parentsRefs.isEmpty || !aliasesRefs.isEmpty
+    }
+}
+
 /// Typed erased protocol
 protocol AnyEntityNode: AnyObject {
+    associatedtype Value
+
+    var ref: Observable<Value> { get }
     var value: Any { get }
+    var metadata: EntityMetadata { get }
+    var storageKey: String { get }
 
-  func nullify()
+    func nullify() -> Bool
+    func removeParent(_ node: any AnyEntityNode)
+    func updateEntityRelationship(_ child: some AnyEntityNode)
+    func enqueue(in: ObserverRegistry)
 }
 
 /// A graph node representing a entity of type `T` and its children. Anytime one of its children is updated the node
 /// will reflect the change on its own value.
 class EntityNode<T>: AnyEntityNode {
+    typealias Value = T
     /// A child subscription used by its EntityNode parent
     struct SubscribedChild {
         /// the child subscription. Use it to unsubscribe to child upates
         let subscription: Subscription
         /// the child node value
-        let node: AnyEntityNode
+        let node: any AnyEntityNode
     }
 
     var value: Any { ref.value }
 
+    var metadata = EntityMetadata()
+    // FIXME: to delete, it's "just" to have a strong ref and avoid nodes to be deleted. Need a better memory management
+    private var childrenNodes: [any AnyEntityNode] = []
+
     var applyChildrenChanges = true
     /// An observable entity reference
     let ref: Observable<T>
 
+    let storageKey: String
+
     private let onChange: ((EntityNode<T>) -> Void)?
     /// last time the ref.value was changed. Any subsequent change must have a higher value to be applied
     /// if nil ref has no stamp and any change will be accepted
     private var modifiedAt: Stamp?
     /// entity children
     private(set) var children: [PartialKeyPath<T>: SubscribedChild] = [:]
 
-    init(ref: Observable<T>, modifiedAt: Stamp?, onChange: ((EntityNode<T>) -> Void)? = nil) {
+    init(ref: Observable<T>, key: String, modifiedAt: Stamp?, onChange: ((EntityNode<T>) -> Void)? = nil) {
         self.ref = ref
         self.modifiedAt = modifiedAt
         self.onChange = onChange
+        self.storageKey = key
+    }
+
+    convenience init(_ entity: T, key: String, modifiedAt: Stamp?, onChange: ((EntityNode<T>) -> Void)? = nil) {
+        self.init(ref: Observable(value: entity), key: key, modifiedAt: modifiedAt, onChange: onChange)
     }
 
-    convenience init(_ entity: T, modifiedAt: Stamp?, onChange: ((EntityNode<T>) -> Void)? = nil) {
-        self.init(ref: Observable(value: entity), modifiedAt: modifiedAt, onChange: onChange)
+    convenience init(_ entity: T, modifiedAt: Stamp?, onChange: ((EntityNode<T>) -> Void)? = nil) where T: Identifiable {
+        let key = "\(T.self)-\(entity.id)"
+        self.init(entity, key: key, modifiedAt: modifiedAt, onChange: onChange)
     }
 
     /// change the entity to a new value. If modifiedAt is nil or > to previous date update the value will be changed
@@ -52,17 +92,56 @@ class EntityNode<T>: AnyEntityNode {
 
         modifiedAt = newModifiedAt ?? modifiedAt
         ref.value = newEntity
-        onChange?(self)
     }
 
-    func nullify() {
+    func nullify() -> Bool {
         if let value = ref.value as? Nullable {
-            try? updateEntity(value.nullified() as! T, modifiedAt: nil)
+            do {
+                try updateEntity(value.nullified() as! T, modifiedAt: nil)
+                return true
+            }
+            catch {
+                return false
+            }
         }
+
+        return false
     }
 
     func removeAllChildren() {
         children = [:]
+        metadata.childrenRefs = [:]
+        childrenNodes = []
+    }
+
+    func removeParent(_ node: any AnyEntityNode) {
+        metadata.parentsRefs.remove(node.storageKey)
+    }
+
+    func updateEntityRelationship<U: AnyEntityNode>(_ child: U) {
+        guard applyChildrenChanges else {
+            return
+        }
+
+        guard let keyPath = metadata.childrenRefs[child.storageKey] else {
+            return
+        }
+
+        if let writableKeyPath = keyPath as? WritableKeyPath<T, U.Value> {
+            ref.value[keyPath: writableKeyPath] = child.ref.value
+            return
+        }
+
+        if let optionalWritableKeyPath = keyPath as? WritableKeyPath<T, U.Value?> {
+            ref.value[keyPath: optionalWritableKeyPath] = child.ref.value
+            return
+        }
+
+        print("CohesionKit: cannot convert \(type(of: keyPath)) to WritableKeyPath<\(T.self), \(U.Value.self)>")
+    }
+
+    func enqueue(in registry: ObserverRegistry) {
+        registry.enqueueChange(for: self)
     }
 
     /// observe one of the node child
@@ -88,20 +167,9 @@ class EntityNode<T>: AnyEntityNode {
         identity keyPath: KeyPath<T, C>,
         update: @escaping (inout T, Element) -> Void
     ) {
-        if let subscribedChild = children[keyPath]?.node as? EntityNode<Element>, subscribedChild == childNode {
-            return
-        }
-
-        let subscription = childNode.ref.addObserver { [unowned self] newValue in
-            guard self.applyChildrenChanges else {
-                return
-            }
-
-            update(&self.ref.value, newValue)
-            self.onChange?(self)
-        }
-
-        children[keyPath] = SubscribedChild(subscription: subscription, node: childNode)
+        metadata.childrenRefs[childNode.storageKey] = keyPath
+        childNode.metadata.parentsRefs.insert(storageKey)
+        childrenNodes.append(childNode)
     }
 }
 

Tests/CohesionKitTests/EntityStoreTests.swift

@@ -411,7 +411,23 @@ extension EntityStoreTests {
         XCTAssertTrue(registry.hasPendingChange(for: AliasContainer<RootFixture>.self))
     }
 
-    func test_update_entityIsIndirectlyUsedByAlias_itEnqueuesAliasInRegistry() {
+    // make sure that when we have A -> B -> C and update C, we enqueue parents B AND A.
+    func test_update_entityIsNested_itEnqueuesAllParents() {
+        let a = AFixture(b: BFixture(c: SingleNodeFixture(id: 1)))
+        let registry = ObserverRegistryStub()
+        let entityStore = EntityStore(registry: registry)
+
+        withExtendedLifetime(entityStore.store(entity: a)) {
+            registry.clearPendingChangesStub()
+
+            _ = entityStore.nodeStore(entity: SingleNodeFixture(id: 1, primitive: "updated"), modifiedAt: nil)
+        }
+
+        XCTAssertTrue(registry.hasPendingChange(for: BFixture.self))
+        XCTAssertTrue(registry.hasPendingChange(for: AFixture.self))
+    }
+
+    func test_update_entityIsInsideAggregagte_aggreateIsAliased_itEnqueuesAliasInRegistry() {
         let aggregate = RootFixture(id: 1, primitive: "", singleNode: SingleNodeFixture(id: 1), listNodes: [])
         let registry = ObserverRegistryStub()
         let entityStore = EntityStore(registry: registry)

Tests/CohesionKitTests/RootFixture.swift

@@ -1,6 +1,24 @@
 import Foundation
 import CohesionKit
 
+struct AFixture: Aggregate {
+    var id: BFixture.ID { b.id }
+    var b: BFixture
+
+    var nestedEntitiesKeyPaths: [PartialIdentifiableKeyPath<Self>] {
+        [.init(\.b)]
+    }
+}
+
+struct BFixture: Aggregate {
+    var id: SingleNodeFixture.ID { c.id }
+    var c: SingleNodeFixture
+
+    var nestedEntitiesKeyPaths: [PartialIdentifiableKeyPath<Self>] {
+        [.init(\.c)]
+    }
+}
+
 struct RootFixture: Aggregate, Equatable {
     let id: Int
     let primitive: String
@@ -56,4 +74,4 @@ struct ListNodeFixture: Identifiable, Equatable {
             PartialIdentifiableKeyPath(parent.appending(path: \.singleNode))
         ]
     }
-}
+}

Tests/CohesionKitTests/Storage/EntityNodeTests.swift

@@ -68,47 +68,45 @@ class EntityNodeTests: XCTestCase {
         }
     }
 
-    func test_observeChild_childChange_entityIsUpdated() throws {
+    func test_observeChild_nodeIsAddedAsParentMetadata() {
         let childNode = EntityNode(startEntity.singleNode, modifiedAt: nil)
-        let newChild = SingleNodeFixture(id: 1, primitive: "updated")
 
         node.observeChild(childNode, for: \.singleNode)
 
-        try childNode.updateEntity(newChild, modifiedAt: nil)
+        XCTAssertTrue(childNode.metadata.parentsRefs.contains(node.storageKey))
+    }
+
+    func test_observeChild_childrenMetadataIsUpdated() {
+        let childNode = EntityNode(startEntity.singleNode, modifiedAt: nil)
 
-        XCTAssertEqual((node.value as? RootFixture)?.singleNode, newChild)
+        node.observeChild(childNode, for: \.singleNode)
+
+        XCTAssertTrue(node.metadata.childrenRefs.keys.contains(childNode.storageKey))
     }
 
-    func test_observeChild_childChange_entityObserversAreCalled() throws {
+    func test_updateEntityRelationship_childIsUpdated() throws {
         let childNode = EntityNode(startEntity.singleNode, modifiedAt: startTimestamp)
         let newChild = SingleNodeFixture(id: 1, primitive: "updated")
-        let entityRef = Observable(value: startEntity)
-        var observerCalled = false
-
-        let subscription = entityRef.addObserver { _ in
-            observerCalled = true
-        }
 
-        node = EntityNode(ref: entityRef, modifiedAt: startTimestamp)
         node.observeChild(childNode, for: \.singleNode)
 
         try childNode.updateEntity(newChild, modifiedAt: nil)
 
-        subscription.unsubscribe()
+        node.updateEntityRelationship(childNode)
 
-        XCTAssertTrue(observerCalled)
+        XCTAssertEqual(node.ref.value.singleNode, newChild)
     }
 
     func test_observeChild_childIsCollection_eachChildIsAdded() {
         let child1 = EntityNode(ListNodeFixture(id: 1), modifiedAt: startTimestamp)
         let child2 = EntityNode(ListNodeFixture(id: 2), modifiedAt: startTimestamp)
         let node = EntityNode(startEntity, modifiedAt: startTimestamp)
 
-        XCTAssertEqual(node.children.count, 0)
+        XCTAssertEqual(node.metadata.childrenRefs.count, 0)
 
         node.observeChild(child1, for: \.listNodes[0])
         node.observeChild(child2, for: \.listNodes[1])
 
-        XCTAssertEqual(node.children.count, 2)
+        XCTAssertEqual(node.metadata.childrenRefs.count, 2)
     }
 }