NNBlockDelegate.h

//
//  NNBlockDelegate.h
//
//  Created by Dominique d'Argent on 1/8/13.
//  Copyright (c) 2013 Dominique d'Argent. All rights reserved.
//

#import <Foundation/Foundation.h>

@interface NNBlockDelegate : NSObject

+ (id)delegateWithOwner:(id)owner;

- (id)initWithOwner:(id)owner ;

- (BOOL)setImplementationForMethodWithSelector:(SEL)selector required:(BOOL)isRequired ofProtocol:(Protocol *)protocol usingBlock:(id)block;
- (BOOL)setImplementationForMethodWithSelector:(SEL)selector ofProtocol:(Protocol *)protocol usingBlock:(id)block;

@end

NNBlockDelegate.m

//
//  NNBlockDelegate.m
//
//  Created by Dominique d'Argent on 1/8/13.
//  Copyright (c) 2013 Dominique d'Argent. All rights reserved.
//

#import "NNBlockDelegate.h"

#import <objc/runtime.h>

static char RETAINED_SELF_KEY;


@interface NNBlockDelegate ()

@property (nonatomic, weak) NSObject *owner;

@end


@implementation NNBlockDelegate

+ (id)delegateWithOwner:(id)owner {
    return [[[self class] alloc] initWithOwner:owner];
}

- (id)initWithOwner:(id)owner {
    self = [self init];
    
    if (self) {
        [self setupWithOwner:owner];
    }
    
    return self;
}

- (BOOL)setImplementationForMethodWithSelector:(SEL)selector ofProtocol:(Protocol *)protocol usingBlock:(id)block {
    BOOL success = NO;
    
    // try optional method first
    success = success || [self setImplementationForMethodWithSelector:selector
                                                    required:NO
                                                  ofProtocol:protocol
                                                  usingBlock:block];
    
    // then try required method
    success = success || [self setImplementationForMethodWithSelector:selector
                                                    required:YES
                                                  ofProtocol:protocol
                                                  usingBlock:block];
    
    return success;
}

- (BOOL)setImplementationForMethodWithSelector:(SEL)selector required:(BOOL)isRequired ofProtocol:(Protocol *)protocol usingBlock:(id)block {
    struct objc_method_description desc = protocol_getMethodDescription(protocol, selector, isRequired, YES);
    
    BOOL success = [self setImplementationForMethodWithDescription:desc usingBlock:block];
    
    // if the method was successfully added, add the protocol to the class
    if (success && !class_conformsToProtocol([self class], protocol)) {
        class_addProtocol([self class], protocol);
    }
    
    return success;
}

#pragma mark - Private methods

- (void)setupWithOwner:(id)owner {
    NSParameterAssert(owner);
    
    self.owner = owner;
    
    // retain self
    objc_setAssociatedObject(self.owner, &RETAINED_SELF_KEY, self, OBJC_ASSOCIATION_RETAIN_NONATOMIC);
    
    // TODO: is this a save class name?
    NSString *className = [NSString stringWithFormat:@"%@_%u_%u",
                           NSStringFromClass([self class]),
                           [self hash],
                           [self.owner hash]];
    
    Class class = NSClassFromString(className);
    if (!class) {
        // create a new sublass
        class = objc_allocateClassPair([self class], [className UTF8String], 0);
        objc_registerClassPair(class);
    }
    else {
        NSLog(@"Could this ever happen?");
    }
    
    // set the isa pointer to the newly created subclass
    object_setClass(self, class);
}

- (BOOL)setImplementationForMethodWithDescription:(struct objc_method_description)desc usingBlock:(id)block {
    NSAssert(self.owner, @"%@ has no owner.", self);
    
    if (!desc.name || !desc.types) {
        return NO;
    }
    
    SEL selector = desc.name;
    const char *types = desc.types;
    
    IMP imp = imp_implementationWithBlock(block);
    
    class_replaceMethod([self class], selector, imp, types);
    
    return YES;
}

#pragma mark - Debug

- (void)dealloc {
    NSLog(@"%@ dealloced.", self);
}

@end