Skip to content

Instantly share code, notes, and snippets.

@krackers
Forked from saagarjha/swizzler.h
Last active August 20, 2023 00:19
Show Gist options
  • Save krackers/aa0d7919e0528ee08bff07df125a7697 to your computer and use it in GitHub Desktop.
Save krackers/aa0d7919e0528ee08bff07df125a7697 to your computer and use it in GitHub Desktop.
Type-safe, RAII swizzler for Objective-C++
// https://twitter.com/_saagarjha/status/1692839886053625917#m
#import <Foundation/Foundation.h>
#import <atomic>
#import <memory>
#import <objc/runtime.h>
#include <objc/message.h>
#import <stdexcept>
#import <iostream>
#import <utility>
using namespace std;
@implementation NSObject (Convenience)
-(id)performSelector:(SEL)selector asClass:(Class)clazz
{
struct objc_super mySuper = {
.receiver = self,
.super_class = class_isMetaClass(object_getClass(self)) //check if we are an instance or Class
? object_getClass(clazz) //if we are a Class, we need to send our metaclass (our Class's Class)
: clazz //if we are an instance, we need to send our Class (which we already have)
};
auto *objc_superAllocTyped = (id (*)(struct objc_super *, SEL)) &objc_msgSendSuper; //cast our pointer so the compiler can sort out the ABI
return (*objc_superAllocTyped)(&mySuper, selector);
}
@end
template <typename Return, typename _Class = id, typename... Arguments>
class Swizzler {
using implementation_type = Return (*)(_Class, SEL, Arguments...);
struct Dispatch {
std::atomic_bool active;
std::atomic<implementation_type> original;
};
SEL _selector;
bool _moved = false;
std::shared_ptr<Dispatch> _dispatch = std::make_shared<Dispatch>();
public:
template <typename Replacement>
Swizzler(Class _class, SEL selector, Replacement replacement) : _selector(selector) {
auto method = class_getInstanceMethod(_class, selector);
if (!method) {
throw std::runtime_error("Could not find method!");
}
auto encoding = method_getTypeEncoding(method);
auto superForwarder = imp_implementationWithBlock([_class, selector](_Class self, Arguments... arguments) {
return reinterpret_cast<implementation_type>(
class_getMethodImplementation(class_getSuperclass(_class), selector))(self, selector, arguments...);
});
_dispatch->active.store(true, std::memory_order_release);
_dispatch->original.store(
reinterpret_cast<implementation_type>(class_addMethod(_class, selector, superForwarder, encoding)
? superForwarder
: method_getImplementation(method)),
std::memory_order_release);
auto& dispatch = _dispatch;
auto replacementBlock = [dispatch, replacement, selector](_Class self, Arguments... arguments) -> Return {
if (dispatch->active.load(std::memory_order_acquire)) {
return replacement(self, arguments...);
} else {
return dispatch->original.load(std::memory_order_acquire)(self, selector, arguments...);
}
};
auto _replacement = imp_implementationWithBlock(replacementBlock);
_dispatch->original.store(
reinterpret_cast<implementation_type>(class_replaceMethod(_class, selector, _replacement, encoding)),
std::memory_order_release);
}
Swizzler(Swizzler &&other) {
_selector = other._selector;
other._moved = true;
}
~Swizzler() {
if (!_moved) {
_dispatch->active.store(false, std::memory_order_release);
}
}
Swizzler(const Swizzler &) = delete;
Swizzler &operator=(const Swizzler &) = delete;
Return operator()(_Class self, Arguments... arguments) const {
return _dispatch->original.load(std::memory_order_acquire)(self, _selector, arguments...);
}
Return operator()(_Class self, SEL _cmd, Arguments... arguments) const {
return _dispatch->original.load(std::memory_order_acquire)(self, _cmd, arguments...);
}
};
int main(int argc, char *argv[]) {
Swizzler<NSString *, NSDateFormatter *, NSDate *> NSDateFormatter_stringFromDate_ {
NSDateFormatter.class, @selector(stringFromDate:), [&](NSDateFormatter * self, NSDate * date) -> NSString * {
if ([NSCalendar.currentCalendar components:NSCalendarUnitWeekday fromDate:date].weekday == 7) {
return @"It is saturday";
} else {
return NSDateFormatter_stringFromDate_(self, date);
}
}
};
NSDateFormatter *dateFormatter = [[NSDateFormatter alloc] init];
dateFormatter.dateStyle = NSDateFormatterMediumStyle;
dateFormatter.timeStyle = NSDateFormatterNoStyle;
dateFormatter.locale = [[NSLocale alloc] initWithLocaleIdentifier:@"en_US"];
std::cout << "foo" << " " << [[dateFormatter stringFromDate: [NSDate date]] UTF8String] << " bar" << "\n";
}
@krackers
Copy link
Author

#import <Foundation/Foundation.h>

#import <atomic>
#import <memory>
#import <objc/runtime.h>
#include <objc/message.h>
#import <stdexcept>
#import <iostream>
#import <utility>

