Lambda.h

#pragma once

#include "Allocator/Allocator.h"
#include "Closure/Function.h"
#include "LambdaView.h"
#include "Meta/MetaFunction.h"
#include "Meta/MetaFunctor.h"
#include "Meta/MetaParameterPack.h"

namespace SE {

namespace Impl {

struct FLambdaStorage {
    void *FuncThis;
    const struct FFuncManager *FuncManager;
};

template<typename T>
inline const FFuncManager &GetFuncManager();

struct FFuncManager {
    void (*const CallCopy)(FLambdaStorage &Storage, const FLambdaStorage &OtherStorage);
    void (*const CallMoveAndDestroy)(FLambdaStorage &Storage, FLambdaStorage &&OtherStorage);
    void (*const CallDestroy)(FLambdaStorage &Storage);

    template<typename T>
    inline static constexpr FFuncManager Create() {
        return FFuncManager{&TemplatedCallCopy<T>, &TemplatedCallMoveAndDestroy<T>, &TemplatedCallDestroy<T>};
    }

private:
    template<typename T>
    static void TemplatedCallCopy(FLambdaStorage &Storage, const FLambdaStorage &OtherStorage) {
        T *Other = (T *)OtherStorage.FuncThis;
        T *This = (T *)FHeapAllocator::Alloc(sizeof(T), alignof(T));
        if constexpr (Meta::IsTrivialCopy<T>) {
            *This = *Other;
        } else {
            Meta::Construct(*This, *Other);
        }
        Storage.FuncThis = This;
        Storage.FuncManager = OtherStorage.FuncManager;
    }
    template<typename T>
    static void TemplatedCallMoveAndDestroy(FLambdaStorage &Storage, FLambdaStorage &&OtherStorage) {
        T *Other = (T *)OtherStorage.FuncThis;
        T *This = (T *)FHeapAllocator::Alloc(sizeof(T), alignof(T));
        if constexpr (Meta::IsTrivialMove<T>) {
            *This = Move(*Other);
        } else {
            Meta::Construct(*This, Move(*Other));
        }
        if constexpr (!Meta::IsTrivialDestruct<T>) {
            Meta::Destruct(*Other);
        }
        FHeapAllocator::Free(Other, sizeof(T), alignof(T));
        OtherStorage.FuncThis = nullptr;
        Storage.FuncThis = This;
        Storage.FuncManager = OtherStorage.FuncManager;
    }
    template<typename T>
    static void TemplatedCallDestroy(FLambdaStorage &Storage) {
        T *This = (T *)Storage.FuncThis;
        if constexpr (!Meta::IsTrivialDestruct<T>) {
            Meta::Destruct(*This);
        }
        FHeapAllocator::Free(This, sizeof(T), alignof(T));
        Storage.FuncThis = nullptr;
    }
};

template<typename T>
Reflection(NoReflection)
inline const FFuncManager &GetFuncManager() {
    static constexpr FFuncManager Manager = FFuncManager::Create<T>();
    return Manager;
}

}  // namespace Impl

template<typename TReturnType, typename... TParams>
class TLambda<TReturnType(TParams...)> {
public:
    TLambda() {
        Func = (void *)&Impl::EmptyFunc<TReturnType, TParams...>;
        Storage.FuncThis = nullptr;
    }

    TLambda(const TLambda &Other) {
        if (Other.Storage.FuncThis) {
            Other.Storage.FuncManager->CallCopy(Storage, Other.Storage);
        } else {
            Storage.FuncThis = nullptr;
        }
        Func = Other.Func;
    }

    TLambda(TLambda &&Other) {
        if (Other.Storage.FuncThis) {
            Other.Storage.FuncManager->CallMoveAndDestroy(Storage, Move(Other.Storage));
        } else {
            Storage.FuncThis = nullptr;
        }
        Func = (void *)&Impl::EmptyFunc<TReturnType, TParams...>;
        Swap(Func, Other.Func);
    }

    TLambda(Meta::FNullptr) {
        Storage.FuncThis = nullptr;
        Func = (void *)&Impl::EmptyFunc<TReturnType, TParams...>;
    }

