Rework ObjectAssociation into ObjectStorage

Author: jordanbaird

Ice/ControlItem/ControlItem.swift

@@ -467,7 +467,7 @@ final class ControlItem {
                 keyEquivalent: ""
             )
             item.target = self
-            Self.sectionStorage[item] = section
+            Self.sectionStorage.weakSet(section, for: item)
             switch name {
             case .visible:
                 break
@@ -516,7 +516,7 @@ final class ControlItem {
 
     /// Toggles the menu bar section associated with the given menu item.
     @objc private func toggleMenuBarSection(for menuItem: NSMenuItem) {
-        Self.sectionStorage[menuItem]?.toggle()
+        Self.sectionStorage.value(for: menuItem)?.toggle()
     }
 
     /// Opens the menu bar search panel.
@@ -567,7 +567,7 @@ private extension ControlItem {
     ///
     /// When one of these menu items is created, its section is stored here.
     /// When its action is invoked, the section is retrieved from storage.
-    static let sectionStorage = ObjectAssociation<MenuBarSection>()
+    static let sectionStorage = ObjectStorage<MenuBarSection>()
 }
 
 // MARK: - Logger

Ice/Main/AppState.swift

@@ -251,7 +251,7 @@ final class AppState: ObservableObject {
         // Store whether the app has previously activated inside an internal
         // context to keep it isolated.
         enum Context {
-            static let hasActivated = ObjectAssociation<Bool>()
+            static let hasActivated = ObjectStorage<Bool>()
         }
 
         func activate() {
@@ -263,10 +263,10 @@ final class AppState: ObservableObject {
             NSApp.setActivationPolicy(policy)
         }
 
-        if Context.hasActivated[self] == true {
+        if Context.hasActivated.value(for: self) == true {
             activate()
         } else {
-            Context.hasActivated[self] = true
+            Context.hasActivated.set(true, for: self)
             Logger.appState.debug("First time activating app, so going through Dock")
             // Hack to make sure the app properly activates for the first time.
             NSRunningApplication.runningApplications(withBundleIdentifier: "com.apple.dock").first?.activate()

Ice/Utilities/ObjectAssociation.swift

@@ -0,0 +1,74 @@
+//
+//  ObjectStorage.swift
+//  Ice
+//
+
+import ObjectiveC
+
+// MARK: - Object Storage
+
+/// A type that uses the Objective-C runtime to store values of a given
+/// type with an object.
+final class ObjectStorage<Value> {
+    /// The association policy to use for storage.
+    ///
+    /// - Note: Regardless of whether a value is stored with a strong or
+    ///   weak reference, the association is made strongly. Weak references
+    ///   are stored inside a `WeakReference` object.
+    private let policy = objc_AssociationPolicy.OBJC_ASSOCIATION_RETAIN_NONATOMIC
+
+    /// The key used for value lookup.
+    ///
+    /// The key is unique to this instance.
+    private var key: UnsafeRawPointer {
+        UnsafeRawPointer(Unmanaged.passUnretained(self).toOpaque())
+    }
+
+    /// Sets the value for the given object.
+    ///
+    /// If the value is an object, it is stored with a strong reference.
+    /// Use ``weakSet(_:for:)`` to store an object with a weak reference.
+    ///
+    /// - Parameters:
+    ///   - value: A value to set.
+    ///   - object: An object to set the value for.
+    func set(_ value: Value?, for object: AnyObject) {
+        objc_setAssociatedObject(object, key, value, policy)
+    }
+
+    /// Retrieves the value stored for the given object.
+    ///
+    /// - Parameter object: An object to retrieve the value for.
+    func value(for object: AnyObject) -> Value? {
+        let value = objc_getAssociatedObject(object, key)
+        return if let container = value as? WeakReference {
+            container.object as? Value
+        } else {
+            value as? Value
+        }
+    }
+}
+
+// MARK: - Weak Storage
+
+/// An object containing a weak reference to another object.
+private final class WeakReference {
+    /// A weak reference to an object.
+    private(set) weak var object: AnyObject?
+
+    /// Creates a weak reference to an object.
+    init(_ object: AnyObject) {
+        self.object = object
+    }
+}
+
+extension ObjectStorage where Value: AnyObject {
+    /// Sets a weak reference to an object.
+    ///
+    /// - Parameters:
+    ///   - value: An object to set a weak reference to.
+    ///   - object: An object to set the weak reference for.
+    func weakSet(_ value: Value?, for object: AnyObject) {
+        objc_setAssociatedObject(object, key, value.map(WeakReference.init), policy)
+    }
+}