using namespace std;

@implementation NSObject (Convenience)

-(id)performSelector:(SEL)selector asClass:(Class)clazz
{
	struct objc_super mySuper = {
		.receiver = self,
		.super_class = class_isMetaClass(object_getClass(self)) //check if we are an instance or Class
		? object_getClass(clazz)                //if we are a Class, we need to send our metaclass (our Class's Class)
		: clazz                                 //if we are an instance, we need to send our Class (which we already have)
	};
	
	auto *objc_superAllocTyped = (id (*)(struct objc_super *, SEL)) &objc_msgSendSuper; //cast our pointer so the compiler can sort out the ABI
	return (*objc_superAllocTyped)(&mySuper, selector);
}

@end

template <typename Return, typename _Class = id, typename... Arguments>
class Swizzler {
	using implementation_type = Return (*)(_Class, SEL, Arguments...);

	struct Dispatch {
		std::atomic_bool active;
		std::atomic<implementation_type> original;
	};

	SEL _selector;
	bool _moved = false;
	std::shared_ptr<Dispatch> _dispatch = std::make_shared<Dispatch>();

  public:
	Swizzler() {}
	template <typename Replacement>
	Swizzler(Class _class, SEL selector, Replacement replacement) : _selector(selector) {
		auto method = class_getInstanceMethod(_class, selector);
		if (!method) {
			throw std::runtime_error("Could not find method!");
		}
		auto encoding = method_getTypeEncoding(method);
		auto superForwarder = imp_implementationWithBlock([_class, selector](_Class self, Arguments... arguments) {
			return reinterpret_cast<implementation_type>(
			    class_getMethodImplementation(class_getSuperclass(_class), selector))(self, selector, arguments...);
		});
		_dispatch->active.store(true, std::memory_order_release);
		_dispatch->original.store(
		    reinterpret_cast<implementation_type>(class_addMethod(_class, selector, superForwarder, encoding)
		                                              ? superForwarder
		                                              : method_getImplementation(method)),
		    std::memory_order_release);

		auto& dispatch = _dispatch;
		auto replacementBlock = [dispatch, replacement, selector](_Class self, Arguments... arguments) -> Return {
			if (dispatch->active.load(std::memory_order_acquire)) {
				return replacement(self, arguments...);
			} else {
				return dispatch->original.load(std::memory_order_acquire)(self, selector, arguments...);
			}
		};
		auto _replacement = imp_implementationWithBlock(replacementBlock);

		_dispatch->original.store(
		    reinterpret_cast<implementation_type>(class_replaceMethod(_class, selector, _replacement, encoding)),
		    std::memory_order_release);
	}

	Swizzler(Swizzler &&other) {
		_selector = other._selector;
		other._moved = true;
	}
	Swizzler& operator=(Swizzler &&other) {
		_selector = other._selector;
		_dispatch = std::move(_dispatch);
		other._moved = true;
		return *this;
	}
	
	~Swizzler() {
		if (!_moved) {
			_dispatch->active.store(false, std::memory_order_release);
		}
	}

	Swizzler(const Swizzler &) = delete;
	Swizzler& operator=(const Swizzler &) = delete;

	Return operator()(_Class self, Arguments... arguments) const {
		return _dispatch->original.load(std::memory_order_acquire)(self, _selector, arguments...);
	}

	Return operator()(_Class self, SEL _cmd, Arguments... arguments) const {
		return _dispatch->original.load(std::memory_order_acquire)(self, _cmd, arguments...);
	}
};

Swizzler<NSString *, NSDateFormatter *, NSDate *> NSDateFormatter_stringFromDate_;

NSString* myDate(NSDateFormatter * self,  NSDate * date)  {
	if ([NSCalendar.currentCalendar components:NSCalendarUnitWeekday fromDate:date].weekday == 7) {
		NSLog(@"%@", [self performSelector:@selector(class) asClass:[NSFormatter class]]);
		return @"It is saturday";
	} else {
		return NSDateFormatter_stringFromDate_(self, date);
	}
}

template <typename RV, typename... Args>
Swizzler<RV, Args...> doSwizzle(id clazz, SEL selector, RV (*replacement)(Args...)) {
	return Swizzler<RV, Args...> {
		clazz, selector, replacement
	};
};

int main(int argc, char *argv[]) {

	NSDateFormatter *dateFormatter = [[NSDateFormatter alloc] init];
	dateFormatter.dateStyle = NSDateFormatterMediumStyle;
	dateFormatter.timeStyle = NSDateFormatterNoStyle;
	dateFormatter.locale = [[NSLocale alloc] initWithLocaleIdentifier:@"en_US"];
	
	NSDateFormatter_stringFromDate_ = doSwizzle([NSDateFormatter class], @selector(stringFromDate:), myDate);
	
	std::cout << "foo" << " " << [[dateFormatter stringFromDate: [NSDate date]] UTF8String] << " bar" << "\n";
}

@krackers
Copy link
Author

#import <Foundation/Foundation.h>

