std::function replacement using inlined storage

Function.hpp

#pragma once

#include <concepts>
#include <core/Utility.hpp>
#include <new>
#include <string.h>

// A movable function adapted from LLVM's unique function. Assumes
// relocatability.
template <typename F>
class Function;

namespace internal
{
template <bool Const, typename R, typename... T>
class FunctionStorage
{
public:
    operator bool() const
    {
        return meta_.caller_functions_ != nullptr;
    }

    R operator()(T... args)
    {
        void* addr = meta_.sso_enabled_ ? data_.sso : data_.heap;
        if (meta_.is_trivial_)
        {
            return ((TrivialStorage*)meta_.caller_functions_)
                ->invoke_(addr, fwd<T>(args)...);
        }
        else
        {
            return ((StandardStorage*)meta_.caller_functions_)
                ->invoke_(addr, fwd<T>(args)...);
        }
    }

    R operator()(T... args) const
    {
        void* addr = (void*)(meta_.sso_enabled_ ? data_.sso : data_.heap);
        if (meta_.is_trivial_)
        {
            return ((TrivialStorage*)meta_.caller_functions_)
                ->invoke_(addr, fwd<T>(args)...);
        }
        else
        {
            return ((StandardStorage*)meta_.caller_functions_)
                ->invoke_(addr, fwd<T>(args)...);
        }
    }

protected:
    FunctionStorage() = default;

    template <typename Caller>
    FunctionStorage(Caller&& caller)
    {
        using caller_t = remove_reference_t<Caller>;

        void* addr = data_.sso;

        // Check if we need to allocate
        if constexpr (sizeof(caller_t) > sso_size)
        {
            static_assert(
                sizeof(caller_t) <= 0xffff,
                "Captured caller size exceeds maximum allowable size");
            data_.heap         = malloc(sizeof(caller_t));
            addr               = data_.heap;
            meta_.sso_enabled_ = false;
            meta_.size_        = (u16)(sizeof(caller_t));
        }
        else
        {
            meta_.sso_enabled_ = true;
        }

        meta_.is_trivial_ = std::is_trivially_destructible_v<caller_t>
                            && std::is_trivially_move_constructible_v<caller_t>
                            && std::is_trivially_copy_constructible_v<caller_t>;
        meta_.caller_functions_ = (void*)&CallbackStorage<caller_t>::storage;

        new (addr) caller_t(fwd<Caller>(caller));
    }

    ~FunctionStorage()
    {
        if (!meta_.caller_functions_)
        {
            return;
        }

        void* addr = data_.sso;
        if (!meta_.sso_enabled_)
        {
            addr = data_.heap;
        }

        if (!meta_.is_trivial_)
        {
            ((StandardStorage*)meta_.caller_functions_)->destroy_(addr);
        }

        if (!meta_.sso_enabled_)
        {
            free(addr);
        }
    }

    FunctionStorage(FunctionStorage const& other)
    {
        if (meta_.caller_functions_ && !meta_.sso_enabled_)
        {
            free(data_.heap);
        }

        meta_ = other.meta_;
        if (!meta_.sso_enabled_)
        {
            data_.heap = malloc(meta_.size_);
            if (meta_.is_trivial_)
            {
                memcpy(data_.heap, other.data_.heap, meta_.size_);
            }
            else
            {
                ((StandardStorage*)meta_.caller_functions_)
                    ->copy_(data_.heap, other.data_.heap);
            }
        }
        else if (meta_.is_trivial_)
        {
            memcpy(&data_.sso, &other.data_.sso, sizeof(data_.sso));
        }
        else
        {
            ((StandardStorage*)meta_.caller_functions_)
                ->copy_(data_.heap, other.data_.heap);
        }
    }

    FunctionStorage(FunctionStorage&& other)
    {
        if (meta_.caller_functions_ && !meta_.sso_enabled_)
        {
            free(data_.heap);
        }

        meta_                         = other.meta_;
        other.meta_.caller_functions_ = nullptr;

        if (!meta_.sso_enabled_)
        {
            data_.heap = other.data_.heap;
        }
        else if (meta_.is_trivial_)
        {
            memcpy(&data_.sso, &other.data_.sso, sizeof(data_.sso));
        }
        else
        {
            ((StandardStorage*)meta_.caller_functions_)
                ->move_(data_.sso, other.data_.sso);
        }
    }

