Store childCancellables inside ViewModelCancellable

Author: rickclephas

KMPObservableViewModelCore/ChildViewModels.swift

@@ -9,19 +9,11 @@ import KMPObservableViewModelCoreObjC
 
 public extension ViewModel {
 
-    private func setChildViewModelPublishers(_ keyPath: AnyKeyPath, _ publishers: AnyHashable?) {
-        if let publishers = publishers {
-            observableViewModelPublishers(for: self).childPublishers[keyPath] = publishers
-        } else {
-            observableViewModelPublishers(for: self).childPublishers.removeValue(forKey: keyPath)
-        }
-    }
-    
     private func setChildViewModel<VM: ViewModel>(
         _ viewModel: VM?,
         at keyPath: AnyKeyPath
     ) {
-        setChildViewModelPublishers(keyPath, observableViewModelPublishers(for: viewModel))
+        viewModelWillChange.cancellable.setChildCancellables(keyPath, ViewModelCancellable.get(for: viewModel))
     }
 
     /// Stores a reference to the `ObservableObject` for the specified child `ViewModel`.
@@ -48,8 +40,8 @@ public extension ViewModel {
         _ viewModels: [VM?]?,
         at keyPath: AnyKeyPath
     ) {
-        setChildViewModelPublishers(keyPath, viewModels?.map { viewModel in
-            observableViewModelPublishers(for: viewModel)
+        viewModelWillChange.cancellable.setChildCancellables(keyPath, viewModels?.map { viewModel in
+            ViewModelCancellable.get(for: viewModel)
         })
     }
 
@@ -95,8 +87,8 @@ public extension ViewModel {
         _ viewModels: Set<VM?>?,
         at keyPath: AnyKeyPath
     ) {
-        setChildViewModelPublishers(keyPath, viewModels?.map { viewModel in
-            observableViewModelPublishers(for: viewModel)
+        viewModelWillChange.cancellable.setChildCancellables(keyPath, viewModels?.map { viewModel in
+            ViewModelCancellable.get(for: viewModel)
         })
     }
 
@@ -142,8 +134,8 @@ public extension ViewModel {
         _ viewModels: [Key : VM?]?,
         at keyPath: AnyKeyPath
     ) {
-        setChildViewModelPublishers(keyPath, viewModels?.mapValues { viewModel in
-            observableViewModelPublishers(for: viewModel)
+        viewModelWillChange.cancellable.setChildCancellables(keyPath, viewModels?.mapValues { viewModel in
+            ViewModelCancellable.get(for: viewModel)
         })
     }
 

KMPObservableViewModelCore/ViewModelCancellable.swift

@@ -12,13 +12,15 @@ internal class ViewModelCancellable {
 
     private var didInit = false
     private weak var cancellable: AnyCancellable?
+    private var childCancellables: Dictionary<AnyKeyPath, AnyHashable>? = [:]
 
     private func get(_ viewModel: any ViewModel) -> AnyCancellable {
         guard didInit else {
             let cancellable = AnyCancellable {
                 if let cancellable = viewModel as? Cancellable {
                     cancellable.cancel()
                 }
+                self.childCancellables = nil
                 viewModel.clear()
             }
             self.cancellable = cancellable
@@ -31,7 +33,21 @@ internal class ViewModelCancellable {
         return cancellable
     }
 
+    func setChildCancellables(_ keyPath: AnyKeyPath, _ cancellables: AnyHashable?) {
+        if let cancellables = cancellables {
+            childCancellables?[keyPath] = cancellables
+        } else {
+            childCancellables?.removeValue(forKey: keyPath)
+        }
+    }
+    
     static func get(for viewModel: any ViewModel) -> AnyCancellable {
         viewModel.viewModelWillChange.cancellable.get(viewModel)
     }
+    
+    static func get(for viewModel: (any ViewModel)?) -> AnyCancellable? {
+        guard let viewModel else { return nil }
+        let cancellable = get(for: viewModel)
+        return cancellable
+    }
 }