Optimizing Swift: Tracking property changes and building a Memoization system

As a Swift consultant, I always look for ways to optimize clients’ code and improve efficiency. I recently built a prototype for The Browser Company that tracks changes in Swift properties and implements a memoization system.

This system re-executes complex computation scopes only when necessary, avoiding redundant re-computations and returning cached values when no property changes have been made.

In this article, I’ll walk you through the inner workings of this prototype and demonstrate how it can enhance your Swift coding experience.

Here’s a final demo from this series where you can see that it will re-compute just the selected Scopes depending on which properties in State I mutate.

Foundations

Let’s start by introducing the foundational components of our prototype: the ChangeCounter, ChangeID, and MemoizableState types. The ChangeCounter is a simple struct containing a public static property named “iteration”:

public struct ChangeCounter {
public static var iteration: UInt64 = 0
}

This UInt64 number increments with each new iteration of our memoization system, allowing us to keep track of iterations and know which iteration we’re working with at any given moment.

On the other hand, the ChangeID struct is responsible for generating unique identifiers for our property tracking needs:

public struct DirtyCheckingID: Hashable {
static var counter: UInt64 = 0
var value: UInt64

init() {
    value = Self.counter
    Self.counter += 1
}
}

With a self-incrementing counter, each instance of ChangeID is assigned a distinct value, allowing seamless tracking and debugging throughout the code.

public protocol MemoizableState {
  func lastUpdate(for id: ChangeID) -> UInt64
}

MemoizableState is a protocol used to keep track of lastUpdate for a given property id. This is required due to the transient nature of PropertyWrappers.

ℹ️

By transient nature, I mean that PropertyWrappers are not carried through like standard types but are syntactic sugar, which complicates passing and re-computing data using them.

@ChangeTracking Property Wrapper

Now that we’ve covered the basics let’s discuss the ChangeTracking property wrapper. This wrapper keeps track of when a property was last updated and provided an efficient way to access it. It uses the ChangeID for identification and  ChangeCounter.iteration to store the last updated iteration:

@propertyWrapper public struct ChangeTracking<Value> {
    public let id = ChangeID()
    public var lastUpdatedIteration: UInt64 = 0

    var _wrappedValue: Value {
        didSet {
            lastUpdatedIteration = DirtyCounter.iteration
        }
    }

    // ...
}

When the value of the wrapped property changes, the didSet observer updates the lastUpdatedIteration. The wrapper also includes methods for accessing child properties without registering them for dirty access, allowing us to filter update triggers for the memoized blocks but more about that later.

Memoizer

The heart of the memoization system lies within the Memoizer class, which uses the ChangeTracking property wrapper and the ChangeCounter struct to determine if a computation scope needs to be re-executed or not:

public class Memoizer {
// ...

@discardableResult public func memoized<State: MemoizableState, Value>(state: State, scopeID: ScopeID, memoize: (State) -> Value) -> MemoizationResult<Value> {
    // ...
}
}

The memoized the function compares the current state’s last update for each property with the property’s previous update from the memoization scope. If any properties have changed, the function executes the memoization block and updates the scope with the new value. If not, it returns the previously cached value without re-executing the block. This optimizes the overall performance and avoids unnecessary computations.

🤔

