Responsible code sharing using the power of protocol extensions

Photo by Maddi Bazzocco / Unsplash

Extensions in Swift are great. We can extend structs, classes, enums, etc., whether we own the code or not, in almost any way we need. This enables us to do easier initializations, mapping, convenience functions, protocol implementation (one of my favorites) and more. There is only one problem.

"with great power comes great responsibility"

It's very easy to clutter the code with countless of public extensions that pollute the application. We can make extensions private or constrain them using i.e. where Self: SomeClass but this doesn't solve all the issues. Another problem is that global extensions act like a hidden dependency and it's not always clear what is going on i.e. while debugging.

Protocol extensions can mitigate this issues. They are constrained only to the code that implements the protocol. Because it's explicit it's easier to control the scope and reason about the code. To present the concept I will use view controller Embedding protocol with extension.

Imagine an application displaying some data to the user. First data is loading, which can fail, and later the data is displayed. Let's divide each state to a separate view controller:

final class PhotosViewController: UIViewController { }
final class LoadingViewController: UIViewController { }
final class ErrorViewController: UIViewController { }
final class PhotoListViewController: UIViewController { }

We want PhotosViewController to be able to embed each of the other view controllers depending on the state of the screen. We could just add the embedding code into the PhotosViewController but there is a high chance we would need it also for other view controllers so it's not great. We could make it an extension but then every view controller would have the embedding capabilities even if they are not needed. We could make some separate Embedder providing this functionality but this is a bit cumbersome and we can do a lot better.

What I propose is to make a protocol first:

protocol Embedding {
    func embed(_ viewController: UIViewController)
    func embed(_ viewController: UIViewController, in containerView: UIView)
    func removeEmbedded(_ viewController: UIViewController)
    func removeAllEmbedded()
    func isEmbedded(_ comparator: (UIViewController?) -> Bool) -> Bool
}

And when this is done we make a protocol extension with all of the code needed for embedding. The thing this whole post is about:

extension Embedding where Self: UIViewController {
    func embed(_ viewController: UIViewController) {
        embed(viewController, in: view)
    }

    func embed(_ viewController: UIViewController, in containerView: UIView) {
        addChild(viewController)
        viewController.view.frame = containerView.frame
        containerView.addSubview(viewController.view)
        viewController.view.fillInSuperview()
        viewController.didMove(toParent: self)
    }

    func removeEmbedded(_ viewController: UIViewController) {
        viewController.willMove(toParent: nil)
        viewController.view.removeFromSuperview()
        viewController.removeFromParent()
    }

    func removeAllEmbedded() {
        children.forEach { [weak self] child in
            self?.removeEmbedded(child)
        }
    }

    func isEmbedded(_ comparator: (UIViewController?) -> Bool) -> Bool {
        return children.filter(comparator).isEmpty == false
    }
}

Notice a constraint in extension Embedding where Self: UIViewController. It's needed because:

  • We want to make sure this protocol can be only implemented by the UIViewControllers.
  • We need access to self in the extension to actually make the embedding.

It's not an extension on UIViewController so the view controllers didn't automagically learn how to embed. We need to grant them this power:

final class PhotosViewController: UIViewController, Embedding { }

By just implementing a protocol our PhotosViewController is now able to embed child view controllers with ease. Methods from protocol are now available for the view controller to use. Let's simulate asynchronous data loading with asyncAfter and try some embedding:

final class PhotosViewController: UIViewController, Embedding {
    override func viewDidLoad() {
        super.viewDidLoad()
        embed(LoadingViewController())
        
        // Simulate loading data delay
        DispatchQueue.main.asyncAfter(deadline: .now() + 1) { [weak self] in
            self?.removeAllEmbedded()
            if isDataLoaded {
                self?.embed(PhotoListViewController())
            } else {
                self?.embed(ErrorViewController())
            }
        }
    }
}

Trying to embed a child view controller in any other view controller i.e. LoadingViewController fails because func embed(_ viewController: UIViewController) is not available. Just as we intended.

While implementing PhotoListViewController we discovered that some user actions can fail and we need to display a nice error there. This couldn't be simpler:

final class PhotoListViewController: UIViewController, Embedding {
    override func viewDidLoad() {
        super.viewDidLoad()
        
        // Simulate error after some user action
        DispatchQueue.main.asyncAfter(deadline: .now() + 1) { [weak self] in
            self?.removeAllEmbedded()
            self?.embed(ErrorViewController())
        }
    }
}

Right after implementing the Embedded protocol for PhotoListViewController it's able to embed child view controllers.

This is just a simple example. You can do a lot more with this concept. It's definitely one of my favorite Swift features!

Thank you for reading!

P.S. You can find more on view controller composition in this great posts by Dave DeLong

P.S.2 For simplicity I'm not introducing view models here but I strongly advise against dealing with network calls inside of a view controller.

Kamil Tustanowski

Kamil Tustanowski

I'm an iOS developer dinosaur who remembers times when Objective-C was "the only way", we did memory management by hand and whole iPhones were smaller than screens in current models.
Poland