Back WeakDictionary with NSMapTable

Author: LinusU

WeakDictionary/WeakDictionary.swift

@@ -8,90 +8,153 @@
 
 import Foundation
 
+fileprivate class AnyHashableBox: NSObject {
+    private let hashable: AnyHashable
+
+    init<T: Hashable>(_ base: T) {
+        hashable = AnyHashable(base)
+    }
+
+    var base: Any {
+        return hashable.base
+    }
+
+    @objc override var hash: Int {
+        return hashable.hashValue
+    }
+
+    @objc override func isEqual(_ object: Any?) -> Bool {
+        return hashable == (object as! AnyHashableBox).hashable
+    }
+}
+
 public struct WeakDictionary<Key: Hashable, Value: AnyObject> {
 
-    private var storage: [Key: WeakDictionaryReference<Value>]
+    private var storage: NSMapTable<AnyHashableBox, Value>
 
     public init() {
-        self.init(storage: [Key: WeakDictionaryReference<Value>]())
+        self.init(storage: NSMapTable(keyOptions: .strongMemory, valueOptions: .weakMemory))
     }
 
     public init(dictionary: [Key: Value]) {
-        var newStorage = [Key: WeakDictionaryReference<Value>]()
-        dictionary.forEach({ key, value in newStorage[key] = WeakDictionaryReference<Value>(value: value) })
+        let newStorage = NSMapTable<AnyHashableBox, Value>(keyOptions: .strongMemory, valueOptions: .weakMemory, capacity: dictionary.capacity)
+        dictionary.forEach({ key, value in newStorage.setObject(value, forKey: AnyHashableBox(key)) })
         self.init(storage: newStorage)
     }
 
-    private init(storage: [Key: WeakDictionaryReference<Value>]) {
+    private init(storage: NSMapTable<AnyHashableBox, Value>) {
         self.storage = storage
     }
 
-    public mutating func reap() {
-        storage = weakDictionary().storage
-    }
-
     public func weakDictionary() -> WeakDictionary<Key, Value> {
         return self[startIndex ..< endIndex]
     }
 
     public func dictionary() -> [Key: Value] {
         var newStorage = [Key: Value]()
 
-        storage.forEach { key, value in
-            if let retainedValue = value.value {
-                newStorage[key] = retainedValue
+        for key in storage.keyEnumerator() {
+            if let value = storage.object(forKey: (key as! AnyHashableBox)) {
+                newStorage[(key as! AnyHashableBox).base as! Key] = value
             }
         }
 
         return newStorage
     }
 }
 
+public class WeakDictionaryIndex<Key: Hashable, Value> {
+    fileprivate let key: AnyHashableBox?
+
+    private let position: UInt
+    private var enumerator: NSEnumerator?
+    private var nextIndex: WeakDictionaryIndex<Key, Value>?
+
+    fileprivate init(enumerator: NSEnumerator?, position: UInt) {
+        self.key = enumerator?.nextObject() as! AnyHashableBox?
+        self.enumerator = enumerator
+        self.position = position
+    }
+
+    fileprivate init() {
+        self.key = nil
+        self.enumerator = nil
+        self.position = UInt.max
+    }
+
+    fileprivate func next() -> WeakDictionaryIndex<Key, Value> {
+        if let nextIndex = self.nextIndex {
+            return nextIndex
+        }
+
+        let nextIndex = WeakDictionaryIndex<Key, Value>(enumerator: enumerator, position: position + 1)
+
+        self.nextIndex = nextIndex
+        self.enumerator = nil
+
+        return nextIndex
+    }
+}
+
+extension WeakDictionaryIndex: Equatable {
+    public static func == (lhs: WeakDictionaryIndex<Key, Value>, rhs: WeakDictionaryIndex<Key, Value>) -> Bool {
+        return lhs.position == rhs.position
+    }
+}
+
+extension WeakDictionaryIndex: Comparable {
+    public static func < (lhs: WeakDictionaryIndex<Key, Value>, rhs: WeakDictionaryIndex<Key, Value>) -> Bool {
+        return lhs.position < rhs.position
+    }
+}
+
 extension WeakDictionary: Collection {
 
-    public typealias Index = DictionaryIndex<Key, WeakDictionaryReference<Value>>
+    public typealias Index = WeakDictionaryIndex<Key, Value>
 
     public var startIndex: Index {
-        return storage.startIndex
+        return Index(enumerator: storage.keyEnumerator(), position: 0)
     }
 
     public var endIndex: Index {
-        return storage.endIndex
+        return Index()
     }
 
     public func index(after index: Index) -> Index {
-        return storage.index(after: index)
+        return index.next()
     }
 
-    public subscript(position: Index) -> (Key, WeakDictionaryReference<Value>) {
-        return storage[position]
+    public subscript(position: Index) -> (Key, Value) {
+        guard let key = position.key else {
+            fatalError("Attempting to access WeakDictionary elements using an invalid index")
+        }
+
+        return (key.base as! Key, storage.object(forKey: key)!)
     }
 
     public subscript(key: Key) -> Value? {
         get {
-            guard let valueRef = storage[key] else {
-                return nil
-            }
-
-            return valueRef.value
+            return storage.object(forKey: AnyHashableBox(key))
         }
 
         set {
             guard let value = newValue else {
-                storage[key] = nil
+                storage.removeObject(forKey: AnyHashableBox(key))
                 return
             }
 
-            storage[key] = WeakDictionaryReference<Value>(value: value)
+            storage.setObject(value, forKey: AnyHashableBox(key))
         }
     }
 
     public subscript(bounds: Range<Index>) -> WeakDictionary<Key, Value> {
-        let subStorage = storage[bounds.lowerBound ..< bounds.upperBound]
-        var newStorage = [Key: WeakDictionaryReference<Value>]()
+        let newStorage = NSMapTable<AnyHashableBox, Value>(keyOptions: .strongMemory, valueOptions: .weakMemory)
 
-        subStorage.filter { _, value in return value.value != nil }
-            .forEach { key, value in newStorage[key] = value }
+        var pos = bounds.lowerBound
+        while pos.key != nil && pos != bounds.upperBound {
+            newStorage.setObject(storage.object(forKey: pos.key), forKey: pos.key)
+            pos = pos.next()
+        }
 
         return WeakDictionary<Key, Value>(storage: newStorage)
     }