Add `DefaultsObsevation.tieToLifetime(of:)` (#13)

Co-authored-by: Sindre Sorhus <sindresorhus@gmail.com>

Defaults.xcodeproj/xcshareddata/xcschemes/Defaults-iOS.xcscheme

@@ -27,6 +27,15 @@
       selectedDebuggerIdentifier = "Xcode.DebuggerFoundation.Debugger.LLDB"
       selectedLauncherIdentifier = "Xcode.DebuggerFoundation.Launcher.LLDB"
       shouldUseLaunchSchemeArgsEnv = "YES">
+      <MacroExpansion>
+         <BuildableReference
+            BuildableIdentifier = "primary"
+            BlueprintIdentifier = "52D6D97B1BEFF229002C0205"
+            BuildableName = "Defaults.framework"
+            BlueprintName = "Defaults-iOS"
+            ReferencedContainer = "container:Defaults.xcodeproj">
+         </BuildableReference>
+      </MacroExpansion>
       <Testables>
          <TestableReference
             skipped = "NO"
@@ -41,17 +50,6 @@
             </BuildableReference>
          </TestableReference>
       </Testables>
-      <MacroExpansion>
-         <BuildableReference
-            BuildableIdentifier = "primary"
-            BlueprintIdentifier = "52D6D97B1BEFF229002C0205"
-            BuildableName = "Defaults.framework"
-            BlueprintName = "Defaults-iOS"
-            ReferencedContainer = "container:Defaults.xcodeproj">
-         </BuildableReference>
-      </MacroExpansion>
-      <AdditionalOptions>
-      </AdditionalOptions>
    </TestAction>
    <LaunchAction
       buildConfiguration = "Debug"
@@ -72,8 +70,6 @@
             ReferencedContainer = "container:Defaults.xcodeproj">
          </BuildableReference>
       </MacroExpansion>
-      <AdditionalOptions>
-      </AdditionalOptions>
    </LaunchAction>
    <ProfileAction
       buildConfiguration = "Release"

Sources/Defaults/Observation.swift

@@ -1,8 +1,28 @@
 import Foundation
 
 /// TODO: Nest this inside `Defaults` if Swift ever supported nested protocols.
-public protocol DefaultsObservation {
+public protocol DefaultsObservation: AnyObject {
 	func invalidate()
+
+	/**
+	Keep this observation alive for as long as, and no longer than, another object exists.
+
+	```
+	Defaults.observe(.xyz) { [unowned self] change in
+		self.xyz = change.newValue
+	}.tieToLifetime(of: self)
+	```
+	*/
+	@discardableResult
+	func tieToLifetime(of weaklyHeldObject: AnyObject) -> Self
+
+	/**
+	Break the lifetime tie created by `tieToLifetime(of:)`, if one exists.
+
+	- Postcondition: The effects of any call to `tieToLifetime(of:)` are reversed.
+	- Note: If the tied-to object has already died, then self is considered to be invalidated, and this method has no logical effect.
+	*/
+	func removeLifetimeTie()
 }
 
 extension Defaults {
@@ -95,6 +115,20 @@ extension Defaults {
 		public func invalidate() {
 			object?.removeObserver(self, forKeyPath: key, context: nil)
 			object = nil
+			lifetimeAssociation?.cancel()
+		}
+
+		private var lifetimeAssociation: LifetimeAssociation? = nil
+
+		public func tieToLifetime(of weaklyHeldObject: AnyObject) -> Self {
+			lifetimeAssociation = LifetimeAssociation(of: self, with: weaklyHeldObject, deinitHandler: { [weak self] in
+				self?.invalidate()
+			})
+			return self
+		}
+
+		public func removeLifetimeTie() {
+			lifetimeAssociation?.cancel()
 		}
 
 		// swiftlint:disable:next block_based_kvo
@@ -104,8 +138,12 @@ extension Defaults {
 			change: [NSKeyValueChangeKey: Any]?, // swiftlint:disable:this discouraged_optional_collection
 			context: UnsafeMutableRawPointer?
 		) {
+			guard let selfObject = self.object else {
+				invalidate()
+				return
+			}
+
 			guard
-				let selfObject = self.object,
 				selfObject == object as? NSObject,
 				let change = change
 			else {

Sources/Defaults/util.swift

@@ -17,3 +17,99 @@ extension Decodable {
 		self.init(jsonData: data)
 	}
 }
+
+final class AssociatedObject<T: Any> {
+	subscript(index: Any) -> T? {
+		get {
+			return objc_getAssociatedObject(index, Unmanaged.passUnretained(self).toOpaque()) as! T?
+		} set {
+			objc_setAssociatedObject(index, Unmanaged.passUnretained(self).toOpaque(), newValue, .OBJC_ASSOCIATION_RETAIN_NONATOMIC)
+		}
+	}
+}
+
+/**
+Causes a given target object to live at least as long as a given owner object.
+*/
+final class LifetimeAssociation {
+	private class ObjectLifetimeTracker {
+		var object: AnyObject?
+		var deinitHandler: () -> Void
+
+		init(for weaklyHeldObject: AnyObject, deinitHandler: @escaping () -> Void) {
+			self.object = weaklyHeldObject
+			self.deinitHandler = deinitHandler
+		}
+
+		deinit {
+			deinitHandler()
+		}
+	}
+
+	private static let associatedObjects = AssociatedObject<[ObjectLifetimeTracker]>()
+	private weak var wrappedObject: ObjectLifetimeTracker?
+	private weak var owner: AnyObject?
+
+	/**
+	Causes the given target object to live at least as long as either the given owner object or the resulting `LifetimeAssociation`, whichever is deallocated first.
+
+	When either the owner or the new `LifetimeAssociation` is destroyed, the given deinit handler, if any, is called.
+
+	```
+	class Ghost {
+		var association: LifetimeAssociation?
+
+		func haunt(_ host: Furniture) {
+			association = LifetimeAssociation(of: self, with: host) { [weak self] in
+				// Host has been deinitialized
+				self?.haunt(seekHost())
+			}
+		}
+	}
+
+	let piano = Piano()
+	Ghost().haunt(piano)
+	// The Ghost will remain alive as long as `piano` remains alive.
+	```
+
+	- Parameter target: The object whose lifetime will be extended.
+	- Parameter owner: The object whose lifetime extends the target object's lifetime.
+	- Parameter deinitHandler: An optional closure to call when either `owner` or the resulting `LifetimeAssociation` is deallocated.
+	*/
+	init(of target: AnyObject, with owner: AnyObject, deinitHandler: @escaping () -> Void = { }) {
+		let wrappedObject = ObjectLifetimeTracker(for: target, deinitHandler: deinitHandler)
+		
+		let associatedObjects = LifetimeAssociation.associatedObjects[owner] ?? []
+		LifetimeAssociation.associatedObjects[owner] = associatedObjects + [wrappedObject]
+		
+		self.wrappedObject = wrappedObject
+		self.owner = owner
+	}
+
+	/**
+	Invalidates the association, unlinking the target object's lifetime from that of the owner object. The provided deinit handler is not called.
+	*/
+	func cancel() {
+		wrappedObject?.deinitHandler = {}
+		invalidate()
+	}
+
+	deinit {
+		invalidate()
+	}
+
+	private func invalidate() {
+		guard
+			let owner = owner,
+			let wrappedObject = wrappedObject,
+			var associatedObjects = LifetimeAssociation.associatedObjects[owner],
+			let wrappedObjectAssociationIndex = associatedObjects.firstIndex(where: { $0 === wrappedObject })
+		else {
+			return
+		}
+
+		associatedObjects.remove(at: wrappedObjectAssociationIndex)
+		LifetimeAssociation.associatedObjects[owner] = associatedObjects
+		self.owner = nil
+	}
+}

Tests/DefaultsTests/DefaultsTests.swift

@@ -256,4 +256,38 @@ final class DefaultsTests: XCTestCase {
 		XCTAssertEqual(Defaults[key2], nil)
 		XCTAssertEqual(Defaults[key3], newString3)
 	}
+
+	func testObserveWithLifetimeTie() {
+		let key = Defaults.Key<Bool>("lifetimeTie", default: false)
+		let expect = expectation(description: "Observation closure being called")
+
+		weak var observation: DefaultsObservation!
+		observation = Defaults.observe(key, options: []) { change in
+			observation.invalidate()
+			expect.fulfill()
+		}.tieToLifetime(of: self)
+
+		Defaults[key] = true
+
+		waitForExpectations(timeout: 10)
+	}
+
+	func testObserveWithLifetimeTieManualBreak() {
+		let key = Defaults.Key<Bool>("lifetimeTieManualBreak", default: false)
+
+		weak var observation: DefaultsObservation? = Defaults.observe(key, options: []) { _ in }.tieToLifetime(of: self)
+		observation!.removeLifetimeTie()
+
+		for i in 1...10 {
+			if observation == nil {
+				break
+			}
+
+			sleep(1)
+
+			if i == 10 {
+				XCTFail()
+			}
+		}
+	}
 }

readme.md

@@ -142,6 +142,27 @@ Defaults[.isUnicornMode] = true
 
 In contrast to the native `UserDefaults` key observation, here you receive a strongly-typed change object.
 
+### Invalidate observations automatically
+
+```swift
+extension Defaults.Keys {
+	static let isUnicornMode = Key<Bool>("isUnicornMode", default: false)
+}
+
+final class Foo {
+	init() {
+		Defaults.observe(.isUnicornMode) { change in
+			print(change.oldValue)
+			print(change.newValue)
+		}.tieToLifetime(of: self)
+	}
+}
+
+Defaults[.isUnicornMode] = true
+```
+
+The observation will be valid until `self` is deinitialized.
+
 ### Reset keys to their default values
 
 ```swift
@@ -284,6 +305,47 @@ Type: `func`
 
 Remove all entries from the `UserDefaults` suite.
 
+### `DefaultsObservation`
+
+Type: `protocol`
+
+Represents an observation of a defaults key.
+
+#### `DefaultsObservation.invalidate`
+
+```swift
+DefaultsObservation.invalidate()
+```
+
+Type: `func`
+
+Invalidate the observation.
+
+#### `DefaultsObservation.tieToLifetime`
+
+```swift
+@discardableResult
+DefaultsObservation.tieToLifetime(of weaklyHeldObject: AnyObject) -> Self
+```
+
+Type: `func`
+
+Keep the observation alive for as long as, and no longer than, another object exists.
+
+When `weaklyHeldObject` is deinitialized, the observation is invalidated automatically.
+
+#### `DefaultsObservation.removeLifetimeTie`
+
+```swift
+DefaultsObservation.removeLifetimeTie()
+```
+
+Type: `func`
+
+Break the lifetime tie created by `tieToLifetime(of:)`, if one exists.
+
+The effects of any call to `tieToLifetime(of:)` are reversed. Note however that if the tied-to object has already died, then the observation is already invalid and this method has no logical effect.
+
 
 ## FAQ