Explicit scopeID is helpful for debugging, but you can leverage the fact that \(#file#line) is always going to be unique scopeID in the codebase.

func memoized<State: MemoizableState, Value>(state: State, file: StaticString = #file, line: UInt = #line, memoize: (State) -> Value) -> MemoizationResult<Value> {
memoized(state: state, scopeID: "\(file):\(line)", memoize: memoize)
}

Memoized function

Let’s examine the memoized function within the Memoizer class in more detail.

@discardableResult public func memoized<State: MemoizableState, Value>(state: State, scopeID: ScopeID, memoize: (State) -> Value) -> MemoizationResult<Value> {
// 1. Check if the scope already exists
guard let scope = scopes[scopeID] else {
    // 2 Create a new scope and add it to the active scopes
    activeScopes.append(.init(id: scopeID))
    // 3 Mark the scope as hot (re-computed)
    hotScopes.insert(scopeID)
    // 4 Execute the memoize closure and cache the result
    let value = memoize(state)
    // 5 Update the scope with the new value
    updateScope(scope: activeScopes.removeLast(), with: value)
    return .init(value: value, updated: true)
}
// ...
}

1. Check if the scope already exists: The function starts by checking if the scope exists in the scopes dictionary using the provided scopeID.
2. Create a new scope and add it to the active scopes: If the scope doesn’t exist, a new scope is created and added to the activeScopes list.
3. Mark the scope as hot (re-computed): The new scope is marked as hot by inserting its scopeID into the hotScopes set, indicating that it was re-computed.
4. Execute the memoize closure and cache the result: The memoize closure is executed, and the result is cached in the new scope.
5. Update the scope with the new value: The updateScope function is called to update the scope with the new value, including the updated lastUpdateCounter value.

In this way, the memoized function ensures that the memoized block is only executed when necessary, avoiding redundant computations and enhancing the performance of your code.

The next thing we want to do is to avoid re-executing scopes in the same mutation iteration (reducer run in TCA):

// 1.Check if the scope has already been executed in the current iteration
if scope.lastUpdateCounter == ChangeCounter.iteration {
// 2 Update the dependency graph
activeScopes.last?.sourceScopes.insert(scope)
scope.derivedScope = activeScopes.last
return .init(value: scope.lastValue as! Value, updated: false)
}

1. Check if the scope has already been executed in the current iteration: If the scope’s lastUpdateCounter is equal to the current iteration of ChangeCounter, the scope has already been executed, and there’s no need to re-execute it.
2. Update the dependency graph: In this case, the function updates the dependency graph between the scopes by adding the current scope to the sourceScopes of the last active scope and setting the derivedScope of the current scope to the last active scope. The function then returns the cached value without executing the memoized block again.

Next, we want to do is verify if the scope needs to be re-computed by checking whether any of the properties have changed from the last run:

let needsToRun = scope.lastUpdateCounter == 0 || scope.allSourceProperties.contains(where: { config in
state.lastUpdate(for: config.id) > config.lastUpdate
})

The final part is to either re-compute or return a cached value. We can do it as follows:

guard needsToRun else {
  return .init(value: scope.lastValue as! Value, updated: false)
}
// 1 Mark the scope as hot
hotScopes.insert(scopeID)
// 2 Add the scope to the active scopes list
activeScopes.append(scope)
// 3 Reset the source properties for the scope
scope.sourceProperties.removeAll()
// 4 Execute the memoize closure and cache the result
let value = memoize(state)
// 5 Remove the scope from the active scopes list
_ = activeScopes.removeLast()
// 6 Update the scope with the new value
updateScope(scope: scope, with: value)
return .init(value: value, updated: true)

Using the system

Now that we understand the underlying components of our memoization system let’s see how to use it in practice. To demonstrate, let’s assume we have a MemoizableState conforming class called MyState:

struct MyState: MemoizableState {
    @ChangeTracking var name: String
    @ChangeTracking var age: Int

    func lastUpdate(for id: ChangeID) -> UInt64 {
        if id == _name.id {
            return _name.lastUpdatedIteration
        } else if id == _age.id {
            return _age.lastUpdatedIteration
        }
        return 0
    }
}

Here, MyState has two properties, name and age, both wrapped with ChangeTracking. The lastUpdate(for:) function returns the last updated iteration for each property (This boilerplate can easily be automated with Sourcery).

With this setup, you can use the Memoizer class to memoize functions based on the state of these properties:

let memoizer = Memoizer.shared
let state = MyState(name: "Alice", age: 30)

let result = memoizer.memoized(state: state, scopeID: "someScopeID") { state -> String in
    print("Executing memoized block")
    return "Name: \(state.name), Age: \(state.age)"
}

The memoized block will only be executed if the name or age properties have changed since the last time it was called, thus avoiding unnecessary re-execution.

Conclusion

Potential benefits

The prototype described in this article offers a powerful way to optimize your Swift code by tracking property changes and memoizing computation scopes. This system can minimize your app’s redundant computations, improving its performance and responsiveness.

Drawbacks

This system adds complexity, and using property wrappers makes it harder for a larger team to maintain.

It could be helpful as one of the tools in your toolbox, but use it with care.

ℹ️

Next time we’ll look at further improvements to the code and how you could integrate it inside TCA and visualize the scope changes.