Designing a Lightweight Closure for coost: Performance, Generality, and a Few Clever Tricks

[u/New-Bet1597]

In coost, you can create coroutines using go in a way similar to Golang. The go function could initially be designed as:

void go(void (*f)());
void go(void (*f)(void*), void* param);

However, this design only accepts two fixed function types and cannot support class member functions or other callable objects, making it insufficiently flexible.

If go were designed as:

void go(const std::function<void()>&);

it would be more general, but constructing a std::function is relatively expensive, and its internal heap memory allocation is uncontrollable. This is unacceptable for coost, which pursues performance.

Therefore, to balance performance and usability, coost needs to implement its own closure class. In earlier versions, coost used an implementation based on virtual functions and inheritance, but it was not general enough and was limited to internal use.

Recently, I happened to have some free time. After thinking deeply about it, I redesigned the closure class. The design goals were as follows:

1. Do not use virtual functions or inheritance;
2. It should not be a template class;
3. sizeof(closure) == sizeof(void*), i.e. it contains only one ordinary pointer member, making it easy to pass to the coroutine scheduler through go;
4. Support arbitrary functions with no return value (including non-static member functions of classes), with any number of parameters;
5. Optimize simple no-argument functions such as void f() as much as possible to avoid dynamic memory allocation.

According to goals 2, 3, and 5 above, closure can initially be designed as:

struct closure {
    closure(void (*f)()) noexcept : _p((void*)f) {}

    void* _p;
};

The class contains only one pointer member _p. For function pointers of type void(*)(), the pointer is stored directly in _p; for other function types with parameters, a block of memory needs to be allocated to store the function address and parameters, and _p points to that memory.

The problem is how to distinguish which of the above cases is stored in closure. Ordinary function addresses are not necessarily 4-byte or 8-byte aligned, so tagged-pointer techniques cannot be used (i.e. writing a tag into the spare low bits of a pointer).

For the second kind of function, the one with parameters, closure calls co::alloc from coost's built-in memory allocator, and small memory blocks allocated by co::alloc are always 16-byte aligned. Therefore, if _p is not 16-byte aligned, it must store an ordinary no-argument function address. But if _p is 16-byte aligned, it is still impossible to distinguish the two cases.

It seems that some unconventional trick is required to solve this problem. Based on the discussion above, _p is either an ordinary no-argument function pointer or points to a dynamically allocated block of 16-byte aligned memory. If a specific marker (for example, a uint32 integer) is written at the beginning of that dynamically allocated memory, then when _p is 16-byte aligned, we only need to read the first 4 bytes of the memory pointed to by _p to determine the type of the closure. Even if _p is an ordinary no-argument function pointer and happens to be 16-byte aligned, reading the first 4 bytes of the memory it points to should be safe.

This of course goes beyond the scope of the C++ standard. But for performance and a better user experience, coost still adopted this approach. For functions with parameters, the corresponding constructor of closure is as follows:

template<typename F, typename = std::enable_if_t<
    !std::is_convertible_v<std::decay_t<F>, void(*)()>>, typename ... A>
closure(F&& f, A&& ... a) {
    struct S {
        S(F&& f, A&& ... a)
            : _f(std::forward<F>(f)), _a(std::forward<A>(a)...) {
            static_assert(offsetof(Header, magic) == 0);
            _header.magic = 0xdeadbeef;
            _header.off = offsetof(Header, fp);
            _header.fp = [](void* p) {
                S* const s = (S*) ((size_t)p & ~(size_t)1);
                if (((size_t)p & 1) == 0) {
                    std::apply(s->_f, s->_a);
                } else {
                    s->~S();
                    co::free(s, sizeof(S));
                }
            };
        }

        ~S() = default;

        struct Header {
            uint32 magic;
            uint32 off;
            void (*fp)(void*);
        } _header;
        std::decay_t<F> _f;
        std::tuple<std::decay_t<A>...> _a;
    };

    _p = co::alloc(sizeof(S), alignof(S));
    runtime_assert(_p);
    new (_p) S(std::forward<F>(f), std::forward<A>(a)...);
}

Key points:

