Update ObservationRegistrar logic and implement usage in ObservableViewModelPublisher

Author: rickclephas

KMMViewModel.xcodeproj/project.pbxproj

@@ -18,7 +18,7 @@
 		1D198B312933C01800EF778D /* KMMViewModelCoreObjC.m in Sources */ = {isa = PBXBuildFile; fileRef = 1D198B2F2933C01800EF778D /* KMMViewModelCoreObjC.m */; };
 		1D198B322933C04400EF778D /* KMMViewModelCoreObjC.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 1D198B272933BFD900EF778D /* KMMViewModelCoreObjC.framework */; };
 		1D2AAC592BB1F528005F1344 /* Observable.swift in Sources */ = {isa = PBXBuildFile; fileRef = 1D2AAC582BB1F528005F1344 /* Observable.swift */; };
-		1D2AAC5B2BB2053D005F1344 /* ViewModelObservationRegistrar.swift in Sources */ = {isa = PBXBuildFile; fileRef = 1D2AAC5A2BB2053D005F1344 /* ViewModelObservationRegistrar.swift */; };
+		1D2AAC5B2BB2053D005F1344 /* ObservationRegistrar.swift in Sources */ = {isa = PBXBuildFile; fileRef = 1D2AAC5A2BB2053D005F1344 /* ObservationRegistrar.swift */; };
 		1D2AAC5D2BB205AC005F1344 /* ObservableProperty.swift in Sources */ = {isa = PBXBuildFile; fileRef = 1D2AAC5C2BB205AC005F1344 /* ObservableProperty.swift */; };
 		1D43F3EC2ABAFCA600EB3DFE /* ObservableViewModelPublisher.swift in Sources */ = {isa = PBXBuildFile; fileRef = 1D43F3EB2ABAFCA600EB3DFE /* ObservableViewModelPublisher.swift */; };
 		1D43F3EE2ABAFD7D00EB3DFE /* ObservableViewModelPublishers.swift in Sources */ = {isa = PBXBuildFile; fileRef = 1D43F3ED2ABAFD7D00EB3DFE /* ObservableViewModelPublishers.swift */; };
@@ -56,7 +56,7 @@
 		1D198B2E2933C01800EF778D /* KMMVMViewModelScope.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = KMMVMViewModelScope.h; sourceTree = "<group>"; };
 		1D198B2F2933C01800EF778D /* KMMViewModelCoreObjC.m */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.objc; path = KMMViewModelCoreObjC.m; sourceTree = "<group>"; };
 		1D2AAC582BB1F528005F1344 /* Observable.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Observable.swift; sourceTree = "<group>"; };
-		1D2AAC5A2BB2053D005F1344 /* ViewModelObservationRegistrar.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ViewModelObservationRegistrar.swift; sourceTree = "<group>"; };
+		1D2AAC5A2BB2053D005F1344 /* ObservationRegistrar.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ObservationRegistrar.swift; sourceTree = "<group>"; };
 		1D2AAC5C2BB205AC005F1344 /* ObservableProperty.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ObservableProperty.swift; sourceTree = "<group>"; };
 		1D43F3EB2ABAFCA600EB3DFE /* ObservableViewModelPublisher.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ObservableViewModelPublisher.swift; sourceTree = "<group>"; };
 		1D43F3ED2ABAFD7D00EB3DFE /* ObservableViewModelPublishers.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ObservableViewModelPublishers.swift; sourceTree = "<group>"; };
@@ -117,12 +117,12 @@
 			children = (
 				1D6641DC2A5175C3000180D7 /* ChildViewModels.swift */,
 				1DDAF21D293545DD0049C114 /* KMMViewModel.swift */,
+				1D2AAC582BB1F528005F1344 /* Observable.swift */,
+				1D2AAC5C2BB205AC005F1344 /* ObservableProperty.swift */,
 				1D0DA80129336AD40057DDAD /* ObservableViewModel.swift */,
 				1D43F3EB2ABAFCA600EB3DFE /* ObservableViewModelPublisher.swift */,
 				1D43F3ED2ABAFD7D00EB3DFE /* ObservableViewModelPublishers.swift */,
-				1D2AAC582BB1F528005F1344 /* Observable.swift */,
-				1D2AAC5A2BB2053D005F1344 /* ViewModelObservationRegistrar.swift */,
-				1D2AAC5C2BB205AC005F1344 /* ObservableProperty.swift */,
+				1D2AAC5A2BB2053D005F1344 /* ObservationRegistrar.swift */,
 			);
 			path = KMMViewModelCore;
 			sourceTree = "<group>";
