Replacing Type Code with Polymorphism

Note: In this post I use the terms class, struct and entity to refer to the same thing. Please bear with me!

Replacing type code with polymorphism is one of those refactoring that are, literally, in the book. However, opportunities to refactor to polymorphism can often times be difficult to identify.

Let’s say that we are working on a SaaS system. The system has different account tiers, with different denomination, different prices, and different capabilities.

To be a little bit more specific, let’s assume just to simplify this post, that our service tracks workouts. There are three account tiers: Free, Plus and Premium, according to the following:

Free account:
– Price: free (duh!)
– Can track: runs, hikes and walks
– Number of workouts: 20 per month

Plus account:
– Price: 2 USD a month
– Can track: Everything supported by the Free Tier plus biking, swimming
– Number of workouts: 50 per month

Premium account:
– Price: 10 USD a month
– Can track: Everything supported by the Plus Tier plus yoga, pilates and roller derby.
– Number of workouts: unlimited.

So, a first attempt at implementing this model might look like this:

enum Workout {
    case running
    case hiking
    case walking
    case biking
    case swimming
    case yoga
    case pilates
    case rollerDerby
}

struct Account {
    enum Tier {
        case free
        case plus
        case premium
    }
    
    let id: String
    let tier: Tier   
}

Now let’s imagine that we want to present the number of workouts that can be still tracked in the current month in one of our views. Given the previous model, we could do something like this:

final class Current: UIViewController {
    /**
     Returns number of available workouts, according to the user's account tier, to be presented on screen
    */
    func workoutsAvailable(account: Account) -> String {
        let returnValue: String
        switch account.tier {
        case .free:
            returnValue = "20"
            break
        case .plus:
            returnValue = "50"
            break
        case .premium:
            returnValue = "unlimited"
            break
        }
        
        return returnValue
    }
}

I guess you immediately recognise an issue here: the logic in the workoutsAvailable function should be part of the Account object.

There are multiple reasons for that. First, well, that logic is part of the Account abstraction. Second, if we need to present the number of available workouts in two different places in the UI, we would have to duplicate that logic.

So, let’s do the right thing and move that code to the Account entity

enum Workout {
    case running
    case hiking
    case walking
    case biking
    case swimming
    case yoga
    case pilates
    case rollerDerby
}

struct Account {
    enum Tier {
        case free
        case plus
        case premium
    }
 
    let id: String
    let tier: Tier
 
    var workoutsAvailable: String {
        let returnValue: String
        switch tier {
        case .free:
            returnValue = "20"
            break
        case .plus:
            returnValue = "50"
            break
        case .premium:
            returnValue = "unlimited"
            break
        }
        
        return returnValue
    }
 }

final class Current: UIViewController {
    /**
     Returns number of available workouts, according to the user's account tier, to be presented on screen
     */
    func workoutsAvailable(account: Account) -> String {
        return account.workoutsAvailable
    }
}

It is obvious that the way clients of the Account entity can present the number of available workouts is simpler now. That is important for many different reasons, but to men the most important is readability. Reading and understanding the code in the Current.workousAvailable function takes a second, because there is no complicated logic to understand in there. And as a bonus, when there is no complicated logic to read and understand, it is more difficult to make mistakes.

But now, we also need to present the monthly price in the UI. given our previous experience with the number of available workouts, we might go directly to add that logic to the Account struct

enum Workout {
    case running
    case hiking
    case walking
    case biking
    case swimming
    case yoga
    case pilates
    case rollerDerby
}


struct Account {
    enum Tier {
        case free
        case plus
        case premium
    }
    
    let id: String
    let tier: Tier
    
    var workoutsAvailable: String {
        let returnValue: String
        switch tier {
        case .free:
            returnValue = "20"
            break
        case .plus:
            returnValue = "50"
            break
        case .premium:
            returnValue = "unlimited"
            break
        }
        
        return returnValue
    }
    
    var price: String {
        let returnValue: String
        switch tier {
        case .free:
            returnValue = "Free"
            break
        case .plus:
            returnValue = "2 USD"
            break
        case .premium:
            returnValue = "10 USD"
            break
        }
        
        return returnValue
    }
}

(As a side note, the price should not be typed as a string, but as another struct called Currency, with two properties: the value and the unit it is measured in. But that’s beyond the scope of this post)

Now, notice how the Account struct is basically has the same logic duplicated in two calculated properties. Now, if we add the supported workout types to the mix, we would have three different properties implementing more or less the same logic.

We are turning our Account abstraction into something that can do three different things, depending on a type parameter.

We have code that affects the behaviour of a class, according to an internal enumeration. That is a good indicator that, maybe, the Account struct should be broken down into different classes, each and every one of them modelling only one of those behaviours.

So here is where the refactoring in the post title comes to help.

What do we really need our accounts to provide? Price, number of workouts and a boolean that tells us if that account can track a particular workout type.

Let’s model that as a protocol:

enum Workout {
    case running
    case hiking
    case walking
    case biking
    case swimming
    case yoga
    case pilates
    case rollerDerby
}

protocol Account {
    var workoutsAvailable: String { get }
    var price: String { get }
    func canTrack(workout: Workout) -> Bool
}

Now, and this also might be a matter for another post, there is something in that protocol that I find a little bit annoying: the lack of symmetry. So I am going to refactor it to:

enum Workout {
    case running
    case hiking
    case walking
    case biking
    case swimming
    case yoga
    case pilates
    case rollerDerby
}

protocol Account {
    func workoutsAvailable() -> String
    func price()-> String
    func canTrack(workout: Workout) -> Bool
}

Now, I shall provide a different implementation of this protocol for each account tier:

enum Workout {
    case running
    case hiking
    case walking
    case biking
    case swimming
    case yoga
    case pilates
    case rollerDerby
}

protocol Account {
    func workoutsAvailable() -> String
    func price()-> String
    func canTrack(workout: Workout) -> Bool
}

struct Free: Account {
    func workoutsAvailable() -> String {
        return "20"
    }
    
    func price()-> String {
        return "Free"
    }
    
    func canTrack(workout: Workout) -> Bool {
        return (workout == .running ||
                workout == .hiking ||
                workout == .walking)
    }
}

struct Plus: Account {
    func workoutsAvailable() -> String {
        return "50"
    }
    
    func price()-> String {
        return "2 USD"
    }
    
    func canTrack(workout: Workout) -> Bool {
        return (Free().canTrack(workout: workout) ||
                workout == .biking ||
                workout == .swimming)
    }
}

struct Premium: Account {
    func workoutsAvailable() -> String {
        return "Unlimited"
    }
    
    func price()-> String {
        return "10 USD"
    }
    
    func canTrack(workout: Workout) -> Bool {
        return true
    }
}

Now, each struct models a simple behaviour. There is more code, true, but the code is dumber, simpler. Each logical unit (struct) model a single behaviour. That makes the code easier to read, easier to understand, and therefore more robust and easier to maintain.

Leave a Reply

Your email address will not be published. Required fields are marked *