    FunctionStorage& operator=(FunctionStorage const& other)
    {
        if (this == &other)
        {
            return *this;
        }

        new (this) FunctionStorage(other);
        return *this;
    }

    FunctionStorage& operator=(FunctionStorage&& other)
    {
        if (this == &other)
        {
            return *this;
        }

        new (this) FunctionStorage(mv(other));
        return *this;
    }

private:
    struct Meta
    {
        bool is_trivial_        = false;
        bool sso_enabled_       = false;
        u17 size_               = 0u;
        void* caller_functions_ = nullptr;
    };

    static constexpr size_t sso_size = STYX_CACHE_LINE_SIZE - sizeof(Meta);

    using invoke_t  = R (*)(void*, param_t<T>...);
    using move_t    = void (*)(void*, void*);
    using copy_t    = void (*)(void*, void*);
    using destroy_t = void (*)(void*);

    template <typename Caller>
    static R invoke(void* addr, param_t<T>... args)
    {
        return (*(Caller*)(addr))(fwd<T>(args)...);
    }

    template <typename Caller>
    static void move(void* a, void* b)
    {
        if constexpr (std::is_move_constructible_v<Caller>)
        {
            if constexpr (std::is_trivially_move_constructible_v<Caller>)
            {
                memcpy(a, b, sizeof(Caller));
            }
            else
            {
                new (a) Caller(mv(*(Caller*)b));
            }
        }
    }

    template <typename Caller>
    static void copy(void* a, void* b)
    {
        if constexpr (std::is_copy_constructible_v<Caller>)
        {
            if constexpr (std::is_trivially_copy_constructible_v<Caller>)
            {
                memcpy(a, b, sizeof(Caller));
            }
            else
            {
                new (a) Caller(*(Caller*)b);
            }
        }
    }

    template <typename Caller>
    static void destroy(void* addr)
    {
        if constexpr (!std::is_trivially_destructible_v<Caller>)
        {
            ((Caller*)addr)->~Caller();
        }
    }

    struct TrivialStorage
    {
        invoke_t invoke_;
    };

    struct StandardStorage
    {
        invoke_t invoke_;
        move_t move_;
        copy_t copy_;
        destroy_t destroy_;
    };

    template <typename Caller>
    struct CallbackStorage
    {
        constexpr static StandardStorage storage{
            &invoke<Caller>,
            &move<Caller>,
            std::is_copy_constructible_v<Caller> ? &copy<Caller> : nullptr,
            &destroy<Caller>};
    };

    template <typename Caller>
        requires(std::is_trivially_destructible_v<Caller>
                 && std::is_trivially_move_constructible_v<Caller>
                 && std::is_trivially_copy_constructible_v<Caller>)
    struct CallbackStorage<Caller>
    {
        constexpr static TrivialStorage storage{&invoke<Caller>};
    };

    alignas(16) union
    {
        void* heap = nullptr;
        byte sso[sso_size];
    } data_;
    Meta meta_;
};

} // namespace internal

template <typename Caller, typename R, typename... T>
concept Invokable = requires(Caller&& c, T&&... args) {
    {
        c(fwd<T>(args)...)
    } -> std::convertible_to<R>;
};

template <typename R, typename... T>
class Function<R(T...)> final : public internal::FunctionStorage<false, R, T...>
{
    using parent_t = internal::FunctionStorage<false, R, T...>;

public:
    Function()                           = default;
    Function(Function const&)            = default;
    Function(Function&&)                 = default;
    Function& operator=(Function const&) = default;
    Function& operator=(Function&&)      = default;

    template <typename Caller>
        requires Invokable<Caller, R, T...>
    Function(Caller&& caller)
        : parent_t{fwd<Caller>(caller)}
    {
    }
};

template <typename R, typename... T>
class Function<R(T...) const> final
    : public internal::FunctionStorage<true, R, T...>
{
    using parent_t = internal::FunctionStorage<true, R, T...>;

public:
    Function()                           = default;
    Function(Function const&)            = default;
    Function(Function&&)                 = default;
    Function& operator=(Function const&) = default;
    Function& operator=(Function&&)      = default;

    template <typename Caller>
        requires Invokable<Caller, R, T...>
    Function(Caller&& caller)
        : parent_t{fwd<Caller>(caller)}
    {
    }
};