@@ -312,7 +312,7 @@
 			buildActionMask = 2147483647;
 			files = (
 				1D43F3EE2ABAFD7D00EB3DFE /* ObservableViewModelPublishers.swift in Sources */,
-				1D2AAC5B2BB2053D005F1344 /* ViewModelObservationRegistrar.swift in Sources */,
+				1D2AAC5B2BB2053D005F1344 /* ObservationRegistrar.swift in Sources */,
 				1D43F3EC2ABAFCA600EB3DFE /* ObservableViewModelPublisher.swift in Sources */,
 				1DDAF21E293545DD0049C114 /* KMMViewModel.swift in Sources */,
 				1D6641DD2A5175C3000180D7 /* ChildViewModels.swift in Sources */,

KMMViewModelCore/Observable.swift

@@ -10,9 +10,7 @@ import Observation
 import KMMViewModelCoreObjC
 
 @available(iOS 17.0, macOS 14.0, tvOS 17.0, watchOS 10.0, *)
-public protocol Observable: Observation.Observable {
-    associatedtype ViewModel: Observable
-    
-    @ViewModelObservationRegistrarBuilder<ViewModel>
-    var viewModelObservationRegistrar: ViewModelObservationRegistrar<ViewModel> { get }
+public protocol Observable: KMMViewModel, Observation.Observable {
+    @ObservationRegistrarBuilder<Self>
+    var observationRegistrar: ObservationRegistrar { get }
 }

KMMViewModelCore/ObservableProperty.swift

@@ -8,29 +8,40 @@
 import Observation
 
 @available(iOS 17.0, macOS 14.0, tvOS 17.0, watchOS 10.0, *)
-public protocol ObservableProperty<ViewModel> {
-    associatedtype ViewModel: Observable
+public protocol ObservableProperty {
 
-    func willSet(registrar: ObservationRegistrar, viewModel: ViewModel)
+    func initialize(subject: any Observable)
 
-    func didSet(registrar: ObservationRegistrar, viewModel: ViewModel)
+    func willSet(registrar: Observation.ObservationRegistrar, subject: any Observable)
 
-    func access(registrar: ObservationRegistrar, viewModel: ViewModel)
+    func didSet(registrar: Observation.ObservationRegistrar, subject: any Observable)
+    
+    func access(registrar: Observation.ObservationRegistrar, subject: any Observable)
 }
 
 @available(iOS 17.0, macOS 14.0, tvOS 17.0, watchOS 10.0, *)
 extension KeyPath: ObservableProperty where Root: Observable {
-    public typealias ViewModel = Root
 
-    public func willSet(registrar: ObservationRegistrar, viewModel: ViewModel) {
-        registrar.willSet(viewModel, keyPath: self)
+    private func root(_ subject: any Observable) -> Root {
+        guard let root = subject as? Root else {
+            fatalError("subject must be of type Root")
+        }
+        return root
+    }
+    
+    public func initialize(subject: any Observable) {
+        _  = root(subject)[keyPath: self]
+    }
+    
+    public func willSet(registrar: Observation.ObservationRegistrar, subject: any Observable) {
+        registrar.willSet(root(subject), keyPath: self)
     }
 
-    public func didSet(registrar: ObservationRegistrar, viewModel: ViewModel) {
-        registrar.didSet(viewModel, keyPath: self)
+    public func didSet(registrar: Observation.ObservationRegistrar, subject: any Observable) {
+        registrar.didSet(root(subject), keyPath: self)
     }
 
-    public func access(registrar: ObservationRegistrar, viewModel: ViewModel) {
-        registrar.access(viewModel, keyPath: self)
+    public func access(registrar: Observation.ObservationRegistrar, subject: any Observable) {
+        registrar.access(root(subject), keyPath: self)
     }
 }

KMMViewModelCore/ObservableViewModelPublisher.swift

@@ -20,8 +20,12 @@ public final class ObservableViewModelPublisher: Publisher {
 
     internal init(_ viewModel: any KMMViewModel, _ objectWillChange: ObservableObjectPublisher) {
         self.viewModel = viewModel
-        viewModel.viewModelScope.setSendObjectWillChange { [weak self] in
-            self?.publisher.send()
+        if #available(iOS 17.0, macOS 14.0, tvOS 17.0, watchOS 10.0, *), let observable = viewModel as? (any Observable) {
+            observable.observationRegistrar.initialize(observable, publisher)
+        } else {
+            viewModel.viewModelScope.setPropertyWillSet { [weak self] _ in
+                self?.publisher.send()
+            }
         }
         objectWillChangeCancellable = objectWillChange.sink { [weak self] _ in
             self?.publisher.send()

KMMViewModelCore/ObservationRegistrar.swift

@@ -0,0 +1,84 @@
+//
+//  ObservationRegistrar.swift
+//  KMMViewModelCore
+//
+//  Created by Rick Clephas on 25/03/2024.
+//
+
+import Foundation
+import Observation
+import Combine
+
+@available(iOS 17.0, macOS 14.0, tvOS 17.0, watchOS 10.0, *)
+public class ObservationRegistrar {
+    
+    private let registrar = Observation.ObservationRegistrar()
+    
+    private let observableProperties: [ObservableProperty]
+    private var properties: [NSObject:[ObservableProperty]] = [:]
+    
+    internal init(_ observableProperties: [ObservableProperty]) {
+        self.observableProperties = observableProperties
+    }
+    
+    private weak var observable: (any Observable)? = nil
+    private var publisher: ObservableObjectPublisher? = nil
+    
+    internal func initialize(_ observable: any Observable, _ publisher: ObservableObjectPublisher) {
+        observable.viewModelScope.setPropertyAccess(access)
+        observable.viewModelScope.setPropertyWillSet(willSet)
+        observable.viewModelScope.setPropertyDidSet(didSet)
+        for observableProperty in observableProperties {
+            Thread.current.threadDictionary["observableProperty"] = observableProperty
+            observableProperty.initialize(subject: observable)
+        }
+        Thread.current.threadDictionary.removeObject(forKey: "observableProperty")
+        self.observable = observable
+        self.publisher = publisher
+    }
+    
+    private func access(_ property: NSObject) {
+        guard let observable else {
+            if let observableProperty = Thread.current.threadDictionary["observableProperty"] as? ObservableProperty {
+                var observableProperties = properties[property] ?? []
+                observableProperties.append(observableProperty)
+                properties[property] = observableProperties
+            }
+            return
+        }
+        guard let properties = properties[property] else { return }
+        for property in properties {
+            property.access(registrar: registrar, subject: observable)
+        }
+    }
+    
+    private func willSet(_ property: NSObject) {
+        guard let observable, let properties = properties[property] else {
+            publisher?.send()
+            return
+        }
+        for property in properties {
+            property.willSet(registrar: registrar, subject: observable)
+        }
+    }
+    
+    private func didSet(_ property: NSObject) {
+        guard let observable, let properties = properties[property] else { return }
+        for property in properties {
+            property.didSet(registrar: registrar, subject: observable)
+        }
+    }
+}
+
+@resultBuilder
+@available(iOS 17.0, macOS 14.0, tvOS 17.0, watchOS 10.0, *)
+public struct ObservationRegistrarBuilder<ViewModel: Observable> {
+    
+    public static func buildExpression<Member>(_ expression: KeyPath<ViewModel, Member>) -> ObservableProperty {
+        return expression
+    }
+    
+    public static func buildBlock(_ components: ObservableProperty...) -> ObservationRegistrar {
+        return ObservationRegistrar(components)
+    }
+}

KMMViewModelCore/ViewModelObservationRegistrar.swift

@@ -14,7 +14,9 @@ __attribute__((swift_name("ViewModelScope")))
 @protocol KMMVMViewModelScope
 - (void)increaseSubscriptionCount;
 - (void)decreaseSubscriptionCount;
-- (void)setSendObjectWillChange:(void (^ _Nonnull)(void))sendObjectWillChange;
+- (void)setPropertyAccess:(void (^ _Nonnull)(NSObject * _Nonnull))propertyAccess;
+- (void)setPropertyWillSet:(void (^ _Nonnull)(NSObject * _Nonnull))propertyWillSet;
+- (void)setPropertyDidSet:(void (^ _Nonnull)(NSObject * _Nonnull))propertyDidSet;
 - (void)cancel;
 @end
 

KMMViewModelCoreObjC/KMMVMViewModelScope.h

@@ -14,7 +14,8 @@ public actual fun <T> MutableStateFlow(
 ): MutableStateFlow<T> = MutableStateFlowImpl(viewModelScope.asImpl(), MutableStateFlow(value))
 
 /**
- * A [MutableStateFlow] that triggers [ViewModelScopeImpl.sendObjectWillChange]
+ * A [MutableStateFlow] that triggers [ViewModelScopeImpl.propertyWillSet],
+ * [ViewModelScopeImpl.propertyDidSet] and [ViewModelScopeImpl.propertyAccess]
  * and accounts for the [ViewModelScopeImpl.subscriptionCount].
  */
 private class MutableStateFlowImpl<T>(
@@ -23,16 +24,21 @@ private class MutableStateFlowImpl<T>(
 ): MutableStateFlow<T> {
 
     override var value: T
-        get() = stateFlow.value
+        get() {
+            viewModelScope.propertyAccess(this)
+            return stateFlow.value
+        }
         set(value) {
-            if (stateFlow.value != value) {
-                viewModelScope.sendObjectWillChange()
-            }
+            val changed = stateFlow.value != value
+            if (changed) viewModelScope.propertyWillSet(this)
             stateFlow.value = value
+            if (changed) viewModelScope.propertyDidSet(this)
         }
 
-    override val replayCache: List<T>
-        get() = stateFlow.replayCache
+    override val replayCache: List<T> get() {
+        viewModelScope.propertyAccess(this)
+        return stateFlow.replayCache
+    }
 
     override val subscriptionCount: StateFlow<Int> =
         SubscriptionCountFlow(viewModelScope.subscriptionCount, stateFlow.subscriptionCount)
@@ -41,10 +47,11 @@ private class MutableStateFlowImpl<T>(
         stateFlow.collect(collector)
 
     override fun compareAndSet(expect: T, update: T): Boolean {
-        if (stateFlow.value == expect && expect != update) {
-            viewModelScope.sendObjectWillChange()
-        }
-        return stateFlow.compareAndSet(expect, update)
+        val changed = stateFlow.value == expect && expect != update
+        if (changed) viewModelScope.propertyWillSet(this)
+        val result = stateFlow.compareAndSet(expect, update)
+        if (changed) viewModelScope.propertyDidSet(this)
+        return result
     }
 
     @ExperimentalCoroutinesApi

kmm-viewmodel-core/src/appleMain/kotlin/com/rickclephas/kmm/viewmodel/StateFlow.kt

@@ -59,20 +59,52 @@ public class ViewModelScopeImpl internal constructor(
         _subscriptionCount.update { it - 1 }
     }
 
-    private var sendObjectWillChange: (() -> Unit)? = null
+    private var propertyAccess: ((NSObject) -> Unit)? = null
 
-    override fun setSendObjectWillChange(sendObjectWillChange: () -> Unit) {
-        if (this.sendObjectWillChange != null) {
+    override fun setPropertyAccess(propertyAccess: (NSObject?) -> Unit) {
+        if (this.propertyAccess != null) {
             throw IllegalStateException("KMMViewModel can't be wrapped more than once")
         }
-        this.sendObjectWillChange = sendObjectWillChange
+        this.propertyAccess = propertyAccess
     }
 
     /**
-     * Invokes the object will change listener set by [setSendObjectWillChange].
+     * Invokes the listener set by [setPropertyAccess].
      */
-    public fun sendObjectWillChange() {
-        sendObjectWillChange?.invoke()
+    public fun propertyAccess(property: Any) {
+        propertyAccess?.invoke(property as NSObject)
+    }
+
+    private var propertyWillSet: ((NSObject) -> Unit)? = null
+
+    override fun setPropertyWillSet(propertyWillSet: (NSObject?) -> Unit) {
+        if (this.propertyWillSet != null) {
+            throw IllegalStateException("KMMViewModel can't be wrapped more than once")
+        }
+        this.propertyWillSet = propertyWillSet
+    }
+
+    /**
+     * Invokes the listener set by [setPropertyWillSet].
+     */
+    public fun propertyWillSet(property: Any) {
+        propertyWillSet?.invoke(property as NSObject)
+    }
+
+    private var propertyDidSet: ((NSObject) -> Unit)? = null
+
+    override fun setPropertyDidSet(propertyDidSet: (NSObject?) -> Unit) {
+        if (this.propertyDidSet != null) {
+            throw IllegalStateException("KMMViewModel can't be wrapped more than once")
+        }
+        this.propertyDidSet = propertyDidSet
+    }
+
+    /**
+     * Invokes the listener set by [setPropertyDidSet].
+     */
+    public fun propertyDidSet(property: Any) {
+        propertyDidSet?.invoke(property as NSObject)
     }
 
     override fun cancel() {