Tip: How To weakly reference self in protocol extensions that are not class-bound

As any iOS developer can tell you, it's often important to use weak in closures to break refernce cycles and avoid memory leaks. Unfortunately, using weak is not always directly possible in protocol extensions that are not class-bound. Suppose you have a protocol for an API client to fetch some information:

protocol APIClient {
    func fetchData() async throws -> Data
}

Unfortunately, not all of your code is ready to rely on async/await exclusively for concurrency so you may want to also expose a Combine compatible interface:

protocol APIClient {
    func fetchData() async throws -> Data
    func fetchDataPublisher() -> Future<Data, Error>
}

At a glance, it's obvious these two protocol requirements are redundant and conforming types shouldn't have to provide two implementations to solve the same problem. We can solve this by providing a default implementation of the Combine interface:

extension APIClient {
    func fetchDataPublisher() -> Future<Data, Error> {
        return Future { promise in
            Task {
                do {
                    let data = try await fetchData()
                    promise(.success(data))
                } catch {
                    promise(.failure(error))
                }
            }
        }
    }
}

Unfortunately, the above may cause unexpected behavior. Even if implementations of APIClient choose to implement fetchData in a way that does not strongly reference self, the Combine compatible implementation will still retain self for the lifetime of the API call.

Unfortunately, it's not possible to simply add [weak self] into the closure because we cannot guarantee that self is a reference type that can be referenced weakly. Instead we get a build error suggesting we limit the protocol to a class:

'weak' must not be applied to non-class-bound 'Self'; consider adding a protocol conformance that has a class bound

While this may be a perfectly good general solution, it's not a good solution here. There's no fundamental difference in the logic we want to use for reference types vs value types. So we need a workaround to weakly reference classes that won't cause errors value types.

Fortunately, the solution is quite simple:

struct WeakBox<T> {
    var object: T? {
        get { storage as? T }
        set { storage = newValue as AnyObject }
    }
    weak private var storage: AnyObject?
    
    init(_ stored: T) {
        self.object = stored
    }
}

WeakBox casts the stored value to AnyObject so we can always use reference semantics, even if the provided object is actually a value type. (We can unconditionally cast as AnyObject because of the some Objective C interoperability requirements.)

Fortunately, the casting a value type as AnyObject does not cause it to be deallocated, as proved by the following pair of unit tests.

final class WeakCaptureTests: XCTestCase {
    func test_doesNotRetainClass() throws {
        class Example {}
        let subject = WeakBox(Example())
        XCTAssertNil(subject.object)
    }
    
    func test_doesRetainStructs() throws {
        let subject = WeakBox(42)
        XCTAssertEqual(subject.object, 42)
    }
}

WeakBox allows us to appropriately leave the decision of whether or not to retain self up to implementations of APIClient rather than our default implementation without duplicating the implementation or restricting it to AnyObject:

struct APIClientDeallocated: Error {}
extension APIClient {
    func fetchDataPublisher() -> Future<Data, Error> {
        let weakReference = WeakBox(self)
        return Future {  promise in
            Task {
                do {
                    guard let data = try await weakReference.object?.fetchData() else {
                        throw APIClientDeallocated()
                    }
                    promise(.success(data))
                } catch {
                    promise(.failure(error))
                }
            }
        }
    }
}

Hopefully this has been interesting and useful. Please feel free to reach out on LinkedIn if you have any questions or comments.