The power of Swift enums
Enumerations, or enums, are a symbolic way to represent a “one of” type. In this post we’ll get a taste of the great flexibility that Swift enums bring to the table, and how they can simplify and clarify our code.
I don’t think it’s untrue to say that a bird can be one of the following: Galah; Kookaburra; or Other.
In C, we might represent the concept like this:
C enums are a pretty leaky abstraction, however. Since they are essentially just ints, we can do strange things like add them together. What should Galah + Kookaburra mean, and why does it equal 1? (If you don’t know the answer to that, ask your nearest greybeard.)
Swift also has enums and they can work very similarly to C:
To make Swift enums work like C, we need to explicitly declare that they are based on Ints, and then extract values with the rawValue property. This is because Swift is “less leaky” than C and quite a bit stricter about enforcing type safety.
Since ints are not appropriate for representing types of birds, in this case we won’t mention the underlying storage type at all. Instead we’ll just write:
But now our program just prints (Enum Value), which is not very useful. How might we log the actual value?
There are a few options. We could use our Int trick as above, but we can do better. We could base our enum on a different raw type, say, Strings:
A more powerful alternative is to employ the Printable protocol:
The protocol declares a property called description which returns a string describing the type. It is conceptually similar to the description method in the NSObject protocol in Objective-C.
This example also shows another interesting feature of enums in Swift: they can contain properties and functions, just like classes.
Wait, what?
Swift enums can also have associated values, and this is where things start to really diverge from the venerable C heritage.
We can add parameters to the cases in our enums. Let’s change our Other case to also have an associated name:
Now let’s fix our description property to return this name. We do this by binding a variable to the pattern in our switch expression:
Now we can give names to any bird not already covered in our enum and log them sensibly:
Binary trees
Let’s say we want to build our own binary tree data type that stores sorted integers. We all need to do that. All the time. (Well, in interviews, anyway.)
If we wanted to use a class, we might do it like this:
We’ll want to throw in a constructor:
So far so good. Nice, simple code. Note that we’re using optional types for the left and right subtrees since some nodes won’t have subtrees (the tree has to stop somewhere).
We’ll also want an insert function that takes a value and returns a new tree with it stored in the appropriate place:
func insert(newValue: Int) -> Tree {
if newValue > value {
if let theRight = right {
return Tree(value: value,
left: left,
right: theRight.insert(newValue))
} else {
return Tree(value: value,
left: left,
right: Tree(value: newValue, left: nil, right: nil))
}
} else {
if let theLeft = left {
return Tree(value: value,
left: theLeft.insert(newValue)
right: right)
} else {
return Tree(value: value,
left: Tree(value: newValue, left: nil, right: nil),
right: right)
}
}
}
This function recursively traverses the tree to find the point where it can insert the new value. Because we’re using optional types for subtrees, we need to contort the code to work around the cases when they do or do not exist.
This code is functional, but not very easy to read. It could be improved by extracting functions and other refactorings, but let’s take a different tack entirely.
Binary trees with enums
Since enums can have associated values, maybe we can just define a tree recursively like so:
Actually, no we can’t. The compiler will complain that “Recursive value types are not allowed.” This is an irritating limitation of the current version of Swift (v1.2 at time of writing).
However, there is a workaround. First let’s declare a new protocol:
Next let’s make our enum extend the protocol and have our associated values refer to it:
This is permitted, since the enum is not then directly recursive. Hopefully future versions of Swift will allow recursive enums, but let’s go with this for now.
We still have the problem that subtrees of a node may not exist, but instead of using optionals, let’s model this with another case in our enum:
Our left and right subtrees are now guaranteed not to be nil. Instead, they are “one of” a node or an empty tree. Consequently, our insertion method can be simplified because the places we deal with an empty tree can be consolidated.
We’ll first need to add the insert function declaration to our protocol:
Then we can implement it in our enum:
func insert(newValue: Int) -> Tree {
switch self {
case .Empty:
return Tree.Node(newValue, left: Tree.Empty, right: Tree.Empty)
case let .Node(value, left, right):
if newValue > value {
return Tree.Node(value, left: left, right: right.insert(newValue))
} else {
return Tree.Node(value, left: left.insert(newValue), right: right)
}
}
}
The logic in this function is much more readable than that in our class-based tree. There is less duplication of code, and the syntactic structure mirrors the algorithmic structure, rather than being obfuscated by special cases.
By encoding special knowledge of the tree data structure into our enum cases, instead of using a generic catch-all optional type, we are able to write clearer and more concise code.
Wrapping up
We’ll finish with a complete listing of the tree type, and while we’re at it let’s also add a handy tree depth property and make the whole thing Printable:
protocol TreeProtocol {
func insert(Int) -> TreeProtocol
var depth: Int { get }
}
enum Tree: TreeProtocol, Printable {
case Empty
case Node(Int, left: TreeProtocol, right: TreeProtocol)
func insert(newValue: Int) -> Tree {
switch self {
case .Empty:
return Tree.Node(newValue, left: Tree.Empty, right: Tree.Empty)
case let .Node(value, left, right):
if newValue > value {
return Tree.Node(value, left: left, right: right.insert(newValue))
} else {
return Tree.Node(value, left: left.insert(newValue), right: right)
}
}
}
var depth: Int {
switch self {
case .Empty:
return 0
case let .Node(_, left, right):
return 1 + max(left.depth, right.depth)
}
}
var description: String {
switch self {
case .Empty:
return "."
case let .Node(value, left, right):
return "[\(left) \(value) \(right)]"
}
}
}
In this post, we’ve dealt only with binary trees of integers. Swift, however, has generics which allows us to make this tree work for any type of value, not just integers. We could also overload some operators to make working with the tree type more readable. But those are topics for a future post.
