// inspired by http://www.chiark.greenend.org.uk/~sgtatham/coroutines.html

#import <Foundation/Foundation.h>


int (^Primes(void))(void)
{
    __block int n = 2;
    __block int i;
    __block BOOL isPrime;
    __block int where = 0;
    return [[^{
        switch(where)
        {
            case 0:
                while(1)
                {
                    isPrime = YES;
                    for(i = 2; i < n; i++)
                        if(n % i == 0)
                            isPrime = NO;
                    if(isPrime)
                    {
                        where = 1;
                        return n;
            case 1:
                        ;
                    }
                    n++;
                }
            default:
                return 0;
        }
    } copy] autorelease];
}

@interface MACoroutineCleanup : NSObject
{
    void (^_block)(void);
}
@property (copy, nonatomic) void (^block)(void);
@end
@implementation MACoroutineCleanup
@synthesize block = _block;
- (id)retain
{
    return [super retain];
}
- (void)release
{
    [super release];
}
- (void)dealloc
{
    if(_block)
        _block();
    [_block release];
    [super dealloc];
}
@end

#define COROUTINE(name, ...) (^name(__VA_ARGS__))(void)
#define COROUTINE_BEGIN \
    __block int COROUTINE_where = -1; \
    MACoroutineCleanup *COROUTINE_cleanupObj = [[MACoroutineCleanup alloc] init]; \
    id mainBlock = ^{ \
        [COROUTINE_cleanupObj self]; \
        switch(COROUTINE_where) \
        { \
            case -1:
#define YIELD(...) \
                do { \
                    COROUTINE_where = __LINE__; \
                    return __VA_ARGS__; \
                    case __LINE__: \
                    ; \
                } while(0)
#define COROUTINE_DUMMY_RETURN(x) \
            default: \
                return x;
#define COROUTINE_CLEANUP \
        } \
    }; \
    COROUTINE_cleanupObj.block = ^{{
#define COROUTINE_END \
        } \
    }; \
    mainBlock = [[mainBlock copy] autorelease]; \
    [COROUTINE_cleanupObj release]; \
    return mainBlock;

int COROUTINE(CoPrimes, void)
{
    __block int n = 2;
    __block int i;
    __block BOOL isPrime;
    
    COROUTINE_BEGIN
        while(1)
        {
            isPrime = YES;
            for(i = 2; i < n; i++)
                if(n % i == 0)
                    isPrime = NO;
            if(isPrime)
                YIELD(n);
            n++;
        }
        COROUTINE_DUMMY_RETURN(0);
    COROUTINE_END
}

int COROUTINE(CoPrimesMultipleYield, void)
{
    __block int n;
    __block int i;
    __block BOOL isPrime;
    
    COROUTINE_BEGIN
        YIELD(2);
        YIELD(3);
        YIELD(5);
        YIELD(7);
        YIELD(11);
        YIELD(13);
        n = 14;
        while(1)
        {
            isPrime = YES;
            for(i = 2; i < n; i++)
                if(n % i == 0)
                    isPrime = NO;
            if(isPrime)
                YIELD(n);
            n++;
        }
        COROUTINE_DUMMY_RETURN(0);
    COROUTINE_END
}

int COROUTINE(Counter, int start, int end)
{
    __block int i;
    COROUTINE_BEGIN
        for(i = start; i <= end; i++)
            YIELD(i);
        while(1)
            YIELD(-1);
        COROUTINE_DUMMY_RETURN(0);
    COROUTINE_END
}

int COROUTINE(CleanupTest, void)
{
    __block NSString *str = nil;
    __block int i = 0;
    COROUTINE_BEGIN
        for(;;)
        {
            [str release];
            str = [[NSString alloc] initWithFormat: @"STRING: %d", ++i];
            YIELD(i);
        }
        COROUTINE_DUMMY_RETURN(0);
    COROUTINE_CLEANUP
        NSLog(@"Cleaning up, string is %@", str);
        [str release];
    COROUTINE_END
}

int main(int argc, char **argv)
{
    NSAutoreleasePool *pool = [NSAutoreleasePool new];
    
    int (^primes)(void) = CoPrimesMultipleYield();
    
    int i;
    for(i = 0; i < 10; i++)
        NSLog(@"%d", primes());
    
    int (^counter)(void) = Counter(5, 10);
    for(i = 0; i < 10; i++)
        NSLog(@"%d", counter());
    
    int (^cleanupTest)(void) = CleanupTest();
    for(i = 0; i < 5; i++)
        NSLog(@"%d", cleanupTest());
    
    [pool release];
    
    return 0;
}
