Ensure @provide values are removed from all @collect values when modu…

…le is unloaded

Sources/Spezi/Capabilities/Communication/CollectPropertyWrapper.swift

@@ -67,6 +67,6 @@ extension Module {
 
 extension _CollectPropertyWrapper: StorageValueCollector {
     public func retrieve<Repository: SharedRepository<SpeziAnchor>>(from repository: Repository) {
-        injectedValues = repository[CollectedModuleValue<Value>.self] ?? []
+        injectedValues = repository[CollectedModuleValues<Value>.self]?.map { $0.value } ?? []
     }
 }

Sources/Spezi/Capabilities/Communication/CollectedModuleValues.swift

@@ -0,0 +1,79 @@
+//
+// This source file is part of the Stanford Spezi open-source project
+//
+// SPDX-FileCopyrightText: 2023 Stanford University and the project authors (see CONTRIBUTORS.md)
+//
+// SPDX-License-Identifier: MIT
+//
+
+import SpeziFoundation
+
+protocol AnyCollectModuleValue {
+    associatedtype Value
+
+    var moduleReference: ModuleReference { get }
+}
+
+protocol AnyCollectModuleValues {
+    associatedtype Value
+
+    var values: [any AnyCollectModuleValue] { get }
+
+    mutating func removeValues(from module: any Module) -> Bool
+
+    func store(into storage: inout SpeziStorage)
+}
+
+
+struct CollectModuleValue<Value>: AnyCollectModuleValue {
+    let value: Value
+    let moduleReference: ModuleReference
+
+    init(_ value: Value) {
+        self.value = value
+
+        guard let module = Spezi.moduleInitContext else {
+            preconditionFailure("Tried to initialize CollectModuleValue with unknown module init context.")
+        }
+        self.moduleReference = ModuleReference(module)
+    }
+}
+
+/// Provides the ``KnowledgeSource`` for any value we store in the ``SpeziStorage`` that is
+/// provided or request from am ``Module``.
+///
+/// For more information, look at the ``Module/Provide`` or ``Module/Collect`` property wrappers.
+struct CollectedModuleValues<ModuleValue>: DefaultProvidingKnowledgeSource {
+    typealias Anchor = SpeziAnchor
+
+    typealias Value = [CollectModuleValue<ModuleValue>]
+
+
+    static var defaultValue: Value {
+        []
+    }
+}
+
+
+extension Array: AnyCollectModuleValues where Element: AnyCollectModuleValue {
+    typealias Value = Element.Value
+
+    var values: [any AnyCollectModuleValue] {
+        self
+    }
+
+    mutating func removeValues(from module: any Module) -> Bool {
+        let previousCount = count
+        removeAll { entry in
+            entry.moduleReference == ModuleReference(module)
+        }
+        return previousCount != count
+    }
+
+    func store(into storage: inout SpeziStorage) {
+        guard let values = self as? [CollectModuleValue<Value>] else {
+            preconditionFailure("Unexpected array type: \(type(of: self))")
+        }
+        storage[CollectedModuleValues<Value>.self] = values
+    }
+}

Sources/Spezi/Capabilities/Communication/ProvidePropertyWrapper.swift

@@ -116,8 +116,7 @@ extension _ProvidePropertyWrapper: StorageValueProvider {
         } else if let wrapperWithArray = self as? CollectionBasedProvideProperty {
             wrapperWithArray.collectArrayElements(into: &repository)
         } else {
-            // concatenation is handled by the `CollectedModuleValue/reduce` implementation.
-            repository[CollectedModuleValue<Value>.self] = [storedValue]
+            repository.appendValues([CollectModuleValue(storedValue)])
         }
 
         collected = true
@@ -127,16 +126,24 @@ extension _ProvidePropertyWrapper: StorageValueProvider {
 
 extension _ProvidePropertyWrapper: CollectionBasedProvideProperty where Value: AnyArray {
     func collectArrayElements<Repository: SharedRepository<SpeziAnchor>>(into repository: inout Repository) {
-        // concatenation is handled by the `CollectedModuleValue/reduce` implementation.
-        repository[CollectedModuleValue<Value.Element>.self] = storedValue.unwrappedArray
+        repository.appendValues(storedValue.unwrappedArray.map { CollectModuleValue($0) })
     }
 }
 
 
 extension _ProvidePropertyWrapper: OptionalBasedProvideProperty where Value: AnyOptional {
     func collectOptional<Repository: SharedRepository<SpeziAnchor>>(into repository: inout Repository) {
         if let storedValue = storedValue.unwrappedOptional {
-            repository[CollectedModuleValue<Value.Wrapped>.self] = [storedValue]
+            repository.appendValues([CollectModuleValue(storedValue)])
         }
     }
 }
+
+
+extension SharedRepository where Anchor == SpeziAnchor {
+    fileprivate mutating func appendValues<Value>(_ values: [CollectModuleValue<Value>]) {
+        var current = self[CollectedModuleValues<Value>.self]
+        current.append(contentsOf: values)
+        self[CollectedModuleValues<Value>.self] = current
+    }
+}

Sources/Spezi/Dependencies/Property/DependencyPropertyWrapper.swift

@@ -33,6 +33,7 @@ private protocol ModuleArrayDependency {
 public class _DependencyPropertyWrapper<Value> { // swiftlint:disable:this type_name
     private let dependencies: DependencyCollection
 
+    /// The dependency value.
     public var wrappedValue: Value {
         if let singleModule = self as? SingleModuleDependency {
             return singleModule.wrappedValue(as: Value.self)

Sources/Spezi/Spezi/Spezi.swift

@@ -68,11 +68,30 @@ private struct ImplicitlyCreatedModulesKey: DefaultProvidingKnowledgeSource {
 /// The ``Module`` documentation provides more information about the structure of modules.
 /// Refer to the ``Configuration`` documentation to learn more about the Spezi configuration.
 ///
+/// ### Dynamically Loading Modules
+///
+/// While the above examples demonstrated how Modules are configured within your ``SpeziAppDelegate``, they can also be loaded and unloaded dynamically based on demand.
+/// To do so you, you need to access the global `Spezi` instance from within your Module.
+///
+/// Below is a short code example:
+/// ```swift
+/// class ExampleModule: Module {
+///     @Application(\.spezi)
+///     var spezi
+///
+///     func userAuthenticated() {
+///         spezi.loadModule(AccountManagement())
+///         // ...
+///     }
+/// }
+/// ```
+///
 /// ## Topics
 ///
 /// ### Properties
 /// - ``logger``
 /// - ``launchOptions``
+/// - ``spezi``
 ///
 /// ### Actions
 /// - ``registerRemoteNotifications``
@@ -182,16 +201,19 @@ public class Spezi {
     }
 
     private func loadModules(_ modules: [any Module]) {
+        precondition(Self.moduleInitContext == nil, "Modules cannot be loaded within the `configure()` method.")
         let existingModules = self.modules
 
         let dependencyManager = DependencyManager(modules, existing: existingModules)
         dependencyManager.resolve()
 
         implicitlyCreatedModules.formUnion(dependencyManager.implicitlyCreatedModules)
 
+        // we pass through the whole list of modules once to collect all @Provide values
         for module in dependencyManager.initializedModules {
-            // we pass through the whole list of modules once to collect all @Provide values
-            module.collectModuleValues(into: &storage)
+            Self.$moduleInitContext.withValue(module) {
+                module.collectModuleValues(into: &storage)
+            }
         }
 
         for module in dependencyManager.initializedModules {
@@ -201,8 +223,6 @@ public class Spezi {
 
         // Newly loaded modules might have @Provide values that need to be updated in @Collect properties in existing modules.
         for existingModule in existingModules {
-            // TODO: do we really want to support that?, that gets super chaotic with unload modules???
-            // TODO: create an issue to have e.g. update functionality (rework that whole thing?), remove that system altogether?
             existingModule.injectModuleValues(from: storage)
         }
     }
@@ -217,6 +237,8 @@ public class Spezi {
     ///
     /// - Parameter module: The Module to unload.
     public func unloadModule(_ module: any Module) {
+        precondition(Self.moduleInitContext == nil, "Modules cannot be unloaded within the `configure()` method.")
+
         guard module.isLoaded(in: self) else {
             return // module is not loaded
         }
@@ -228,7 +250,7 @@ public class Spezi {
 
         implicitlyCreatedModules.remove(ModuleReference(module))
 
-        // TODO: remove @Collect values that were previously provided by this Module
+        removeCollectValues(for: module)
 
         // re-injecting all dependencies ensures that the unloaded module is cleared from optional Dependencies from
         // pre-existing Modules.
@@ -256,6 +278,29 @@ public class Spezi {
         }
     }
 
+    private func removeCollectValues(for module: any Module) {
+        let valueContainers = storage.collect(allOf: (any AnyCollectModuleValues).self)
+
+        var changed = false
+        for var container in valueContainers {
+            let didChange = container.removeValues(from: module)
+            guard didChange else {
+                continue
+            }
+
+            changed = true
+            container.store(into: &storage)
+        }
+
+        guard changed else {
+            return
+        }
+
+        for module in modules {
+            module.injectModuleValues(from: storage)
+        }
+    }
+
     /// Initialize a Module.
     ///
     /// Call this method to initialize a Module, injecting necessary information into Spezi property wrappers.

Tests/SpeziTests/DependenciesTests/DependencyTests.swift

@@ -15,17 +15,23 @@ import XCTRuntimeAssertions
 private final class TestModule1: Module {
     @Dependency var testModule2 = TestModule2()
     @Dependency var testModule3: TestModule3
+
+    @Provide var num: Int = 1
 }
 
 private final class TestModule2: Module {
     @Dependency var testModule4 = TestModule4()
     @Dependency var testModule5 = TestModule5()
     @Dependency var testModule3: TestModule3
+
+    @Provide var num: Int = 2
 }
 
 private final class TestModule3: Module, DefaultInitializable {
     let state: Int
 
+    @Provide var num: Int = 3
+
     convenience init() {
         self.init(state: 0)
     }
@@ -37,9 +43,13 @@ private final class TestModule3: Module, DefaultInitializable {
 
 private final class TestModule4: Module {
     @Dependency var testModule5 = TestModule5()
+
+    @Provide var num: Int = 4
 }
 
-private final class TestModule5: Module {}
+private final class TestModule5: Module {
+    @Provide var num: Int = 5
+}
 
 private final class TestModule6: Module {}
 
@@ -62,6 +72,8 @@ private final class TestModuleItself: Module {
 
 private final class OptionalModuleDependency: Module {
     @Dependency var testModule3: TestModule3?
+
+    @Collect var nums: [Int]
 }
 
 private final class OptionalDependencyWithRuntimeDefault: Module {
@@ -122,9 +134,6 @@ final class DependencyTests: XCTestCase {
         XCTAssert(testModule1.testModule2.testModule3 === testModule3)
         XCTAssert(testModule1.testModule2.testModule4.testModule5 === testModule5)
         XCTAssert(optionalModuleDependency.testModule3 === testModule3)
-
-        // TODO: test updated view modifiers
-        // TODO: test updated @Collect properties!
     }
 
     func testUnloadingDependencies() throws {
@@ -134,15 +143,21 @@ final class DependencyTests: XCTestCase {
 
         let spezi = Spezi(standard: DefaultStandard(), modules: [optionalModule])
 
+        // test loading and unloading of @Collect/@Provide property values
+        XCTAssertEqual(optionalModule.nums, [])
 
         spezi.loadModule(module3)
+        XCTAssertEqual(optionalModule.nums, [3])
+
         spezi.loadModule(module1)
+        XCTAssertEqual(optionalModule.nums, [3, 5, 4, 2, 1])
 
         try XCTRuntimePrecondition {
             spezi.unloadModule(module3) // cannot unload module that other modules still depend on
         }
 
         spezi.unloadModule(module1)
+        XCTAssertEqual(optionalModule.nums, [3])
 
         var modules = spezi.modules
         func getModule<M: Module>(_ module: M.Type = M.self) throws -> M {
@@ -160,6 +175,7 @@ final class DependencyTests: XCTestCase {
         XCTAssert(optionalModuleLoaded.testModule3 === module3Loaded)
 
         spezi.unloadModule(module3)
+        XCTAssertEqual(optionalModule.nums, [])
 
         modules = spezi.modules