Refactoring Examples in Swift

Protocol.swift

// Example of refactoring with protocols
// Original javascript version is from Refactoring (2nd Edition) by Martin Fowler
import Foundation

struct Play {
    enum Genre {
        case tragedy
        case comedy
    }

    let name: String
    let type: Genre
}

struct Performance {
    let playID: String
    let audience: Int
}

struct Invoice {
    let customer: String
    let performances: [Performance]
}

typealias Catalog = [String : Play]

struct StatementData {

    struct PerformanceData {
        let play: Play
        let audience: Int
        let amount: Int
        let volumeCredits: Int

        init(_ calculator: PerformanceCalculator, play: Play) {
            self.play = play
            audience = calculator.audience
            amount = calculator.amount
            volumeCredits = calculator.volumeCredits
        }
    }

    let customer: String
    let performances: [PerformanceData]

    var totalVolumeCredits: Int {
        return performances.reduce(0) {
            result, performance in
            return result + performance.volumeCredits
        }
    }

    var totalAmount: Int {
        return performances.reduce(0) {
            result, performance in
            return result + performance.amount
        }
    }

    init(invoice: Invoice, plays: Catalog) {

        func  playFor(_ performance: Performance) -> Play {
            guard let play = plays[performance.playID] else {
                fatalError("Unknown play")
            }
            return play
        }

        customer = invoice.customer
        performances = invoice.performances.map {
            let calculator = StatementData.createPerformanceCalculator(performance: $0, play: playFor($0))
            return PerformanceData(calculator, play: playFor($0))
        }
    }

    static func createPerformanceCalculator(performance: Performance, play: Play) -> PerformanceCalculator {
        switch play.type {
        case .tragedy:
            return TragedyCalculator(audience: performance.audience)
        case .comedy:
            return ComedyCalculator(audience: performance.audience)
        }
    }
}

protocol PerformanceCalculator {
    var audience: Int { get }
    var amount: Int { get }
    var volumeCredits: Int { get }
}

extension PerformanceCalculator {
    var volumeCredits: Int {
        return max(audience - 30, 0)
    }
}

struct TragedyCalculator: PerformanceCalculator {
    let audience: Int

    var amount: Int {
        var result = 40000
        if audience > 30 {
            result += 1000 * (audience - 30)
        }
        return result
    }
}

struct ComedyCalculator: PerformanceCalculator {
    let audience: Int

    var amount: Int {
        var result = 30000
        if audience > 20 {
            result += 10000 + 500 * (audience - 20)
        }
        result += 300 * audience
        return result
    }

    var volumeCredits: Int {
        return max(audience - 30, 0) + audience / 5
    }
}

func statement(invoice: Invoice, plays: Catalog) -> String {
    return renderPlainText(data: StatementData(invoice: invoice, plays: plays))
}

func renderPlainText(data: StatementData) -> String {
    var result = "Statement for \(data.customer)\n"
    for performance in data.performances {
        result += "  \(performance.play.name): \(usd(performance.amount)) (\(performance.audience) seats)\n"
    }

    result += "Amount owed is \(usd(data.totalAmount))\n"
    result += "You earned \(data.totalVolumeCredits) credits\n"
    return result
}

func htmlStatement(invoice: Invoice, plays: Catalog) -> String {
    return renderHTML(data: StatementData(invoice: invoice, plays: plays))
}

func renderHTML(data: StatementData) -> String {
    var result = "<h1>Statement for \(data.customer)</h1>\n"
    result += "<table>\n"
    result += "<tr><th>play</th><th>seats</th><th>cost</th></tr>\n"
    for performance in data.performances {
        result += "<tr><td>\(performance.play.name)</td>"
        result += "<td>\(performance.audience)</td>"
        result += "<td>\(usd(performance.amount))</td></tr>\n"
    }
    result += "</table>\n"
    result += "<p>Amount owed is <em>\(usd(data.totalAmount))</em></p>\n"
    result += "<p>You earned <em>\(data.totalVolumeCredits)</em> credits</p>\n"
    return result
}

func usd(_ number: Int) -> String {
    let format = NumberFormatter()
    format.locale = Locale(identifier: "en-US")
    format.numberStyle = .currency
    return format.string(from: NSNumber(value: number/100)) ?? "N/A"
}

let plays: Catalog = [
    "hamlet" : Play(name: "Hamlet", type: .tragedy),
    "as-like" : Play(name: "As You Like It", type: .comedy),
    "othello" : Play(name: "Othello", type: .tragedy)
]

let performances = [
    Performance(playID: "hamlet", audience: 55),
    Performance(playID: "as-like", audience: 35),
    Performance(playID: "othello", audience: 40)
]

let invoice = Invoice(customer: "BigCo", performances: performances)

let text = statement(invoice: invoice, plays: plays)
print(text)

let html = htmlStatement(invoice: invoice, plays: plays)
print(html)

