An important part of maintaining any app, framework or system is dealing with legacy code. No matter how well architected a system is, legacy will always be built up over time - it can be because of changes in an underlying SDK, because of an expanded feature set, or simply because no one on the team really knows how a particular part works.

I'm a big proponent of dealing with legacy code on an ongoing basis, rather than waiting for a system to become so entangled in itself that it has to be completely rewritten. While complete rewrites can sound tempting (the classic "We rewrote it from the ground up") they are - in my experience - rarely worth it. Usually what ends up happening is that the existing bugs and problems simply get replaced with new ones 😅.

Rather than suffering through all of the stress, risk and pain of completely rewriting a huge system from scratch, let's take a look at a technique that I usually use when dealing with legacy code - that lets you replace a problematic system class by class, rather than having to do it all at once.

1. Pick your target

The first thing we have to do is pick what part of our app that needs refactoring. It can be a subsystem that is often causing problems and bugs, that makes implementing new features harder than it needs to be, or something that most people on the team are afraid to touch because it's so complex.

Let's say that one such subsystem in our app is the way we persist models. It consists of a ModelStorage class which in turn has many different dependencies and types that it uses for things like serialization, caching and file system access.

Rather than picking this whole system as our target, and start by rewriting ModelStorage itself - we will try to identify a single class that we can replace in isolation (that is, it doesn't have a lot of dependencies on its own). As an example, let's say we pick a Database class that ModelStorage uses to talk to our database of choice.

2. Identify the API

Exactly how our target class works under the hood is not super important. What is more important is to define what it's supposed to do by looking at its public-facing API. We'll then make a list of all the methods and properties that are not marked as private or fileprivate. For our Database class, we come up with the following:

func saveObject<O: Saveable>(_ object: O, forKey key: String) throws
func loadObject<O: Saveable>(forKey key: String) -> O?

3. Extract into a protocol

Next up, we'll take the API of our target class, and extract it into a protocol. This will later enable us to have multiple implementations of the same API, which will in turn enable us to iteratively replace the target class with a new one.

protocol Database: class {
    func saveObject<O: Saveable>(_ object: O, forKey key: String) throws
    func loadObject<O: Saveable>(forKey key: String) -> O?
}

Two things to note about the above; first is that we add the class constraint to the protocol. This is to enable us to keep doing things like keeping weak references to the type and using other class-only features, such as identity-based ones.

Secondly, we name our protocol with the exact same name as our target class. This will initially cause some compiler errors, but will later make the replacement process a lot simpler - especially if our target class is used in many different parts of our app.

4. Rename the target

Time to get rid of those compiler errors. First, let's rename our target class and clearly mark it as legacy. The way I usually do this is to simply prefix the class name with "Legacy" - so our Database class will become LegacyDatabase.

Once you perform that rename and build your project, you will still have some compiler errors left. Since Database is now a protocol, it can't be instantiated, so you will get errors like:

'Database' cannot be constructed because it has no accessible initializers

To fix this, do a find-and-replace across your whole project, replacing Database( with LegacyDatabase(. Your project should now build again like normal 👍.

5. Add a new class

Now that we have a protocol defining the expected API of our target class and we have moved the legacy implementation into a legacy class - we can start replacing it. To do that, we'll create a new class called NewDatabase, which will conform to the Database protocol:

class NewDatabase: Database {
    func saveObject<O: Saveable>(_ object: O, forKey key: String) throws {
        // Leave empty for now
    }

    func loadObject<O: Saveable>(forKey key: String) -> O? {
        // Leave empty for now
        return nil
    }
}

6. Write migration tests

Before we start implementing our replacement class with shiny new code, let's take a step back and setup a test case to help us ensure that the migration from our legacy class to our new one goes smoothly.

One big risk with all refactors is that you end up missing some detail in how the API is supposed to work, resulting in bugs & regressions. While tests won't remove all of those risks, setting tests up that run against both our legacy and new implementation will definitely make the process more robust.

Let's start by creating a test case - DatabaseMigrationTests - that has a method to perform a given test on both LegacyDatabase and NewDatabase:

class DatabaseMigrationTests: XCTestCase {
    func performTest(using closure: (Database) throws -> Void) rethrows {
        try closure(LegacyDatabase())
        try closure(NewDatabase())
    }
}

Then, let's write a test to verify that our API works as expected regardless of which implementation that is used:

func testSavingAndLoadingObject() throws {
    try performTest { database in
        let object = User(id: 123, name: "John")
        try database.saveObject(object, forKey: "key")

        let loadedObject: User? = database.loadObject(forKey: "key")
        XCTAssertEqual(object, loadedObject)
    }
}

Since we haven't implemented NewDatabase yet, the above test will fail for now. So next up is getting the test to pass by writing the new implementation in a way that is compatible with the legacy one.

7. Write the new implementation

Since NewDatabase is a completely new implementation, while still being able to be used across our whole app - just like our previous one - we are free to write it in any way we want. We can use techniques such as dependency injection, or even start using some new framework internally.

As an example, let's fill in NewDatabase with an implementation that uses JSON serialized objects stored on the file system:

import Files
import Unbox
import Wrap

class NewDatabase: Database {
    private let folder: Folder

    init(folder: Folder) {
        self.folder = folder
    }

    func saveObject<O: Saveable>(_ object: O, forKey key: String) throws {
        let json = try wrap(object) as Data
        let fileName = O.fileName(forKey: key)
        try folder.createFile(named: fileName, contents: json)
    }

    func loadObject<O: Saveable>(forKey key: String) -> O? {
        let fileName = O.fileName(forKey: key)
        let json = try? folder.file(named: fileName).read()
        return json.flatMap { try? unbox(data: $0) }
    }
}

8. Replace the legacy implementation

Now that we have a new implementation, we run our migration tests to make sure that it works just the same way as the legacy one. Once all tests pass, we are now free to replace LegacyDatabase with NewDatabase.

We'll do a find-and-replace across our project, replacing all occurrences of LegacyDatabase( with NewDatabase(. We'll also have to pass the folder: parameter in all places. Once done, we run all of our app's tests, do manual QA (for example, by shipping this version to our beta testers) to make sure that everything works well.

9. Remove the protocol

Once we're confident that our new implementation works just as well as the old one, we can safely make NewDatabase our one and only implementation. To do that, we rename NewDatabase to Database and remove the protocol called Database.

We have to do one final find-and-replace, to replace all occurrences of NewDatabase( with simply Database(, and we should now no longer have any references to NewDatabase in our project.

10. Finish up

We're almost done! All that's left is to finish up, by either removing our migration tests, or refactoring them into proper unit tests for our new implementation (depending on if our original Database class had unit tests in the first place).

If you want to keep them around, the easiest way to do so is to rename the test case to DatabaseTests, and to simply call the closure in performTest once, like this:

class DatabaseTests: XCTestCase {
    func performTest(using closure: (Database) throws -> Void) rethrows {
        try closure(Database(folder: .temporary))
    }
}

That way you don't have to rewrite or change any of your actual test methods 👌.

Finally, we can delete LegacyDatabase from our project - and we have successfully replaced a legacy class with a shiny new one - all with minimum impact and risk to the rest of our app 🎉. We can now continue using this technique to replace other parts of the ModelStorage system class by class.

Conclusion

Even though this technique is hardly a silver bullet for refactors and replacing legacy code, I think doing it this way (or some similar way) can really help reduce the risk that is usually involved in doing this kind of work.

It does require a little bit more planning up-front before starting to refactor a big system, but I still think it's worth it to do refactors iteratively like this, rather than having to rewrite everything at once.

What do you think? What are some of your favorite refactoring techniques, and do you find replacing legacy code this way to be useful? Let me know, along with any questions, comments or feedback you might have - on Twitter @johnsundell.

Finally, if you're looking for more posts to read - check out "Best of the first 6 months of Swift by Sundell", in which I list the very best content from the previous 26 posts.

Thanks for reading 🚀