compose.cpp

#include <iostream>
#include <cstring>
#include <sys/mman.h>

extern "C" {
	extern const char compose_helper_nomem[];
	extern const size_t compose_helper_nomem_size;
	extern const char compose_helper_mem[];
	extern const size_t compose_helper_mem_size;

	__asm__(
		"compose_helper_nomem:\n"
		"endbr64\n"
		"mov %rsi, %rcx\n"
		"mov %rdi, %rdx\n"
		"mov compose_helper_nomem_size(%rip), %rdi\n"
		"mov compose_helper_nomem_size+8(%rip), %rsi\n"
		"jmp *compose_helper_nomem_size+16(%rip)\n"
		"compose_helper_nomem_size:\n"
		".8byte .-compose_helper_nomem\n"
		"compose_helper_mem:\n"
		"endbr64\n"
		"mov %rdx, %r8\n"
		"mov %rsi, %rcx\n"
		"mov compose_helper_mem_size(%rip), %rsi\n"
		"mov compose_helper_mem_size+8(%rip), %rdx\n"
		"jmp *compose_helper_mem_size+16(%rip)\n"
		"compose_helper_mem_size:\n"
		".8byte .-compose_helper_mem"
	);
}

template <typename A, typename B, typename C>
C compose_at_once(C (*outer)(B), B (*inner)(A), A input) {
	return outer(inner(input));
}

template <typename A, typename B, typename C>
C (*compose(bool c_is_mem, C (*outer)(B), B (*inner)(A)))(A) {
	const char *compose_helper = c_is_mem ? compose_helper_mem : compose_helper_nomem;
	size_t compose_helper_size = c_is_mem ? compose_helper_mem_size : compose_helper_nomem_size;
	C (*compose_at_once_ptr)(C (*)(B), B (*)(A), A) = compose_at_once<A, B, C>; // hack because you can't take the address of an address
	C (*mem)(A) = reinterpret_cast<C (*)(A)>(mmap(nullptr, compose_helper_size + 3*sizeof(void *), PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_ANONYMOUS, -1, 0));
	char *cmem = reinterpret_cast<char *>(mem);
	std::memcpy(cmem, compose_helper, compose_helper_size);
	std::memcpy(cmem + compose_helper_size, &outer, sizeof(void *));
	std::memcpy(cmem + compose_helper_size + sizeof(void *), &inner, sizeof(void *));
	std::memcpy(cmem + compose_helper_size + 2*sizeof(void *), &compose_at_once_ptr, sizeof(void *));
	mprotect(cmem, compose_helper_size + 3*sizeof(void *), PROT_READ|PROT_EXEC);
	return mem;
}

typedef struct {
	long x;
	long y;
} z2_t;

typedef struct {
	long x;
	long y;
	long z;
} z3_t;

template <typename T>
T triple(T x) {
	return 3 * x;
}

template <typename T>
T add_ten(T x) {
	return x + 10;
}

z2_t halve_x2(z2_t in) {
	z2_t retval;
	retval.x = in.x/2;
	retval.y = in.y;
	return retval;
}

z3_t halve_x3(z3_t in) {
	z3_t retval;
	retval.x = in.x/2;
	retval.y = in.y;
	retval.z = in.z;
	return retval;
}

int main(void) {
	int (*triple_then_add_ten_i)(int) = compose(false, add_ten<int>, triple<int>);
	std::cout << triple_then_add_ten_i(5) << std::endl;
	float (*triple_then_add_ten_f)(float) = compose(false, add_ten<float>, triple<float>);
	std::cout << triple_then_add_ten_f(5.1f) << std::endl;
	z2_t (*quarter_x2)(z2_t) = compose(false, halve_x2, halve_x2);
	z2_t in2 = {400,200};
	z2_t out2 = quarter_x2(in2);
	std::cout << out2.x << " " << out2.y << std::endl;
	z3_t (*quarter_x3)(z3_t) = compose(true, halve_x3, halve_x3);
	z3_t in3 = {400,200,300};
	z3_t out3 = quarter_x3(in3);
	std::cout << out3.x << " " << out3.y << " " << out3.z << std::endl;
	return 0;
}