Proof of concept: reference counted closures

rc_closure.d

/**
 * Reference counted closures.
 *
 * version = with_new_trait supports postblitting and destruction of upvalues,
 * but depends on a hypothetical new language-level trait.
 *
 * The new trait, __traits(upvalues, func), would take an alias to a function
 * and return an alias sequence containing a description of its upvalues:
 *
 * (T1, ptrdiff_t offset1, T2, ptrdiff_t offset2, ...)
 *
 * Where each pair is the type of the upvalue and the upvalue's offset from
 * the context pointer.
 */
module rc_closure;

import std.experimental.allocator : IAllocator;

// TODO: simple change that templatizes the type of Payload.dispose,
// so the RefCountedClosure destructor can propagate attributes.

version = proof_of_concept;
// version = with_new_trait;

struct RefCountedClosure(F)
	if(is(F == delegate))
{
	import std.traits : ReturnType;

	private:
	struct Payload
	{
		uint count;
		void* context; // Needed because offsets can be both positive and negative
		typeof(F.init.funcptr) funcPtr;
		IAllocator allocator;
		version(with_new_trait)
			void function(Payload*) dispose;
		else
			size_t contextSize;
	}

	Payload* payload = null;

	this(Payload* payload)
	{
		this.payload = payload;
	}

	F asDelegate() @property @trusted
	{
		F dg;
		dg.funcptr = payload.funcPtr;
		dg.ptr = payload.context;
		return dg;
	}

	public:
	this(Other)(RefCountedClosure!Other other)
	{
		this.opAssign(other);
	}

	this(this)
	{
		if(payload)
			++payload.count;
	}

	void opAssign(typeof(this) rhs)
	{
		import std.algorithm.mutation : swap;
		swap(this.payload, rhs.payload);
	}

	void opAssign(Other)(RefCountedClosure!Other rhs)
		if(is(Other : F))
	{
		auto thisPayload = this.payload;
		auto rhsPayload = rhs.payload;

		this.payload = cast(Payload*)rhsPayload;
		rhs.payload = cast(RefCountedClosure!Other.Payload*)thisPayload;
	}

	~this()
	{
		if(payload && --payload.count == 0)
		{
			version(with_new_trait)
				payload.dispose(payload);
			else
			{
				immutable deallocSize = Payload.sizeof + payload.contextSize;
				payload.allocator.deallocate(() @trusted {
					return (cast(void*)payload)[0 .. deallocSize];
				}());
			}
		}
	}

	///
	auto ref ReturnType!F opCall(Args...)(auto ref Args args)
		if(is(typeof(F.init(args))))
	{
		return asDelegate()(args);
	}

	///
	uint refCount() @property
	{
		return payload? payload.count : 0;
	}
}

version(proof_of_concept)
RefCountedClosure!F rcClosure(F)(scope F dg)
{
	import std.experimental.allocator : allocatorObject;
	import std.experimental.allocator.mallocator : Mallocator;
	return rcClosure(dg, allocatorObject(Mallocator.instance));
}

version(proof_of_concept)
RefCountedClosure!F rcClosure(F)(scope F dg, IAllocator allocator)
	if(is(F == delegate))
{
	import core.exception : onOutOfMemoryError;

	alias Payload = typeof(return).Payload;
	enum contextSize = 256; // Arbitrary size

	if(auto chunk = allocator.allocate(Payload.sizeof + contextSize))
	{
		auto payload = cast(Payload*)chunk.ptr;
		*payload = Payload(0, payload + 1, dg.funcptr, allocator, contextSize);

		// Copy context
		auto source = cast(ubyte*)dg.ptr - contextSize / 2;
		auto destination = cast(ubyte*)payload.context;
		destination[0 .. contextSize] = source[0 .. contextSize];

		payload.context += contextSize / 2;

		return typeof(return)(payload);
	}
	else
	{
		onOutOfMemoryError();
		assert(false);
	}
}

version(with_new_trait)
{
	// TODO: make the offsets runtime arguments, further reducing bloat
	private void destroyUpvalues(upvalues...)(void* base)
	{
		foreach_reverse(i, offset; upvalues)
		{
			static if(i % 2 != 0)
			{
				alias T = upvalues[i - 1];
				static if(hasElaborateDestructor!T)
				{
					auto p = () @trusted {
						return cast(T*)(base + offset);
					}();
					destroy(*p);
				}
			}
		}
	}

	private void copyUpvalues(upvalues...)(void* source, void* destination)
	{
		import std.conv : emplaceRef;

		foreach(i, T; upvalues)
		{
			static if(i % 2 == 0)
			{
				enum offset = upvalues[i + 1];
				auto typedSource = () @trusted {
					return cast(T*)(source + offset);
				}();
				auto typedDestination = () @trusted {
					return cast(T*)(destination + offset);
				}();
				emplaceRef(*typedDestination, *typedSource);
			}
		}
	}
}

version(with_new_trait)
private RefCountedClosure!F rcClosureImpl(F, upvalues...)(
	IAllocator allocator,
	scope F dg)
{
	alias Payload = typeof(return).Payload;
	static if(Upvalues.length)
	{
		enum mostNegativeOffset = upvalues[1];
		enum mostPositiveOffset = upvalues[$ - 1];
		enum contextSize = abs(mostPositiveOffset - mostNegativeOffset);
	}
	else
		enum contextSize = 0;

	if(auto chunk = allocator.allocate(Payload.sizeof + contextSize))
	{
		auto payload = cast(Payload*)chunk.ptr;
		payload.count = 1;
		payload.funcPtr = dg.funcptr;
		payload.allocator = allocator;
		payload.dispose = (payload) {
			destroyUpvalues!upvalues(payload.context.ptr);
			enum deallocSize = Payload.sizeof + contextSize;
			payload.allocator.deallocate(() @trusted {
				return (cast(void*)payload)[0 .. deallocSize];
			}());
		};
		payload.context = cast(void*)(payload + 1) + abs(mostNegativeOffset);
		copyUpvalues!upvalues(dg.ptr, payload.context);
		return typeof(return)(payload);
	}
	else
		onOutOfMemoryError();
}

///
version(with_new_trait)
RefCountedClosure!F rcClosure(alias func)(IAllocator allocator)
	if(is(typeof(&func) == delegate))
{
	return rcClosureImpl!(typeof(&func), __traits(upvalues, func))(allocator, &func);
}

///
version(with_new_trait)
RefCountedClosure!F rcClosure(alias func)()
	if(is(typeof(&func) == delegate))
{
	import std.experimental.allocator : allocatorObject;
	import std.experimental.allocator.mallocator : Mallocator;
	return rcClosure!func(allocatorObject(Mallocator.instance));
}

version(proof_of_concept)
unittest
{
	{
		immutable a = 2;
		immutable b = 8;
		auto dg = rcClosure(() => a + b);
		assert(dg() == 10);
		assert(dg() == 10);
	}

	{
		static auto rc_closure_nested()
		{
			int x = 42;
			return rcClosure(() => ++x);
		}

		auto dg = rc_closure_nested();
		assert(dg() == 43);
		auto dgCopy = dg;
		assert(dg() == 44);
		assert(dgCopy() == 45);
	}

	immutable i = 42;
	RefCountedClosure!(int delegate() pure nothrow @safe @nogc) dg =
		rcClosure(() => i);
	assert(dg() == 42);

	RefCountedClosure!(int delegate()) dg2 = dg;
	assert(dg2() == 42);

	dg2 = dg;
	assert(dg2() == 42);
}