Defer in Objective-C

Some time ago I was learning Go and I came across a very useful feature: defer. As its name suggests it defers the execution of a piece of code until the current function returns.

defer is very useful when dealing with resources or locks. For example suppose you have a class with an array of people that can be accessed from different threads. You need to use a lock to avoid modifying the array while it's being used by another thread. Here is some code you might end up writing:

- (Person*)personAtIndex:(NSUInteger)index {
    [_lock lock];
    Person* person = [_array objectAtIndex:index];
    [_lock unlock];
    return person;
}

- (BOOL)updatePerson:(Person*)person {
    [_lock lock];

    NSUInteger index = [_array indexOfObject:person];
    if (index == NSNotFound) {
        [_lock unlock];
        return NO;
    }

    Person* oldPerson = [_array objectAtIndex:index];
    [oldPerson markDeleted];

    [_array removeObjectAtIndex:index];
    [_array addObject:person];
    [_lock unlock];
    return YES;
}

Notice how calls to lock and unlock are spread all over, not very maintainable. Also, it would be easy to forget calling unlock if a new branch of code is needed.

Now let's rewrite this code using defer:

- (Person*)personAtIndex:(NSUInteger)index {
    [_lock lock];
    defer(^() {
        [_lock unlock];
    });

    return [_array objectAtIndex:index];
}

- (BOOL)updatePerson:(Person*)person {
    [_lock lock];
    defer(^() {
        [_lock unlock];
    });

    NSUInteger index = [_array indexOfObject:person];
    if (index == NSNotFound)
        return NO;

    Person* oldPerson = [_array objectAtIndex:index];
    [oldPerson markDeleted];

    [_array removeObjectAtIndex:index];
    [_array addObject:person];
    return YES;
}

Now all lock invocations are at the top of the methods. Much nicer.

The Code

To implement defer in Objective-C we take advantage of weak pointers. Weak pointers are set to nil and their referenced object released when there are no strong references to the object. But the trick is that this only happens once per run loop. This means that the current method and it's callers will finish executing before the weak object is released. The only thing we have to do is create an object that calls a block on dealloc:

@implementation AIDefer

+ (instancetype)defer:(void (^)())block {
    AIDefer* defer =  [[AIDefer alloc] init];
    defer.block = block;
    return [defer autorelease];
}

- (void)dealloc {
    if (_block)
        _block();
    [super dealloc];
}

@end

Notice that I'm not using ARC, this is because if you activate ARC the compiler will optimize away the autorelease and release the object immediately. Now we create a helper function that creates instances of this class with the given block and keeps a weak reference to the object

void defer(void (^block)()) {  
    static AIDefer* __weak d;
    d = [AIDefer defer:block];
}

Take a look at the repository and start using the code. If you use CocoaPods You can also add pod 'AIDefer' to your Podfile.