Weak Dictionary Values in Swift

When dealing with dictionaries and arrays in Swift, it's very important to know that any reference types used as a key or value will be retained by default. Let's see some interesting cases where this behavior is problematic, and how we can use types like NSMapTable, NSHashTable and NSPointerArray to create weak collections that perform better than their Swift counterparts.

To visualize why the default memory management behavior of Swift dictionaries can sometimes be bad for you, let's attempt to build a system that uses dictionaries to store instances of some arbitrary types. We can use RouterService's dependency management system as an example:

protocol Dependency: AnyObject {
    static var identifier: String { get }
}

final class DependencyStore {

    var factories = [String: () -> Dependency]()
    var cachedDependencies = [String: Dependency]()

    func register<T: Dependency>(dependencyFactory: @escaping () -> T) {
        let key = T.identifier
        factories[key] = dependencyFactory
    }

    func getInstanceOf<T: Dependency>(dependencyType type: T.Type) -> T? {
        let key = T.identifier
        let cachedInstance = cachedDependencies[key]
        if cachedInstance != nil {
            return cachedInstance as? T
        } else {
            let newInstance = factories[key]?()
            cachedDependencies[key] = newInstance
            return newInstance as? T
        }
    }
}

In this dependency injection system, the app is able to register a Dependency by providing a closure that creates such dependency. When a feature states that this dependency is needed in order for it to work, the system creates an instance of it and stores it for future uses.

However, what happens when a feature doesn't need these dependencies anymore? Because dictionary values are retained, the dependency instances will live forever in the app, even if the features that needed them are long gone.

To solve this, you could create some complex system where features could warn the DependencyStore that they are going to close, or, you could create a weak dictionary. A weak dictionary works exactly like a normal one, with the difference being that it doesn't retain its values. If cachedDependencies was a weak dictionary, its values would automatically go out of scope if nobody else in the app was referencing them. This way, dependencies would correctly be created and disposed of as necessary by the app s lifecycle.

Creating a Weak Dictionary in Pure Swift

A way to solve this problem in pure Swift is to create some sort of wrapper class that has a weak reference inside, and stored our values there instead:

final class WeakRef {
    weak var value: AnyObject?
}

var cachedDependencies = [String: WeakRef]()

func getInstanceOf<T: Dependency>(dependencyType type: T.Type) -> T? {
    let key = T.identifier
    let cachedInstance = cachedDependencies[key]?.value
    if cachedInstance != nil {
        print("Using old value!")
        return cachedInstance as? T
    } else {
        print("Using new value!")
        let newInstance = factories[key]?()
        let ref = WeakRef()
        ref.value = newInstance
        cachedDependencies[key] = ref
        return newInstance as? T
    }
}

With some prints added to our function, we can confirm that it works by playing with pointers in a playground:

class MockDependency: Dependency {
    static let identifier: String = "mock"
}

let store = DependencyStore()

store.register(dependencyFactory: { MockDependency() })

var instance: MockDependency?

autoreleasepool {
    instance = store.getInstanceOf(dependencyType: MockDependency.self)
    // Using new value!
    _ = store.getInstanceOf(dependencyType: MockDependency.self)
    // Using old value!
    _ = store.getInstanceOf(dependencyType: MockDependency.self)
    // Using old value!
    instance = nil
}

instance = store.getInstanceOf(dependencyType: MockDependency.self)
// Using new value!
_ = store.getInstanceOf(dependencyType: MockDependency.self)
// Using old value!

(Note: An autoreleasepool block is used to force iOS to deallocate the instance after setting it to nil. Check my article about autoreleasepool to learn more about it!)

While this works, it doesn't fully solve the problem. Although the instance of our dependencies will be correctly disposed of, our cache dictionary will still contain their WeakRef wrappers, which is an unnecessary memory overhead for us:

print(cachedDependencies)
// ["mock": WeakRef(value: nil)]

If we want to completely get rid of the dictionary entries, we must bring the big guns. Let's take a look at three classes that Foundation provides for this purpose.

NSMapTable

The NSMapTable class at first glance looks just like a plain NSDictionary / Swift Dictionary<>, but it supports a broader range of memory semantics.

NSMapTable is configured through its initializer, which allows you to pick different semantics for memory management and key equality. Here's how we can rewrite DependencyStore to use an NSMapTable that retains values weakly:

import Foundation

final class DependencyStore {

    var factories = [String: () -> Dependency]()

    var cachedDependencies = NSMapTable<NSString, AnyObject>.init(
        keyOptions: .copyIn,
        valueOptions: .weakMemory
    )

    func register<T: Dependency>(dependencyFactory: @escaping () -> T) {
        let key = T.identifier
        factories[key] = dependencyFactory
    }

    func getInstanceOf<T: Dependency>(dependencyType type: T.Type) -> T? {
        let key = T.identifier
        let cachedInstance = cachedDependencies.object(forKey: key as NSString)
        if cachedInstance != nil {
            print("Using old value!")
            return cachedInstance as? T
        } else {
            print("Using new value!")
            let newInstance = factories[key]?()
            cachedDependencies.setObject(newInstance, forKey: key as NSString)
            return newInstance as? T
        }
    }
}

If we try to run the previous example again we'll get the same output, but with the lovely difference that printing NSMapTable will reveal that the entire dictionary entry is gone when the dependencies are disposed!

print(cachedDependencies)
// NSMapTable {
// }

Here are some options supported by NSMapTable. Because the options argument is an OptionSet, you can use several of them at the same. Here's my article about OptionSets if you're looking for more info on them!

• .strongMemory - Retains the data.
• .weakMemory - Do not retain the data.
• .copyIn - Create a copy of the data through NSCopying.
• .objectPointerPersonality - When used for keys, equality is determined based on pointer equality instead of NSObject hash equality.

As an important disclaimer, note that NSMapTable is an Objective-C class. Besides only being able to use reference types with it, the equality of keys is not determined by the Swift Hashable protocol. If you do end up having to use these classes, I recommend using nothing but NSStrings as key types to prevent weird bugs.

Additionally, you should be mindful when using this class with weak keys. Although the functionality is similar, weak keys are only removed when the map table is resized.

NSHashTable

For completion purposes, let's take a look at other Foundation classes that provide the same functionality. NSHashTable works exactly the same as an NSMapTable, but instead of providing additional memory options for dictionaries, it does so for a Set:

let stringRef = "AnUniqueString" as NSString

let weakSetOfStrings = NSHashTable<NSString>(options: .weakMemory)
weakSetOfStrings.add(stringRef)
weakSetOfStrings.contains(stringRef)

NSPointerArray

Once again, NSPointerArray is very similar to NSMapTable, but with the difference being that it represents a plain array instead.

Unfortunately, unlike the other types, NSPointerArray doesn't support adding/retrieving objects directly. To use it, you must convert to and from raw pointer types.

let dependency = MockDependency()

let weakArray = NSPointerArray(options: .weakMemory)

let pointer = Unmanaged.passUnretained(dependency).toOpaque()
weakArray.addPointer(pointer)

let storedValue = Unmanaged<MockDependency>
                      .fromOpaque(weakArray.pointer(at: 0)!)
                      .takeUnretainedValue()

Conclusion

Besides the obvious memory management differences between disposing of a value when it's not needed anymore versus keeping it alive forever, using weak collections when applicable can be very useful for performance reasons. As shown in the examples, while a bare Swift weak wrapper keeps existing when its value is gone, classes like NSMapTable clean their entries completely when values are disposed of. This reduced overhead can make a big difference in your app's lifecycle.