diff --git a/Sources/TSAO.swift b/Sources/TSAO.swift index 07cf5f4..6750723 100644 --- a/Sources/TSAO.swift +++ b/Sources/TSAO.swift @@ -49,21 +49,37 @@ public struct AssocMap { } } - fileprivate init(_policy: objc_AssociationPolicy) { + fileprivate init(_policy: objc_AssociationPolicy, useWeakObjectContainer: Bool = false) { + precondition(!useWeakObjectContainer || _policy == .OBJC_ASSOCIATION_RETAIN) self._policy = _policy + self.useWeakObjectContainer = useWeakObjectContainer } private let _policy: objc_AssociationPolicy private let _key: _AssocKey = _AssocKey() + private let useWeakObjectContainer: Bool private func _get(_ object: AnyObject) -> AnyObject? { let p = Unmanaged.passUnretained(_key).toOpaque() - return objc_getAssociatedObject(object, p) as AnyObject? + if useWeakObjectContainer { + return (objc_getAssociatedObject(object, p) as? WeakObjectContainer)?.object + } else { + return objc_getAssociatedObject(object, p) as AnyObject? + } } private func _set(_ object: AnyObject, _ value: AnyObject?) { let p = Unmanaged.passUnretained(_key).toOpaque() - objc_setAssociatedObject(object, p, value, _policy) + if useWeakObjectContainer { + let container = (objc_getAssociatedObject(object, p) as? WeakObjectContainer) ?? { + let new = WeakObjectContainer() + objc_setAssociatedObject(object, p, new, _policy) + return new + }() + container.object = value + } else { + objc_setAssociatedObject(object, p, value, _policy) + } } } @@ -74,6 +90,11 @@ extension AssocMap where ValueType: AnyObject { public init(assign: ()) { self.init(_policy: .OBJC_ASSOCIATION_ASSIGN) } + + /// Initializes an `AssocMap` with a safe assign policy. Accessing the value will return nil if object has been deallocated. + public init(safeAssign: ()) { + self.init(_policy: .OBJC_ASSOCIATION_RETAIN, useWeakObjectContainer: true) + } } extension AssocMap where ValueType: NSCopying { @@ -116,3 +137,7 @@ private final class _AssocValueBox { _storage = v } } + +private final class WeakObjectContainer { + weak var object: AnyObject? +} diff --git a/Tests/TSAOTests/Tests.swift b/Tests/TSAOTests/Tests.swift index 10da0a4..85d5600 100644 --- a/Tests/TSAOTests/Tests.swift +++ b/Tests/TSAOTests/Tests.swift @@ -16,6 +16,7 @@ let tupleMap = AssocMap<(Int, CGRect)>() let objectRetainMap = AssocMap() let objectAssignMap = AssocMap(assign: ()) +let objectSafeAssignMap = AssocMap(safeAssign: ()) class TSAOTests: XCTestCase { var helper: NSObject! @@ -100,4 +101,16 @@ class TSAOTests: XCTestCase { } XCTAssertNil(object) } + + func testSafeAssignValue() { + weak var object: NSObject? + autoreleasepool { + let obj = NSObject() + objectSafeAssignMap[helper] = obj + object = obj + XCTAssertEqual(obj, objectSafeAssignMap[helper]) + } + XCTAssertNil(object) + XCTAssertNil(objectSafeAssignMap[helper]) + } }