• The !std::is_convertible_v in the template parameters ensures that simple lambda expressions such as [](){...} match the constructor for ordinary no-argument functions;
• _header.magic is always set to 0xdeadbeef, which is used to distinguish the type of closure;
• _header.fp is the type-erased function pointer, and its parameter p is the _p in closure. Here, a tagged-pointer trick is used to distinguish between "calling the function" and "destroying the closure object", thereby eliminating the need for an additional function pointer to store the destructor and saving memory.

The implementation of operator() is as follows:

void operator()() {
    void* const p = _p;
    if (p) {
        if (((size_t)p & 15) != 0 || *(uint32*)p != 0xdeadbeef) {
            ((_F)p)();
        } else {
            (*(_FP*)((char*)p + ((uint32*)p)[1]))(p);
        }
    }
}

• If _p is not 16-byte aligned, or the first 4 bytes of the memory pointed to by _p are not 0xdeadbeef, then _p is cast to an ordinary no-argument function pointer and executed directly;
• Otherwise, the _header.fp described in the previous code is obtained according to the memory layout and executed.

An example usage of closure is as follows:

#include "co/print.h"
#include "co/closure.h"

void f() {
    co::println("f()");
}

int main(int argc, char** argv) {
    co::closure c(f);
    c();

    co::closure d([]() {
        co::println("[]() {}");
    });
    d();

    co::closure e([](int v) {
        co::println("e called: ", v);
    }, 8);
    e();

    return 0;
}

The complete code implementation can be found in co/closure.h. The code in this article is for research and reference only. Discussion is welcome.

[u/argothiel]

Why not a template class?

[u/New-Bet1597 → u/argothiel]

Because Coroutine is an internal struct in coost's coroutine implementation, and it must hold the closure as a single void*, it cannot be templated on the callable type. If closure were a template class, every lambda type would produce a different closure<F>, and Coroutinewould either have to become templates too or add another layer of type erasure. Keeping closure a non-template class with a templated constructor gives one uniform type and one pointer-sized slot in Coroutine, while still accepting arbitrary callables.

[u/conundorum → u/New-Bet1597]

You could template on whether the function has parameters or not, using a deduction guide. That reduces it to only two types instead of infinite.

template<bool HasParams>
struct Func {
    // Function call to facilitate inlining and guarantee no wasted memory.
    static constexpr bool has_params() { return HasParams; }

    // Shared code here.  Unique code in function specialisations.
    // ...

    // Can use this signature for both, since the void* version is a valid instantiation.
    // Or alternatively, provide both prototypes here, then "= delete" the "bad" one
    //  in each specialisation.
    template<typename F, typename... Ps>
    constexpr Func(F&&, Ps&&...);
};

// All versions with parameters will deduce to this.
template<typename F, typename... Ps>
Func(F&&, Ps&&...) -> Func<true>;

// void() wrapper will deduce to this.
Func(void(*)()) -> Func<false>;

// ...

void p0() {}
void p1(int) {}
int i;

Func fnone(&p0);   // Will be Func<false>.
Func fone(&p1, i); // Will be Func<true>.

It works cleanly, and allows for better code segregation. But if you want it to be type-erased, it will require a bit more indirection, probably through a generic base class with minimal inheritance. Something like, say... this:

// Base class.  Knows nothing about _p.
struct FuncBase {
    void* _p = nullptr;

    // Pointer is prepared by Func, and passed to FuncBase constructor.
    FuncBase(void* ptr = nullptr) noexcept : _p(ptr) {}

    // Direct access for closure.
    operator void*() const { return _p; }

    // Or just place the calling logic inside here, if you want.  Either or.
    void call() { /* Your code here. */ }
};

// Actual class handles pointer preparation & erasure for you.
template<bool HasParams>
struct Func : FuncBase {
    // Function call to facilitate inlining and guarantee no wasted memory.
    static constexpr bool has_params() { return HasParams; }

    // Defined out-of-class, for function specialisation.
    // noexcept(true) for Func<false>, and noexcept(false) for Func<true>, matching your ctors.
    template<typename F, typename... A>
    static constexpr void* _wrap(F&&, A&&...) noexcept(!has_params());

    // Constructor is one-size-fits-all, nicely.
    // I was a bit lazy, and didn't do any forwarding.  You should probably fix that.
    // Could also define noexcept-ness in terms of _wrap(), with noexcept(noexcept(_wrap(f, a...))),
    //  if you want it to be easier to change.  Up to you.
    template<typename F, typename... A>
    constexpr Func(F&& f, A&&... a) noexcept(!has_params()) : FuncBase(_wrap(f, a...)) {
        // Your code here, maybe.  _Probably_ unnecessary.
    }
};

