One of the main limitations of Swift Extensions is the impossibility to use stored properties. Let’s see a workaround to achieve our goals.

Overview

Swift Extensions allow us to add new functionality to an existing class, structure, enumeration, or protocol. We often use them in our projects, and there are moments where we would like having a way to keep the reference of some objects inside these extensions. Unfortunately, Swift doesn’t provide a straightforward way to do it. In this article I will explain how achieve it with the current API provided by the environment.

First of all, I will show you how I would like using the stored properties and then I will show you the workaround which we need to use. Happy reading!

The Dream

Let’s use, as example, a protocol ToggleProtocol, which has a method toggle. Then, we let UIButton implement this protocol to change the background image depending on the toggle state:

protocol ToggleProtocol {
    func toggle()
}

enum ToggleState {
    case on
    case off
}

extension UIButton: ToggleProtocol {

    private(set) var toggleState = ToggleState.off

    func toggle() {
        toggleState = toggleState == .on ? .off : .on

        if toggleState == .on {
            // Shows background for status on
        } else {
            // Shows background for status off
        }
    }
}

Unfortunately, there is a problem. If we compile this code, the compiler will throw an error in the row of private(set) var toggleState = ToggleState.off:

error: extensions may not contain stored properties.

It means that Swift doesn’t support stored properties inside the extension. Therefore, we cannot use the toggleState property to keep the internal state of our toggle button. For this reason, we need a workaround.

The Workaround

The workaround to use stored properties inside the extensions is using the methods objc_getAssociatedObject and objc_setAssociatedObject which allow us to store an object associated to a key.

objc_getAssociatedObject

Returns the value associated with a given object for a given key.

You can find the documentation in this Apple web page.

This method wants two parameters:

1. object: Any!: The source object for the association. We can pass self as argument since the extension is the place where we want to manage the association.
2. key: UnsafeRawPointer!: A pointer which is the key associated to the object to get.

Finally, this method returns the object which we want.

objc_setAssociatedObject

Sets an associated value for a given object using a given key and association policy.

You can find the documentation in this Apple web page.

This method wants four parameters:

1. object: Any!: The source object for the association. We can pass self as argument since the extension is the place where we want to manage the association.
2. key: UnsafeRawPointer!: A pointer which is the key associated to the object to save.
3. value: Any!: The object to save.
4. policy: objc_AssociationPolicy: The policy used to save the object. It can be:
   • OBJC_ASSOCIATION_ASSIGN: It saves the object with a weak reference, in this way we don’t increase the retain count.
   • OBJC_ASSOCIATION_RETAIN_NONATOMIC: It saves the object not atomically with a strong reference.
   • OBJC_ASSOCIATION_COPY_NONATOMIC: It saves the object not atomically creating its copy.
   • OBJC_ASSOCIATION_RETAIN: It saves the object atomically with a strong reference.
   • OBJC_ASSOCIATION_COPY: It saves the object atomically creating its copy.

Note

If you were wondering the difference between atomic and not atomic:

• Atomic: The object is thread-safe and always has a valid state even if a thread reads it while another one is changing its value.
• Not Atomic: It’s the opposite of atomic. It’s not thread-safe. Of course, it’s faster since atomic has to manage a strategy to keep the object thread-safe. It should be the preferred way if you don’t have to use threads.

That’s all. Now, we can use these two methods to refactor our example used in “The Dream”:

struct AssociatedKeys {
    static var toggleState: UInt8 = 0
}

protocol ToggleProtocol {
    func toggle()
}

enum ToggleState {
    case on
    case off
}

extension UIButton: ToggleProtocol {

    private(set) var toggleState: ToggleState {
        get {
            guard let value = objc_getAssociatedObject(self, &AssociatedKeys.toggleState) as? ToggleState else {
                return .off
            }
            return value
        }
        set(newValue) {
            objc_setAssociatedObject(self, &AssociatedKeys.toggleState, newValue, objc_AssociationPolicy.OBJC_ASSOCIATION_RETAIN_NONATOMIC)
        }
    }

    func toggle() {
        toggleState = toggleState == .on ? .off : .on

        if toggleState == .on {
            // Shows background for status on
        } else {
            // Shows background for status off
        }
    }
}

AssociatedKeys contains the associated keys. It means that if we had two associated keys we may have used it like this:

struct AssociatedKeys {
    static var toggleState: UInt8 = 0
    static var anotherState: UInt8 = 0
}

extension UIButton: ToggleProtocol {
    // ...

    private(set) var anotherState: ToggleState {
        get {
            guard let value = objc_getAssociatedObject(self, &AssociatedKeys.anotherState) as? ToggleState else {
                return .off
            }
            return value
        }
        set(newValue) {
            objc_setAssociatedObject(self, &AssociatedKeys.anotherState, newValue, objc_AssociationPolicy.OBJC_ASSOCIATION_RETAIN_NONATOMIC)
        }
    }

    // ...
}

Since the key must be a pointer (UnsafeRawPointer), we send the address of AssociatedKeys.toggleState with &.

We can refactor a little bit objc_getAssociatedObject using a new function with a generic object type and a default value:

func objc_getAssociatedObject<T>(_ object: Any!, _ key: UnsafeRawPointer!, defaultValue: T) -> T {
    guard let value = objc_getAssociatedObject(object, key) as? T else {
        return defaultValue
    }
    return value
}

Thanks to this method we can change our stored properties in the example with:

private(set) var toggleState: ToggleState {
    get {
        return objc_getAssociatedObject(self, &AssociatedKeys.toggleState, defaultValue: .off)
    }
    set(newValue) {
        objc_setAssociatedObject(self, &AssociatedKeys.toggleState, newValue, objc_AssociationPolicy.OBJC_ASSOCIATION_RETAIN_NONATOMIC)
    }
}

Conclusion

This workaround may seem a not so clean solution, but unfortunately it’s one of the cleanest way to manage a stored property inside an extension, since Swift doesn’t provide a more “Swifty” way to do it.

Update 19/06:

Pablo Roca provided in the comment a better approach for having a clean objc_getAssociatedObject method:

protocol PropertyStoring {

    associatedtype T

    func getAssociatedObject(_ key: UnsafeRawPointer!, defaultValue: T) -> T
}

extension PropertyStoring {
    func getAssociatedObject(_ key: UnsafeRawPointer!, defaultValue: T) -> T {
        guard let value = objc_getAssociatedObject(self, key) as? T else {
            return defaultValue
        }
        return value
    }
}

protocol ToggleProtocol {
    func toggle()
}

enum ToggleState {
    case on
    case off
}

extension UIButton: ToggleProtocol, PropertyStoring {

    typealias T = ToggleState

    private struct CustomProperties {
        static var toggleState = ToggleState.off
    }

    var toggleState: ToggleState {
        get {
            return getAssociatedObject(&CustomProperties.toggleState, defaultValue: CustomProperties.toggleState)
        }
        set {
            return objc_setAssociatedObject(self, &CustomProperties.toggleState, newValue, .OBJC_ASSOCIATION_RETAIN)
        }
    }

    func toggle() {
        toggleState = toggleState == .on ? .off : .on

        if toggleState == .on {
            // Shows background for status on
        } else {
            // Shows background for status off
        }
    }
}

let a = UIButton()
print(a.toggleState)
a.toggleState = .on
print(a.toggleState)

Thank you very much for sharing your approach.