Python-style generators in Objective-C

gistfile1.m

@interface CDYieldEnumerator : NSEnumerator
@end

typedef BOOL(^YieldBlock)(id);
#define YIELD_BLOCK ^(YieldBlock yield)
#define YIELD(OBJECT) do { if (!yield(OBJECT)) return; } while (0)

@implementation CDYieldEnumerator {
  dispatch_queue_t _queue;
  dispatch_semaphore_t _queueLock;
  dispatch_semaphore_t _yieldLock;
  BOOL _started;

  id _nextObject;

  void (^_block)(YieldBlock);
}

- (instancetype)initWithBlock:(void(^)(YieldBlock yield))block
{
  if ((self = [super init])) {
    _queue = dispatch_queue_create([[NSString stringWithFormat:@"CDYieldEnumerator@%p", (__bridge void *)self] UTF8String], NULL);
    _block = [block copy];
    _queueLock = dispatch_semaphore_create(0);
    _yieldLock = dispatch_semaphore_create(0);
  }
  return self;
}


- (void)dealloc
{
  [self cleanUp];
}


- (void)cleanUp
{
  _queue = nil;
  _block = nil;
  dispatch_semaphore_signal(_queueLock);
}


- (void)start
{
  __weak __typeof(self) weakSelf = self;
  void (^block)(YieldBlock) = _block;
  dispatch_semaphore_t yieldLock = _yieldLock;
  dispatch_semaphore_t queueLock = _queueLock;
  dispatch_async(_queue, ^{
    block(^BOOL(id object) {
      {
        __strong __typeof (weakSelf) strongSelf = weakSelf;
        strongSelf->_nextObject = object;
      }
      dispatch_semaphore_signal(yieldLock);
      dispatch_semaphore_wait(queueLock, DISPATCH_TIME_FOREVER);
      {
        __strong __typeof(weakSelf) strongSelf = weakSelf;
        return strongSelf && strongSelf->_queue;
      }
      // return NO if work needs to be cancelled
      // another approach would be to throw an exception and catch it immediately in this method
    });
    dispatch_semaphore_signal(yieldLock);
  });
}


- (id)nextObject
{
  if (!_queue) return nil;

  _nextObject = nil;

  if (!_started) {
    _started = YES;
    [self start];
  } else {
    // Let the queue know we need the next element
    dispatch_semaphore_signal(_queueLock);
  }

  dispatch_semaphore_wait(_yieldLock, DISPATCH_TIME_FOREVER);

  if (!_nextObject) {
    [self cleanUp];
  }

  return _nextObject;
}

@end

gistfile2.m

+ (void)test
{
  NSEnumerator *enumerator = [[CDYieldEnumerator alloc] initWithBlock:YIELD_BLOCK {
    YIELD(@1);
    YIELD(@2);
    for (NSNumber *number in @[@3, @4, @5]) {
      YIELD(number);
    }
    YIELD(@6);
    [self expensiveComputation]; // Will never be executed
    YIELD(@7);
    
    [@[@8] enumerateObjectsUsingBlock:^(NSNumber *number, NSUInteger idx, BOOL *stop) {
      // This would break the generator since YIELD can call return
      // A possible workaround is exceptions, but they would be slower and create other problems (auto breakpoints, mem management)
      YIELD(number);
    }];
  }];
  
  for (int i = 0; i < 6; i++) NSLog([enumerator nextObject]); // Logs 1 2 3 4 5 6
}