How to avoid implicit retain cycles when using function references

One very nice thing about Swift is that functions are first-class values. They can be assigned to variables, curried, and you can pass them as arguments: If a function takes a closure as a parameter, you can pass a reference to a function in place of an in-line closure, as long as the types match. We may think of functions as just special kinds of closures with a name. But beware! Therein lies dragons.

But first, let’s look at functions first class values! The following two are equivalent:

let numbers = [1, 2, 3, 4, 5, 6, 7, 8]

func isEven(number: Int) -> Bool {
    return number % 2 == 0
}

let even = numbers.filter({ $0 % 2 == 0 }) // in-line closure
let alsoEven = numbers.filter(isEven) // function reference

I really like this pattern, and use it everywhere. In an app I’m working on, we’re keeping an open socket for real time communication with the backend. The socket object is an event emitter and subscriber. That is, you can register closures to be executed when certain events are received. This lends itself nicely to the functional approach:

func createSession() {
    socket.emit("start_session")
    socket.once("session_started") { [unowned self] in
        let eventHandlers = [
            socket.on("chat_message_received", callback: self.onChatMessageReceived),
            socket.on("client_state_changed", callback: self.onClientStateChanged),
            //...
        ]
        socket.once("session_ended") {
            eventHandlers.each(socket.off)
        }
    }
}
func onChatMessageReceived(message: ChatMessage) {
    //handle chat message
}
func onClientStateChange(state: ClientState) {
    //handle client state
}
//...

However, there’s a problem with the code above. Can you spot the bug?

The problem

That’s right! It has a retain cycle. Even though we capture self unonwed in the outer closure, the callbacks that gets registered with the socket, are captured strongly. So even when the object above is released from where ever it was created, the object is still kept alive. OH NOES!

At best it will take up unrecoverable memory, but at worst it will keep active event handlers around, and interfere with the network traffic. This may cause strange and seemingly unpredictable bugs. Let’s try to solve it.

A naïve solution

A naïve approach to solving this, is to inline the functions and pass them as closures instead:

func createSession() {
    socket.emit("start_session")
    socket.once("session_started") { [unowned self] in
        let eventHandlers = [
            socket.on("chat_message_received") { self.onChatMessageReceived($0) },
            socket.on("client_state_changed") { self.onClientStateChanged($0) },
            //...
        ]
        socket.once("session_ended") {
            eventHandlers.each(socket.off)
        }
    }
}

Above, we’ve solved the problem by wrapping the functions in in-line closures. Since the closures captures self unowned, the cycles are broken:

In the case above self is already unowned in the current scope, and every function only takes one argument. Thus, the overhead isn’t too bad.

But what if that wasn’t the case? Notice how onChatMessageReceived(_:) and onClientStateChanged(_:) have different signatures? Even though they’re similar, the event subscriber in our socket object is a generic function, that will automatically deserialize objects that conform to a certain protocol. Only JSON is sent over the socket, but we have added a generic method that parses the JSON into proper model objects.

But we also have a “raw” version of the function, that returns the event, a raw data object and an ack class, that can be used to acknowledge the event. Consider this:

socket.on("chat_message_received") { [unowned self] in self.onChatMessageReceived($0, data: $1, ack: $2) }
socket.on("client_state_changed") { [unowned self] in self.onClientStateChanged($0, data: $1, ack: $2) }
socket.on("some_other_event") { [unowned self] in self.onOtherEvent($0, data: $1, ack: $2) }
socket.on("even_more_events") { [unowned self] in self.onMoreEvents($0, data: $1, ack: $2) }
//...

It’s a mess!

The [unowned self] in and argument list is just boilerplate, repeated over, with no other purpose than to pass arguments from one function (the in-line closure) on to the next (the actual event handler functions). Is there a way of wrapping the closures above in a helper method that:

• return a function with the correct signature
• captures self unowned
• is generic over argument list and return values?

Instance functions are curried class functions!

In Swift, instance functions are just curried class functions that takes an instance as its first argument, and implicitly make that first argument available to the function body as self. Thus, the following two are equivalent:

let numbers = [1, 2, 3, 4, 5, 6, 7, 8]
numbers.contains(3) //true
Array.contains(numbers)(3) //true

Also, these are equivalent:

let eventHandler1 = self.onChatMessageReceived
let eventHandler2 = self.dynamicType.onChatMessageReceived(self)

And of course we already have this beast:

let eventHandler3 = { [unowned self] in self.onChatMessageReceived($0, data: $1, ack: $2) }

Functional-generic memory management voodoo

What if we took self.dynamicType.onChatMessageReceived from the second event handler above, but without the argument (self), and passed that as a parameter to our wrapper function, together with a reference to self? Maybe we could then capture that reference unowned, and pass the unowned instance reference to the class function to get an instance function, without creating a retain cycle. Phew! Let’s try!

func unown<T: AnyObject, U, V>(instance: T, _ classFunction: T->U->V) -> U->V {
    return { [unowned instance] args in
        let instanceFunction = classFunction(instance)
        return instanceFunction(args)
    }
}

Now, this gives us a fourth way of accessing a function reference, that are similar to the three above. But where the first two captured self strongly, and the third had a lot of boiler plate and explicit argument lists, this last ones does not:

let eventHandler4 = unown(self, self.dynamicType.onChatMessageReceived)

Since this is generic, Swift’s type inference system will do the right thing, and make sure it can be used everywhere you need to pass a reference to a function as an argument, and avoid creating retain cycles, regardless if the closure expects a JSON in NSData, a deserialized model object, or even an NSNotification:

class SomeClass {
    var observer: AnyObject?
    init() {
        observer = NSNotificationCenter.defaultCenter().addObserverForName(
            "notification",
            object: nil,
            queue: NSOperationQueue.mainQueue(),
            usingBlock: unown(self, SomeClass.notificationHandler) // avoid ref. cycle!
        )
    }
    deinit {
        NSNotificationCenter.defaultCenter().removeObserver(observer)
    }
    func notificationHandler(note: NSNotification) {
        //handle notification
    }
}

Swift evolution?

Although this works, it is far from optimal. I think this patterns is common enough to warrant language support for it. In fact, I most often want this behavior and would like function references to have unowned memory capture semantics by default. But ideally, we should have some kind of keyword that let us use something like this:

let eventHandler5 = @unowned(self).onChatMessageReceived