Added a safe assign initializer.

Sources/TSAO.swift

@@ -49,21 +49,37 @@ public struct AssocMap<ValueType> {
         }
     }
 
-    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<ValueType> {
         _storage = v
     }
 }
+
+private final class WeakObjectContainer {
+    weak var object: AnyObject?
+}

Tests/TSAOTests/Tests.swift

@@ -16,6 +16,7 @@ let tupleMap = AssocMap<(Int, CGRect)>()
 
 let objectRetainMap = AssocMap<NSObject>()
 let objectAssignMap = AssocMap<NSObject>(assign: ())
+let objectSafeAssignMap = AssocMap<NSObject>(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])
+    }
 }