AssociatedObject.swift

import Foundation

enum AssociatedObject {
    /// Wrapper for Associated Objects that allows us to use Swift value types.
    internal final class Wrapper<T>: NSObject, NSCopying {
        let value: T

        init(_ value: T) {
            self.value = value
        }

        func copy(with zone: NSZone? = nil) -> Any {
            if let copyingValue = value as? NSCopying {
                // swiftlint:disable:next force_cast
                return Wrapper(copyingValue.copy(with: zone) as! T)
            }

            return Wrapper(value)
        }
    }

    /// Association policy for properties, we only support a subset from `objc_AssociationPolicy`.
    enum Policy {
        /// Retain, nonatomic
        case retain

        /// Copy, nonatomic
        case copy

        /// Specifies a weak reference to the associated object
        case assign

        func asAssociationPolicy() -> objc_AssociationPolicy {
            switch self {
            case .assign:
                return .OBJC_ASSOCIATION_ASSIGN
            case .copy:
                return .OBJC_ASSOCIATION_COPY_NONATOMIC
            case .retain:
                return .OBJC_ASSOCIATION_RETAIN_NONATOMIC
            }
        }
    }

    /// Sets associated object for specified `key`.
    static func set<T>(_ object: T?, on target: AnyObject, forKey key: UnsafeRawPointer, policy: Policy) {
        if object is NSObjectProtocol {
            objc_setAssociatedObject(target, key, object, policy.asAssociationPolicy())
            return
        }
        let wrapper = Wrapper(object)
        objc_setAssociatedObject(target, key, wrapper, policy.asAssociationPolicy())

    }

    /// Returns associated object for `key`.
    static func get<T>(from target: AnyObject, forKey key: UnsafeRawPointer) -> T? {
        if let object = objc_getAssociatedObject(target, key) as? T {
            return object
        } else if let object = objc_getAssociatedObject(target, key) as? Wrapper<T> {
            return object.value
        } else {
            return nil
        }
    }
}

extension NSObject {
    func associateObject<T>(_ object: T, forKey key: UnsafeRawPointer, policy: AssociatedObject.Policy = .retain) {
        AssociatedObject.set(object, on: self, forKey: key, policy: policy)
    }
    func getAssociatedObject<T>(forKey key: UnsafeRawPointer) -> T? {
        return AssociatedObject.get(from: self, forKey: key)
    }
}