Using Swift Extensions To Clean Up Our Code

In the previous article, we talked about ARC and what performance implications it imposes on our apps. Today, we are going to take a look at how we can use extensions to make our code more readable.

What Are Extensions?

Swift extensions are normally used whenever you want to add functionality to some part of a program where you don’t have access to the source code. For example, if you wanted to extend the Swift Standard Library’s Int struct, you could do it like so:

extension Int {
func multiply(by multiplier: Int) -> Int {
return self * multiplier
}
}
10.multiply(by: 2)
// Returns 20

While the above example may not be very useful, there are endless possibilities to the extensions you could create for your app that would make your code cleaner and easier to maintain. We are going to take a look at how we can use extensions to make our code more readable, even if we are working on a project where we actually have access to all the source code.

We are going to implement the foundation of a Binary Search Tree, which means that we will construct protocols for the components we need, provide dummy implementations (that is, the implementations do not actually work, they just make the compiler happy) and then we are going to look at how we can make it more readable.

Step 1: Designing the Appropriate Protocols

A Binary Search Tree is a data structure that is used to store and look up data. It relies on a bunch of nodes, which all store a piece of data and references to two other nodes on the next level of the tree structure. The tree itself has a few methods to help the user add, remove, and find data in the structure. Let’s start by looking at two protocols that could be a good fit for our purposes:

protocol SearchTree {
associatedtype DataType: Comparable

func add(_ data: DataType) -> Bool
func contains(_ data: DataType) -> Bool
func find(_ data: DataType) -> DataType?
func delete(_ data: DataType) -> DataType?
func remove(_ data:DataType) -> Bool
}

protocol BinaryNode {
associatedtype DataType: Comparable

var value: DataType { get set }
var parent: Self? { get set }
var left: Self? { get set }
var right: Self? { get set }
var isLeaf: Bool { get }
}

extension BinaryNode {
var isLeaf: Bool {
return self.left == nil && self.right == nil
}
}

Let’s break down what is happening here.

Our SearchTree protocol specifies that a Search Tree has to implement a few functions to conform to our protocol. It also specifies that the type of data we choose to store in our tree must implement the Comparable protocol (it has to, or we won’t be able to find the data we’re looking for later on).

Our BinaryNode protocol specifies the properties that need to be present in a binary tree node for us to be able to do our manipulations and calculations. It also specifies that the data type stored in the node must conform to the Comparable protocol. We also provide a default implementation for the isLeaf computed property through a protocol extension.

Step 2: Making Dummy Implementations

Let’s put our protocols to work in a dummy implementation of a Binary Search Tree and see what it looks like:

class BST<DataType: Comparable>: SearchTree {
private typealias NodeType = BSN<DataType>
private var root: NodeType?

init(rootValue: DataType) {
self.root = NodeType(value: rootValue)
}

init(rootValue: DataType, leftTree: BST, rightTree: BST) {
self.root = NodeType(value: rootValue, left: leftTree.root, right: rightTree.root)
}

private init(rootNode: NodeType) {
self.root = rootNode
}

func add(_ data: DataType) -> Bool {
return false
}

func contains(_ data: DataType) -> Bool {
return false
}

func find(_ data: DataType) -> DataType? {
return nil
}

func delete(_ data: DataType) -> DataType? {
return nil
}

func remove(_ data:DataType) -> Bool {
return false
}

private final class BSN<DataType: Comparable>: BinaryNode {
var value: DataType
var parent: BSN<DataType>?
var left: BSN<DataType>?
var right: BSN<DataType>?

init(value: DataType) {
self.value = value
}

init(value: DataType, left: BSN?, right: BSN?) {
self.value = value
self.left = left
self.right = right
}
}
}

I don’t want to be that guy, but… this looks terrible. And just imagine when there are actual implementations of the SearchTree protocol methods. A Search Tree is a structure that lends itself well for recursive method implementations, which means that the functions that are specified in the protocol will act as wrappers around yet another method, doubling the method declarations we need.

The main thing with this layout is that it makes it very difficult for us to get an overview. Where do we start looking if we want to change one of the SearchTree protocol implementations? Where do we put new methods if we want to be able to find them quickly later on?

Step 3: Refactoring Using Extensions

Extensions were not really designed to solve any readability problems. They were designed to solve a whole other problem set, but the fact that we can use them this way is a nice side effect. Let’s give it a go!

class BinarySearchTree<DataType: Comparable> {
fileprivate typealias NodeType = BinarySearchNode<DataType>
private var root: NodeType?

init(rootValue: DataType) {
self.root = NodeType(value: rootValue)
}

init(rootValue: DataType, leftTree: BinarySearchTree, rightTree: BinarySearchTree) {
self.root = NodeType(value: rootValue, left: leftTree.root, right: rightTree.root)
}

private init(rootNode: NodeType) {
self.root = rootNode
}
}

extension BinarySearchTree: SearchTree {
func add(_ data: DataType) -> Bool {
return false
}

func contains(_ data: DataType) -> Bool {
return false
}

func find(_ data: DataType) -> DataType? {
return nil
}

func delete(_ data: DataType) -> DataType? {
return nil
}

func remove(_ data:DataType) -> Bool {
return false
}
}

private extension BinarySearchTree {
final class BinarySearchNode<DataType: Comparable>: BinaryNode {
var value: DataType
var parent: BinarySearchNode<DataType>?
var left: BinarySearchNode<DataType>?
var right: BinarySearchNode<DataType>?

init(value: DataType) {
self.value = value
}

init(value: DataType, left: BinarySearchNode?, right: BinarySearchNode?) {
self.value = value
self.left = left
self.right = right
}
}
}

This looks much better. The way the code is laid out suggests a pattern as to where you can find any property or method you may be looking for. Anything that only belongs to this one class will be in our class declaration scope. This will include initializers, properties and methods that are not part of any protocol. This may also contain overridden methods from super classes (if any super classes exist).

Our first extension scope deals with everything that has to do with the conformance to the SearchTree protocol. This scope will also contain all of the methods needed for the recursive calls.

The second extension scope does something pretty neat. It’s a private extension, meaning that whatever is inside it will only be visible inside the scope of the file. Basically, you’re applying fileprivate access control to whatever the extension contains. This means that we can keep implementation details in a scope where it makes sense to use them, and keep all those details hidden from the outside world. In this case, our Search Tree knows everything about the BinarySearchNode, but no one outside will even know it exists.

That’s it for this week! Feel free to comment if you have questions, and follow to get notifications about future articles. I’d love to hear how you use extensions for purposes other than what it was originally designed for!