#import <atomic>
#import <memory>
#import <objc/runtime.h>
#include <objc/message.h>
#import <stdexcept>
#import <iostream>
#import <utility>

using namespace std;

@implementation NSObject (Convenience)

-(id)performSelector:(SEL)selector asClass:(Class)clazz
{
	struct objc_super mySuper = {
		.receiver = self,
		.super_class = class_isMetaClass(object_getClass(self)) //check if we are an instance or Class
		? object_getClass(clazz)                //if we are a Class, we need to send our metaclass (our Class's Class)
		: clazz                                 //if we are an instance, we need to send our Class (which we already have)
	};
	
	auto *objc_superAllocTyped = (id (*)(struct objc_super *, SEL)) &objc_msgSendSuper; //cast our pointer so the compiler can sort out the ABI
	return (*objc_superAllocTyped)(&mySuper, selector);
}

@end

template <typename Return, typename... Arguments>
class Swizzler {
	using implementation_type = Return (*)(id, SEL, Arguments...);

	struct Dispatch {
		std::atomic_bool active;
		std::atomic<implementation_type> original;
	};

	SEL _selector;
	bool _moved = false;
	std::shared_ptr<Dispatch> _dispatch = std::make_shared<Dispatch>();

  public:
	Swizzler() {}
	template <typename Replacement>
	Swizzler(Class _class, SEL selector, Replacement replacement) : _selector(selector) {
		auto method = class_getInstanceMethod(_class, selector);
		if (!method) {
			throw std::runtime_error("Could not find method!");
		}
		auto encoding = method_getTypeEncoding(method);
		auto superForwarder = imp_implementationWithBlock([_class, selector](id self, Arguments... arguments) {
			return reinterpret_cast<implementation_type>(
			    class_getMethodImplementation(class_getSuperclass(_class), selector))(self, selector, arguments...);
		});
		_dispatch->active.store(true, std::memory_order_release);
		_dispatch->original.store(
		    reinterpret_cast<implementation_type>(class_addMethod(_class, selector, superForwarder, encoding)
		                                              ? superForwarder
		                                              : method_getImplementation(method)),
		    std::memory_order_release);

		auto& dispatch = _dispatch;
		auto replacementBlock = [dispatch, replacement, selector](id self, Arguments... arguments) -> Return {
			if (dispatch->active.load(std::memory_order_acquire)) {
				return replacement(self, arguments...);
			} else {
				return dispatch->original.load(std::memory_order_acquire)(self, selector, arguments...);
			}
		};
		auto _replacement = imp_implementationWithBlock(replacementBlock);

		_dispatch->original.store(
		    reinterpret_cast<implementation_type>(class_replaceMethod(_class, selector, _replacement, encoding)),
		    std::memory_order_release);
	}

	Swizzler(Swizzler &&other) {
		_selector = other._selector;
		other._moved = true;
	}
	Swizzler& operator=(Swizzler &&other) {
		_selector = other._selector;
		_dispatch = std::move(_dispatch);
		other._moved = true;
		return *this;
	}
	
	~Swizzler() {
		if (!_moved) {
			_dispatch->active.store(false, std::memory_order_release);
		}
	}

	Swizzler(const Swizzler &) = delete;
	Swizzler& operator=(const Swizzler &) = delete;

	Return operator()(id self, Arguments... arguments) const {
		return _dispatch->original.load(std::memory_order_acquire)(self, _selector, arguments...);
	}

	Return operator()(id self, SEL _cmd, Arguments... arguments) const {
		return _dispatch->original.load(std::memory_order_acquire)(self, _cmd, arguments...);
	}
};

Swizzler<NSString *, NSDate *> NSDateFormatter_stringFromDate_;

NSString* myDate(id self,  NSDate * date)  {
	if ([NSCalendar.currentCalendar components:NSCalendarUnitWeekday fromDate:date].weekday == 7) {
		NSLog(@"%@", [self performSelector:@selector(class) asClass:[NSFormatter class]]);
		return @"It is saturday";
	} else {
		return NSDateFormatter_stringFromDate_(self, date);
	}
}

template <typename RV, typename _Class, typename... Args>
Swizzler<RV, Args...> doSwizzle(id clazz, SEL selector, RV (*replacement)(_Class, Args...)) {
	return Swizzler<RV, Args...> {
		clazz, selector, replacement
	};
};

int main(int argc, char *argv[]) {

	NSDateFormatter *dateFormatter = [[NSDateFormatter alloc] init];
	dateFormatter.dateStyle = NSDateFormatterMediumStyle;
	dateFormatter.timeStyle = NSDateFormatterNoStyle;
	dateFormatter.locale = [[NSLocale alloc] initWithLocaleIdentifier:@"en_US"];
	
	NSDateFormatter_stringFromDate_ = doSwizzle([NSDateFormatter class], @selector(stringFromDate:), myDate);
	
	std::cout << "foo" << " " << [[dateFormatter stringFromDate: [NSDate date]] UTF8String] << " bar" << "\n";
}

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment