A nice little type-safe wrapper around ObjC associated objects

AssociatedObjects.swift

//
//  AssociatedObjects.swift
//
//  Created by Harlan Haskins on 11/22/24.
//

import Foundation
internal import ObjectiveC

/// The ownership semantics by which an Objective-C associated object will be stored.
@frozen
public enum AssociationPolicy {
    /// Specifies an unowned reference to the associated object.
    /// - Warning: This policy does not zero out references when the associated object is deallocated.
    /// - Note: Equivalent to `OBJC_ASSOCIATION_ASSIGN`
    case unowned
    
    /// Specifies a strong reference to the associated object.
    /// - Note: Equivalent to `OBJC_ASSOCIATION_RETAIN_NONATOMIC`
    case strong
    
    /// Specifies a strong reference to the associated object with atomic access.
    /// - Note: Equivalent to `OBJC_ASSOCIATION_RETAIN`
    case strongAtomic
    
    /// Specifies a copy of the associated object.
    /// - Note: Equivalent to `OBJC_ASSOCIATION_COPY_NONATOMIC`
    case copy
    
    /// Specifies a copy of the associated object with atomic access.
    /// - Note: Equivalent to `OBJC_ASSOCIATION_COPY`
    case copyAtomic
    
    /// Converts the Swift enum case to the corresponding `objc_AssociationPolicy`.
    fileprivate var objcPolicy: objc_AssociationPolicy {
        switch self {
        case .unowned: .OBJC_ASSOCIATION_ASSIGN
        case .strong: .OBJC_ASSOCIATION_RETAIN_NONATOMIC
        case .strongAtomic: .OBJC_ASSOCIATION_RETAIN
        case .copy: .OBJC_ASSOCIATION_COPY_NONATOMIC
        case .copyAtomic: .OBJC_ASSOCIATION_COPY
        }
    }
}

/// A key for accessing values stored as associated objects on Objective-C objects.
///
/// Similar to SwiftUI's `EnvironmentKey`, this protocol enables type-safe access to associated objects
/// by defining both the value type and storage policy for a given key. This allows you to extend objects
/// with additional properties while maintaining strong type safety and clear ownership semantics.
///
/// Example usage:
/// ```swift
/// // Define a custom key
/// private enum ThemeKey: AssociatedObjectKey {
///     static let association: AssociationPolicy = .strong
///     typealias Value = Theme
/// }
///
/// extension UIViewController {
///     var theme: Theme {
///         get { associatedObjects[ThemeKey.self] ?? .default }
///         set { associatedObjects[ThemeKey.self] = newValue }
///     }
/// }
/// ```
public protocol AssociatedObjectKey {
    /// The type of value associated with this key.
    associatedtype Value

    /// The memory management policy to use when storing the associated value.
    ///
    /// This property determines how the runtime manages the lifecycle of the associated object:
    /// - `.strong`: Creates a strong reference to the value (non-atomic)
    /// - `.strongAtomic`: Creates a strong reference with atomic access
    /// - `.copy`: Stores a copy of the value (non-atomic)
    /// - `.copyAtomic`: Stores a copy of the value with atomic access
    /// - `.unowned`: Creates an unowned reference (caution: may lead to dangling pointers)
    ///
    /// Choose the policy based on your ownership and thread-safety requirements.
    static var association: AssociationPolicy { get }
}

extension AssociatedObjectKey {
    /// By default, all associated objects are strong, nonatomic references.
    public static var association: AssociationPolicy {
        .strong
    }
}

extension NSObject {
    /// A type-safe interface for accessing associated objects on an `NSObject` instance.
    ///
    /// `AssociatedObjects` provides a SwiftUI Environment-like interface for working with the Objective-C
    /// runtime's associated object system.
    ///
    /// https://nshipster.com/associated-objects/
    /// 
    /// It offers two ways to store and retrieve associated objects:
    ///
    /// 1. Using a strongly-typed key that conforms to `AssociatedObjectKey`:
    /// ```swift
    /// private enum ThemeKey: AssociatedObjectKey {
    ///     static let association: AssociationPolicy = .strong
    ///     typealias Value = Theme
    /// }
    ///
    /// viewController.associatedObjects[ThemeKey.self] = theme
    /// ```
    ///
    /// 2. Using the type itself as the key with a default strong, non-atomic policy:
    /// ```swift
    /// viewController.associatedObjects[Theme.self] = theme
    /// ```
    public struct AssociatedObjects {
        /// The object on which associated objects will be stored.
        private var subject: AnyObject
        
        /// Creates a new associated objects accessor for the given object.
        ///
        /// - Parameter subject: The object that will store the associated objects.
        public init(_ subject: AnyObject) {
            self.subject = subject
        }
        
        /// Converts a type to a unique key for use with the Objective-C runtime.
        ///
        /// - Parameter type: The type to use as a key.
        /// - Returns: A pointer that uniquely identifies the type.
        private func key<T>(for type: T.Type) -> UnsafeRawPointer {
            unsafeBitCast(type, to: UnsafeRawPointer.self)
        }
        
        /// Accesses associated objects using a strongly-typed key.
        ///
        /// This subscript provides access to associated objects with custom association policies
        /// defined by the key type.
        ///
        /// - Parameter type: The key type that defines both the value type and association policy.
        /// - Returns: The associated value if it exists and matches the expected type, or nil.
        public subscript<Key: AssociatedObjectKey>(_ type: Key.Type) -> Key.Value? {
            get {
                return objc_getAssociatedObject(self, key(for: type)) as? Key.Value
            }
            nonmutating set {
                objc_setAssociatedObject(self, key(for: type), newValue, type.association.objcPolicy)
            }
        }
        
        /// Accesses associated objects using the value type as the key.
        ///
        /// This subscript provides a simplified interface for cases where custom association policies
        /// aren't needed. It uses a strong, non-atomic reference by default.
        ///
        /// - Parameter type: The type to use as both the key and the value type.
        /// - Returns: The associated value if it exists and matches the expected type, or nil.
        public subscript<ValueType>(_ type: ValueType.Type) -> ValueType? {
            get {
                return objc_getAssociatedObject(self, key(for: type)) as? ValueType
            }
            nonmutating set {
                objc_setAssociatedObject(self, key(for: type), newValue, .OBJC_ASSOCIATION_RETAIN_NONATOMIC)
            }
        }
    }
    
    /// Provides access to associated objects stored on this instance.
    ///
    /// https://nshipster.com/associated-objects/
    ///
    /// For example:
    /// ```swift
    /// extension UIViewController {
    ///     var theme: Theme? {
    ///         get { associatedObjects[Theme.self] }
    ///         set { associatedObjects[Theme.self] = newValue }
    ///     }
    /// }
    /// ```
    public var associatedObjects: AssociatedObjects {
        AssociatedObjects(self)
    }
}