|
8 | 8 |
|
9 | 9 | import Foundation |
10 | 10 |
|
| 11 | +fileprivate class KeyBox<T: Hashable>: NSObject { |
| 12 | + let key: T |
| 13 | + |
| 14 | + init(_ key: T) { |
| 15 | + self.key = key |
| 16 | + super.init() |
| 17 | + } |
| 18 | + |
| 19 | + @objc override var hash: Int { |
| 20 | + return key.hashValue |
| 21 | + } |
| 22 | + |
| 23 | + @objc override func isEqual(_ object: Any?) -> Bool { |
| 24 | + return key == (object as! KeyBox<T>).key |
| 25 | + } |
| 26 | +} |
| 27 | + |
11 | 28 | public struct WeakDictionary<Key: Hashable, Value: AnyObject> { |
12 | 29 |
|
13 | | - private var storage: [Key: WeakDictionaryReference<Value>] |
| 30 | + private var storage: NSMapTable<KeyBox<Key>, Value> |
14 | 31 |
|
15 | 32 | public init() { |
16 | | - self.init(storage: [Key: WeakDictionaryReference<Value>]()) |
| 33 | + self.init(storage: NSMapTable(keyOptions: .strongMemory, valueOptions: .weakMemory)) |
17 | 34 | } |
18 | 35 |
|
19 | 36 | public init(dictionary: [Key: Value]) { |
20 | | - var newStorage = [Key: WeakDictionaryReference<Value>]() |
21 | | - dictionary.forEach({ key, value in newStorage[key] = WeakDictionaryReference<Value>(value: value) }) |
| 37 | + let newStorage = NSMapTable<KeyBox<Key>, Value>(keyOptions: .strongMemory, valueOptions: .weakMemory, capacity: dictionary.capacity) |
| 38 | + dictionary.forEach({ key, value in newStorage.setObject(value, forKey: KeyBox<Key>(key)) }) |
22 | 39 | self.init(storage: newStorage) |
23 | 40 | } |
24 | 41 |
|
25 | | - private init(storage: [Key: WeakDictionaryReference<Value>]) { |
| 42 | + private init(storage: NSMapTable<KeyBox<Key>, Value>) { |
26 | 43 | self.storage = storage |
27 | 44 | } |
28 | 45 |
|
29 | | - public mutating func reap() { |
30 | | - storage = weakDictionary().storage |
31 | | - } |
32 | | - |
33 | 46 | public func weakDictionary() -> WeakDictionary<Key, Value> { |
34 | 47 | return self[startIndex ..< endIndex] |
35 | 48 | } |
36 | 49 |
|
37 | 50 | public func dictionary() -> [Key: Value] { |
38 | 51 | var newStorage = [Key: Value]() |
39 | 52 |
|
40 | | - storage.forEach { key, value in |
41 | | - if let retainedValue = value.value { |
42 | | - newStorage[key] = retainedValue |
| 53 | + for key in storage.keyEnumerator() { |
| 54 | + if let value = storage.object(forKey: (key as! KeyBox<Key>)) { |
| 55 | + newStorage[(key as! KeyBox<Key>).key] = value |
43 | 56 | } |
44 | 57 | } |
45 | 58 |
|
46 | 59 | return newStorage |
47 | 60 | } |
48 | 61 | } |
49 | 62 |
|
| 63 | +public class WeakDictionaryIndex<Key: Hashable, Value> { |
| 64 | + fileprivate let key: KeyBox<Key>? |
| 65 | + |
| 66 | + private let position: UInt |
| 67 | + private var enumerator: NSEnumerator? |
| 68 | + |
| 69 | + fileprivate lazy var next: WeakDictionaryIndex<Key, Value> = { |
| 70 | + defer { self.enumerator = nil } |
| 71 | + return .init(enumerator: enumerator, position: position + 1) |
| 72 | + }() |
| 73 | + |
| 74 | + fileprivate convenience init(enumerator: NSEnumerator?, position: UInt = 0) { |
| 75 | + if let key = enumerator?.nextObject() as! KeyBox<Key>? { |
| 76 | + self.init(key: key, enumerator: enumerator, position: position) |
| 77 | + } else { |
| 78 | + self.init(key: nil, enumerator: nil, position: .max) |
| 79 | + } |
| 80 | + } |
| 81 | + |
| 82 | + fileprivate convenience init() { |
| 83 | + self.init(key: nil, enumerator: nil, position: .max) |
| 84 | + } |
| 85 | + |
| 86 | + private init(key: KeyBox<Key>?, enumerator: NSEnumerator?, position: UInt) { |
| 87 | + self.key = key |
| 88 | + self.enumerator = enumerator |
| 89 | + self.position = position |
| 90 | + } |
| 91 | +} |
| 92 | + |
| 93 | +extension WeakDictionaryIndex: Equatable { |
| 94 | + public static func == (lhs: WeakDictionaryIndex<Key, Value>, rhs: WeakDictionaryIndex<Key, Value>) -> Bool { |
| 95 | + return lhs.position == rhs.position |
| 96 | + } |
| 97 | +} |
| 98 | + |
| 99 | +extension WeakDictionaryIndex: Comparable { |
| 100 | + public static func < (lhs: WeakDictionaryIndex<Key, Value>, rhs: WeakDictionaryIndex<Key, Value>) -> Bool { |
| 101 | + return lhs.position < rhs.position |
| 102 | + } |
| 103 | +} |
| 104 | + |
50 | 105 | extension WeakDictionary: Collection { |
51 | 106 |
|
52 | | - public typealias Index = DictionaryIndex<Key, WeakDictionaryReference<Value>> |
| 107 | + public typealias Index = WeakDictionaryIndex<Key, Value> |
53 | 108 |
|
54 | 109 | public var startIndex: Index { |
55 | | - return storage.startIndex |
| 110 | + return Index(enumerator: storage.keyEnumerator()) |
56 | 111 | } |
57 | 112 |
|
58 | 113 | public var endIndex: Index { |
59 | | - return storage.endIndex |
| 114 | + return Index() |
60 | 115 | } |
61 | 116 |
|
62 | 117 | public func index(after index: Index) -> Index { |
63 | | - return storage.index(after: index) |
| 118 | + return index.next |
64 | 119 | } |
65 | 120 |
|
66 | | - public subscript(position: Index) -> (Key, WeakDictionaryReference<Value>) { |
67 | | - return storage[position] |
| 121 | + public subscript(position: Index) -> (Key, Value) { |
| 122 | + guard let key = position.key else { |
| 123 | + fatalError("Attempting to access WeakDictionary elements using an invalid index") |
| 124 | + } |
| 125 | + |
| 126 | + return (key.key, storage.object(forKey: key)!) |
68 | 127 | } |
69 | 128 |
|
70 | 129 | public subscript(key: Key) -> Value? { |
71 | 130 | get { |
72 | | - guard let valueRef = storage[key] else { |
73 | | - return nil |
74 | | - } |
75 | | - |
76 | | - return valueRef.value |
| 131 | + return storage.object(forKey: KeyBox<Key>(key)) |
77 | 132 | } |
78 | 133 |
|
79 | 134 | set { |
80 | 135 | guard let value = newValue else { |
81 | | - storage[key] = nil |
| 136 | + storage.removeObject(forKey: KeyBox<Key>(key)) |
82 | 137 | return |
83 | 138 | } |
84 | 139 |
|
85 | | - storage[key] = WeakDictionaryReference<Value>(value: value) |
| 140 | + storage.setObject(value, forKey: KeyBox<Key>(key)) |
86 | 141 | } |
87 | 142 | } |
88 | 143 |
|
89 | 144 | public subscript(bounds: Range<Index>) -> WeakDictionary<Key, Value> { |
90 | | - let subStorage = storage[bounds.lowerBound ..< bounds.upperBound] |
91 | | - var newStorage = [Key: WeakDictionaryReference<Value>]() |
| 145 | + let newStorage = NSMapTable<KeyBox<Key>, Value>(keyOptions: .strongMemory, valueOptions: .weakMemory) |
92 | 146 |
|
93 | | - subStorage.filter { _, value in return value.value != nil } |
94 | | - .forEach { key, value in newStorage[key] = value } |
| 147 | + var pos = bounds.lowerBound |
| 148 | + while pos.key != nil && pos != bounds.upperBound { |
| 149 | + newStorage.setObject(storage.object(forKey: pos.key), forKey: pos.key) |
| 150 | + pos = pos.next |
| 151 | + } |
95 | 152 |
|
96 | 153 | return WeakDictionary<Key, Value>(storage: newStorage) |
97 | 154 | } |
|
0 commit comments