Dynamic proxy

ZTSDynamicProxy.h

@interface ZTSDynamicProxy : NSProxy

+ (instancetype)dynamicProxyWithObject:(id)object;

@property (nonatomic, strong) id zts_object;

@end

ZTSDynamicProxy.m

static inline id ZTSDynamicProxyForObject(Class objectClass) {
    Class baseClass = [ZTSDynamicProxy class];
    
    NSString *newClassName = [NSString stringWithFormat:@"%@_%@", objectClass, baseClass];
    Class dynamicProxyClass = NSClassFromString(newClassName);
    if (dynamicProxyClass) {
        return [dynamicProxyClass alloc];
    }
    dynamicProxyClass = objc_allocateClassPair(baseClass, [newClassName UTF8String], 0);

    class_setIvarLayout(dynamicProxyClass, class_getIvarLayout(objectClass));
    class_setWeakIvarLayout(dynamicProxyClass, class_getWeakIvarLayout(objectClass));
    ZTSCopyAllIvars(objectClass, dynamicProxyClass);
    
    objc_registerClassPair(dynamicProxyClass);
    return [dynamicProxyClass alloc];
}

static inline void ZTSCopyAllIvars(Class source, Class destination) {
    unsigned int numberOfIvars = 0;
    Ivar *ivarList = class_copyIvarList(source, &numberOfIvars);
    NSLog(@"%@", destination);
    for (NSUInteger i = 0; i < numberOfIvars; ++i) {
        Ivar ivar = ivarList[i];
        const char *ivarName = ivar_getName(ivar);
        const char *ivarTypes = ivar_getTypeEncoding(ivar);
        NSUInteger ivarSize, ivarAlignment;
        NSGetSizeAndAlignment(ivarTypes, &ivarSize, &ivarAlignment);
        NSLog(@"  adding %@->%s", source, ivarName);
        BOOL didAddIvar = class_addIvar(destination, ivarName, ivarSize, (uint8_t)ivarAlignment, ivarTypes);
        NSCAssert(didAddIvar, @"Failed to add ivar %s", ivarName);
    }
    free(ivarList);
}

static inline void ZTSCopyAllIvarValues(id source, id destination, BOOL throw) {
    unsigned int nubmerOfIvars = 0;
    Ivar *ivarList = class_copyIvarList(object_getClass(source), &nubmerOfIvars);
    for (NSUInteger i = 0; i < nubmerOfIvars; ++i) {
        Ivar ivar = ivarList[i];
        void *originalValue = (__bridge void *)object_getIvar(source, ivar);
        BOOL didSetIvar = ZTSetIvar(destination, ivar, originalValue);
        if (throw) {
            NSCAssert(didSetIvar, @"Could not assign ivar %s on %@", ivar_getName(ivar), destination);
        }
    }
    free(ivarList);
}

static inline BOOL ZTSetIvar(id destination, Ivar ivar, void *originalValue) {
    ptrdiff_t ivarOffset = ivar_getOffset(ivar);
    void **ivarPosition = ((__bridge void*)destination + ivarOffset);
    *ivarPosition = originalValue;
    return YES;
}

@interface NSInvocation ()
- (void)invokeUsingIMP:(IMP)implementation;
@end

@implementation ZTSDynamicProxy

+ (instancetype)dynamicProxyWithObject:(id)object {
    ZTSDynamicProxy *instance = ZTSDynamicProxyForObject(object_getClass(object));
    instance.zts_object = object;
    return instance;
}

- (BOOL)respondsToSelector:(SEL)aSelector {
    return [self.zts_object respondsToSelector:aSelector];
}

- (NSMethodSignature *)methodSignatureForSelector:(SEL)sel {
    return [self.zts_object methodSignatureForSelector:sel];
}

- (void)forwardInvocation:(NSInvocation *)invocation {
    // Iterate over registered invocation stubs,
    // for (SEL selector in stubbedSelectors) {
    //     if (invocation.selector == selector) {
    //         void *stubbedValue = [self stubbedValueForSelector:selector];
    //         [invocation setReturnValue:&stubbedValue];
    //         return;
    //     }
    // }
    
    id target = self;
    id originalObject = self.zts_object;
    
    // Copy all ivars from the original object to the target to sync them properly
    ZTSCopyAllIvarValues(originalObject, target, YES);
    
    // Invocation itself
    invocation.target = target;
    
    Method method = class_getInstanceMethod(object_getClass(originalObject), invocation.selector);
    IMP implementaion = method_getImplementation(method);
    [invocation invokeUsingIMP:implementaion];
    
    // Copy all ivars from target to the original object to maintain a valid state
    ZTSCopyAllIvarValues(target, originalObject, NO);
}

- (void)setZts_object:(id)zts_object {
    objc_setAssociatedObject(self, @selector(zts_object), zts_object, OBJC_ASSOCIATION_RETAIN_NONATOMIC);
}

- (id)zts_object {
    return objc_getAssociatedObject(self, @selector(zts_object));
}

@end