Subclass.swift

// Example of refactoring with subclassing
// Original javascript version is from Refactoring (2nd Edition) by Martin Fowler

import Foundation

struct Play {
    enum Genre {
        case tragedy
        case comedy
    }

    let name: String
    let type: Genre
}

struct Performance {
    let playID: String
    let audience: Int
}

struct Invoice {
    let customer: String
    let performances: [Performance]
}

typealias Catalog = [String : Play]

struct StatementData {

    struct PerformanceData {
        let play: Play
        let audience: Int
        let amount: Int
        let volumeCredits: Int

        init(_ calculator: PerformanceCalculator, play: Play) {
            self.play = play
            audience = calculator.audience
            amount = calculator.amount
            volumeCredits = calculator.volumeCredits
        }
    }

    let customer: String
    let performances: [PerformanceData]

    var totalVolumeCredits: Int {
        return performances.reduce(0) {
            result, performance in
            return result + performance.volumeCredits
        }
    }

    var totalAmount: Int {
        return performances.reduce(0) {
            result, performance in
            return result + performance.amount
        }
    }

    init(invoice: Invoice, plays: Catalog) {

        func  playFor(_ performance: Performance) -> Play {
            guard let play = plays[performance.playID] else {
                fatalError("Unknown play")
            }
            return play
        }

        customer = invoice.customer
        performances = invoice.performances.map {
            let calculator = StatementData.createPerformanceCalculator(performance: $0, play: playFor($0))
            return PerformanceData(calculator, play: playFor($0))
        }
    }

    static func createPerformanceCalculator(performance: Performance, play: Play) -> PerformanceCalculator {
        switch play.type {
        case .tragedy:
            return TragedyCalculator(audience: performance.audience)
        case .comedy:
            return ComedyCalculator(audience: performance.audience)
        }
    }
}

class PerformanceCalculator {
    let audience: Int
    
    var amount: Int {
        fatalError("subclass responsibility")
    }

    var volumeCredits: Int {
        return max(audience - 30, 0)
   }

    init(audience: Int) {
        self.audience = audience
    }
}

class TragedyCalculator: PerformanceCalculator {
    override var amount: Int {
        var result = 40000
        if audience > 30 {
            result += 1000 * (audience - 30)
        }
        return result
    }
}

class ComedyCalculator: PerformanceCalculator {
    override var amount: Int {
        var result = 30000
        if audience > 20 {
            result += 10000 + 500 * (audience - 20)
        }
        result += 300 * audience
        return result
    }

    override var volumeCredits: Int {
        return super.volumeCredits + audience / 5
    }
}

func statement(invoice: Invoice, plays: Catalog) -> String {
    return renderPlainText(data: StatementData(invoice: invoice, plays: plays))
}

func renderPlainText(data: StatementData) -> String {
    var result = "Statement for \(data.customer)\n"
    for performance in data.performances {
        result += "  \(performance.play.name): \(usd(performance.amount)) (\(performance.audience) seats)\n"
    }

    result += "Amount owed is \(usd(data.totalAmount))\n"
    result += "You earned \(data.totalVolumeCredits) credits\n"
    return result
}

func htmlStatement(invoice: Invoice, plays: Catalog) -> String {
    return renderHTML(data: StatementData(invoice: invoice, plays: plays))
}

func renderHTML(data: StatementData) -> String {
    var result = "<h1>Statement for \(data.customer)</h1>\n"
    result += "<table>\n"
    result += "<tr><th>play</th><th>seats</th><th>cost</th></tr>\n"
    for performance in data.performances {
        result += "<tr><td>\(performance.play.name)</td>"
        result += "<td>\(performance.audience)</td>"
        result += "<td>\(usd(performance.amount))</td></tr>\n"
    }
    result += "</table>\n"
    result += "<p>Amount owed is <em>\(usd(data.totalAmount))</em></p>\n"
    result += "<p>You earned <em>\(data.totalVolumeCredits)</em> credits</p>\n"
    return result
}

func usd(_ number: Int) -> String {
    let format = NumberFormatter()
    format.locale = Locale(identifier: "en-US")
    format.numberStyle = .currency
    return format.string(from: NSNumber(value: number/100)) ?? "N/A"
}

let plays: Catalog = [
    "hamlet" : Play(name: "Hamlet", type: .tragedy),
    "as-like" : Play(name: "As You Like It", type: .comedy),
    "othello" : Play(name: "Othello", type: .tragedy)
]

let performances = [
    Performance(playID: "hamlet", audience: 55),
    Performance(playID: "as-like", audience: 35),
    Performance(playID: "othello", audience: 40)
]

let invoice = Invoice(customer: "BigCo", performances: performances)

let text = statement(invoice: invoice, plays: plays)
print(text)

let html = htmlStatement(invoice: invoice, plays: plays)
print(html)