    template<typename T>
    requires(!Meta::IsSame<Meta::TRemoveRef<T>, TLambda> && Meta::IsCallableAsFunctor<T, TReturnType(TParams...)>) TLambda(T &&Functor) {
        if (Meta::IsFunctionNull(Functor)) {
            Storage.FuncThis = nullptr;
            Func = (void *)&Impl::EmptyFunc<TReturnType, TParams...>;
        } else {
            Initialize(Meta::Forward<T>(Functor));
        }
    }

    TLambda &operator=(const TLambda &Other) {
        if (Storage.FuncThis) {
            Storage.FuncManager->CallDestroy(Storage);
        }
        if (Other.Storage.FuncThis) {
            Other.Storage.FuncManager->CallCopy(Storage, Other.Storage);
        }
        Func = Other.Func;
        return *this;
    }

    TLambda &operator=(TLambda &&Other) {
        if (Storage.FuncThis) {
            Storage.FuncManager->CallDestroy(Storage);
        }
        if (Other.Storage.FuncThis) {
            Other.Storage.FuncManager->CallMoveAndDestroy(Storage, Move(Other.Storage));
        }
        Func = (void *)&Impl::EmptyFunc<TReturnType, TParams...>;
        Swap(Func, Other.Func);
        return *this;
    }

    TLambda &operator=(Meta::FNullptr) {
        if (Storage.FuncThis) {
            Storage.FuncManager->CallDestroy(Storage);
        }
        Func = (void *)&Impl::EmptyFunc<TReturnType, TParams...>;
        return *this;
    }

    template<typename T>
    requires(Meta::IsCallableAsFunctor<T, TReturnType(TParams...)>) TLambda &operator=(T &&Functor) {
        if (Storage.FuncThis) {
            Storage.FuncManager->CallDestroy(Storage);
        }
        if (Meta::IsFunctionNull(Functor)) {
            Func = (void *)&Impl::EmptyFunc<TReturnType, TParams...>;
        } else {
            Initialize(Meta::Forward<T>(Functor));
        }
        return *this;
    }

    ~TLambda() {
        if (Storage.FuncThis) {
            Storage.FuncManager->CallDestroy(Storage);
        }
    }

    SE_INLINE TReturnType operator()(TParams... Params) const {
        // Branch misprediction is cheaper than function call.
        if (Storage.FuncThis) {
            var ProxyFuncPtr = (TReturnType(*)(void *, TParams...))Func;
            return ProxyFuncPtr(Storage.FuncThis, Meta::Forward<TParams>(Params)...);
        } else {
            var RawFuncPtr = (TReturnType(*)(TParams...))Func;
            return RawFuncPtr(Meta::Forward<TParams>(Params)...);
        }
    }
    
    void *Func;
    Impl::FLambdaStorage Storage;

private:

    // Serves two purposes:
    // 1. For lambdas reinterprets the captures class.
    // 2. For raw functions with arguments that are different, performs the conversion through forwarding.
    template<typename T>
    static TReturnType ProxyFunc(void* ThisPtr, TParams... Params) {
        // T functor could get inlined in here.
        return (*(T*)ThisPtr)(Meta::Forward<TParams>(Params)...);
    }

    template<typename T>
    void Initialize(T &&Functor) {
        using FunctionParams = Meta::TFunctionParams<T>;
        using FunctionReturnType = Meta::TFunctionReturnType<T>;
        if constexpr ((Meta::HasCallOperator<T> && !Meta::IsCapturelessLambda<T>) || (!Meta::IsSame<FunctionReturnType, TReturnType> || !Meta::SameParameterPacks<FunctionParams, Meta::TParameterPack<TParams...>>)) {
            Func = (void *)&ProxyFunc<T>;
            Storage.FuncManager = &Impl::GetFuncManager<T>();
            var This = (T *)FHeapAllocator::Alloc(sizeof(T), alignof(T));
            ::new ((void *)This) T(Meta::Forward<T>(Functor));
            Storage.FuncThis = This;
        } else {
            TFunction<TReturnType(TParams...)> *Decayed = Functor;
            Func = (void *)Decayed;
            Storage.FuncThis = nullptr;
        }
    }
};

}  // namespace SE