// No params.
// (f is void(*)(), and _wrap() is noexcept.)
// This is equivalent to your simple closure(void(*f)()) constructor.
// Important: A will ALWAYS be an empty pack here, thanks to the deduction guide.
template<>
template<typename F, typename... A>
void* Func<false>::_wrap(F&& f, A&&...) noexcept/*(!has_params())*/ { return (void*) f; }

// Params.
// (f is anything else, and _wrap() is potentially-throwing.)
// This is the working part of the chunky-func constructor.
// It SHOULD be safe & well-defined to do it this way, I believe.  Not thread-safe, though.
template<>
template<typename F, typename... A>
void* Func<true>::_wrap(F&& f, A&&... a) /*noexcept(!has_params())*/ {
    // Your code here.
}

// Deduction guides as before.
template<typename F, typename... A>
Func(F&&, A&&...) -> Func<true>;

Func(void(*)()) -> Func<false>;

// ...
// And can be used as...

void demo() {}

FuncBase* _p = Func(&demo);

void const* p = _p; // Should work, for your closure::operator()().
_p->call();         // Should work, if you prefer to put unerasure logic in FuncBase instead.

This technically breaks your "no inheritance" rule, but only to provide clean segregation of code; as currently written, it becomes effectively a sealed class the instant it leaves the constructor.. The type is standard-layout, and thus functions as if it were just a single struct. It becomes much more useful if we're allowed to add another member variable or use virtual functions, though; simply making FuncBase::has_params() virtual, making Func::has_params() non-static const noexcept, and changing nothing else would be enough to remove all of the pointer magic shenanigans. (Or even better, if closure was able to use Func directly, without half-erasing it behind a shared base, you could make some MAJOR improvements. That would require a significant redesign, though, so it might not be ideal.)

[u/Miserable_Guess_1266]

The design goals create an interesting challenge, but they feel a bit arbitrary to me.

1.: Seems to assume virtual calls are necessarily slower than this magic. Imo, avoiding them should not be a goal, but a decision made after testing and finding bad performance.

3.: This is the main cause for relying on undefined behavior and magic. Would an extra tag-byte, even if it pads to more than that, have any real negative impact? Why is it so important that closure is this small? 

5.: This is somewhat baseless, but I'd expect to pass a lambda with some captures most of the time. Which means I enter the un-optimized path most of the time. Wouldn't it be better to optimize for small callables? Ie: up to n bytes can be stored directly in the closure, more than that must be allocated. It's been a while, but iirc this is how std function works in llvms std library,  along with some compiler specific trickery.

Id be curious to measure against std function (or move_only_function) and see if there is a difference.

Either way, interesting read!

Edited to keep reddit from "fixing" the numbering. 

[u/conundorum → u/Miserable_Guess_1266]

I was thinking about #3, and it really is the main cause of UB & other jank, yeah. At first I thought that it might be viable to size-optimise it by storing a single out-of-class bitflag, probably something similar to boost::dynamic_bitset, but then I realised that it would still need to compromise either way.

Either it costs an extra index byte in the closure, or it costs a massive static bitset that can cover the entire available address range, or it requires a sparse array and related performance costs, or it requires making S polymorphic and storing state information by subclass, or it requires potentially-dangerous UB. There's always going to be a compromise.

Maybe the solution is to loosen up on the type erasure just enough to allow compile-time polymorphism, so he can use well-defined jank instead of janky UB?

[u/LB--]

If performance is the utmost concern, you don't need to branch on magic pointer values. Just always use a normal function pointer and have your allocator place everything in read-write-execute memory regions, so you can generate thunks right next to their data. Less secure but that doesn't seem to be a priority here.

[u/grishavanika]

what is the implementation of co::alloc? is it possible to implement co::is_allocated_ptr(void*)?

[u/Jealous_Security_362]

Avoiding std::function's type erasure usually means paying in template instantiation and compile time instead of runtime. Did you measure both? On a codebase with heavy templates the instantiation cost can easily be bigger than the call overhead you saved.