Skip to content

Instantly share code, notes, and snippets.

@harlanhaskins
Created November 23, 2024 16:26
Show Gist options
  • Save harlanhaskins/9766ed3ad5e33791e2ef7c42aa20a3d0 to your computer and use it in GitHub Desktop.
Save harlanhaskins/9766ed3ad5e33791e2ef7c42aa20a3d0 to your computer and use it in GitHub Desktop.
A nice little type-safe wrapper around ObjC associated objects
//
// 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)
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment