ViewModel injection in ViewControllers with Storyboards

Disclaimer: I wrote this post while working at Wolox, I don’t work there anymore but I left the text untouched.

When applying MVVM architecture in iOS, we face some problems regarding the binding and unbinding of a ViewModel to its corresponding ViewController. We consider that the ideal situation would be to inject the view model on initialization, but when using storyboards, this isn’t possible. Having this in mind, how should we inject the ViewModel to its corresponding ViewController?

At Wolox, we use the following approach: The ViewController gets its ViewModel injected and when the view is loaded, we proceed to bind it. Given that ViewControllers cannot receive parameters on initialization when using storyboards, we are left with two options: either assuming that there will always be a ViewModel (explicit dereferenced property) or having the ViewModel as optional. We prefer the second approach, as it isn’t advisable to explicitly dereference a property (due to type safety).

So for each ViewController, we ended up with the same implementation:

public class MyViewController: UIViewController {
    var viewModel: MyViewModel? {
        willSet(maybeViewModel) {
            if let _ = self.viewModel {
                unbindViewModel()
            }
        }

        didSet {
            if isViewLoaded(), let viewModel = self.viewModel {
                bindViewModel(viewModel)
            }
        }
    }
}

We can read this as “when assigning a new viewModel, if we already had a viewModel, first unbind it. When it is assigned, if the view is loaded, proceed to bind the viewModel.”

The problem with this approach is that we end up adding these lines of code in every ViewController with an associated viewModel. It would be great to abstract this behaviour.

First idea

The first approach is to create a class ViewController which inherits from UIViewController and subclass it. There is a problem with this approach: we have many types of ViewControllers which already inherited from UIViewController (for example UITableViewController, UICollectionViewController, etc…), so this forces us to subclass each of them, which isn’t the best approach as we will end up with repeated code.

Second idea

One would then think, let’s use one of the coolest swift features: protocols. We create a ViewModelBindable protocol that looks like this:

public protocol ViewModelBindable: class {

    typealias ViewModel

    var viewModel: ViewModel? { get set }

    func bindViewModel(viewModel: ViewModel)

    func unbindViewModel(viewModel: ViewModel)

    func shouldBindViewModel() -> Bool

}

Here we are saying that ViewModelBindable is a protocol that can only be implemented by a class (in our case, UIView or UIViewController) and that class has to define what type the viewModel is (typealias ViewModel), a property viewModel of type ViewModel, and how the ViewController will bind/unbind to that property.

Now that we could define the protocol, the next step will be adding the ViewModelBindable conformance for our ViewControllers. One would consider adding a default implementation of the protocol (yet another cool Swift feature), something like this:

extension ViewModelBindable where Self: AnyObject {

    public var viewModel: ViewModel? {
        // store view model
    }

}

Unfortunately, this doesn’t compile, we cannot have stored properties on extensions, only computed properties. So now we are facing the same problem as before: we created the protocol ViewModelBindable, but the viewModel remains in each ViewController.

Proposed solution

Luckily, there is a way to solve this and it is using a cool Objective-C feature: associated objects. This allows us to store the viewModel in an extension.

extension ViewModelBindable where Self: AnyObject {

    public var viewModel: ViewModel? {

        get {
            return getAssociatedObject(self, key: &AssociatedKey)
        }

        set(newViewModel) {
            if let viewModel = newViewModel {
               setAssociatedObject(self, value: viewModel, key:
               &AssociatedKey, policy:
           objc_AssociationPolicy.OBJC_ASSOCIATION_RETAIN_NONATOMIC)

            if shouldBindViewModel() {
                    registerBinding(viewModel)
                }
            }
        }
    }

    private func registerUnbinding(viewModel: ViewModel) {
        unbindViewModel(viewModel)
    }

    private func registerBinding(viewModel: ViewModel) {
        bindViewModel(viewModel)
    }

}

Note that for this to work, we need to box view models when they are structs (since associated objects only work for classes).

private var AssociatedKey: UInt = 1

private final class AssociatedObjectBox<T> {
    let value: T

    init(_ x: T) {
        value = x
    }
}

private func lift<T>(x: T) -> AssociatedObjectBox<T> {
    return AssociatedObjectBox(x)
}

private func setAssociatedObject<T>(object: AnyObject, value: T, key: UnsafePointer<Void>, policy: objc_AssociationPolicy) {
    if let v: AnyObject = value as? AnyObject {
        objc_setAssociatedObject(object, key, v,  policy)
    } else {
        objc_setAssociatedObject(object, key, lift(value), policy)
    }
}

private func getAssociatedObject<T>(object: AnyObject, key: UnsafePointer<Void>) -> T? {
    if let v = objc_getAssociatedObject(object, key) as? T {
        return v
    } else if let v = objc_getAssociatedObject(object, key) as? AssociatedObjectBox<T> {
        return v.value
    } else {
        return nil
    }
}

Isn’t it great? We abstracted the view model using protocols and associated objects. Now each ViewController only needs to define which type the viewModel is and how it should be bound.

For example:

public final class MyViewController: UIViewController {
    // Your implementation
}

extension MyViewController: ViewModelBindable {

    public func bindViewModel(viewModel: MyViewModel) {
        // binding logic
    }

}

Note that Xcode infers from the type signature of bindViewModel that for MyViewController ViewModel type is MyViewModel.

We could also overwrite the default implementation of unbindViewModel(), which may be useful for example when binding/unbinding cells.

Here’s the full implementation

Now that we could abstract the binding/unbinding of view models, new interesting things come into light. For example, we are working on a tool which tracks the binding/unbinding of ViewModels for detecting possible memory leaks. Sounds great, doesn’t it?