Circularly referenced lambdas in C++

A handy C++ header that lets you create lambdas that reference each other.
June Rhodes posted on Oct 22, 2020

I decided to start up a new Redpoint Games blog as I wanted somewhere to put all the esoteric and weird Unreal Engine knowledge that I’ve accumulated over the past year.

Without further ado, here is how to make circularly referenced lambdas in C++; that is, lambdas that reference each other or themselves.

The problem

When writing test code for asynchronous behaviour, I’d frequently run into a problem where lambdas are dependent on each other, like so:

void SomeTest(std::function<void()> OnDone)
{
    auto SomeInterface = ...;

    bool* bDidCallEventListener = new bool(false);

    auto Cleanup = [bDidCallEventListenerCallbackHandleTimerHandleOnDone](){
        SomeInterface->RemoveEventListener(CallbackHandle);
        FTicker::GetCoreTicker().RemoveTicker(TimerHandle);
        this->TestTrue("Event listener gets called", *bDidCallEventListener);
        delete bDidCallEventListener;
        OnDone();
    };

    FDelegateHandle CallbackHandle = SomeInterface->AddEventListener(
        FDelegate::CreateLambda([bDidCallEventListenerCleanup](){
            bDidCallEventListener = true;
            Cleanup();
        }));

    FDelegateHandle TimerHandle = FTicker::GetCoreTicker().AddTicker(
        FTimerDelegate::CreateLambda([Cleanup](){
            Cleanup();
            return false;
        }), 5.0f /* 5 second time out */);

    SomeInterface->DoAsyncCallThatShouldFireEventListener(...);
}

This obviously doesn’t work. We can’t shift the FDelegateHandle declarations above Cleanup, because the delegate handle is on the stack, and when the clean up lambda gets created, it’ll capture the current value of the delegate handle. Even if we later assign to CallbackHandle, that won’t get propagated into Cleanup.

At the same time, the lambda we’re passing into AddEventListener needs to be able to reference Cleanup so that the test ends and passes immediately when the event listener fires. We could workaround the problem in this particular scenario by only cleaning up in the timer, but not only does this make the test slower, it’s not applicable to all cases with circularly dependent lambdas.

Move it to the heap

For trivial cases where we need a lambda to reference itself, we can create a double pointer to the lambda. We pass that pointer into the lambda, and after we’ve declared the lambda, we then move it onto the heap with std::move.

#include <iostream>
#include <string>
#include <functional>

int main()
{
    // Create a pointer on the heap.
    std::function<void(int)>** LambdaRef = new std::function<void(int)>*();

    std::function<void(int)> LambdaImpl = [LambdaRef](int InvocationCount){
        // Typically this callback would be making an asynchronous request
        // that itself would return at a later point, and then you would either
        // re-invoke the lambda or clean up based on the result of the
        // asynchronous request: e.g. if you were retrying requests on failure.
        //
        // To keep the example of moving lambdas to the heap though, this
        // lambda just re-invokes itself 3 times.

        std::cout << "Invoked: " << InvocationCount << std::endl;

        if (InvocationCount < 3)
        {
            (**LambdaRef)(InvocationCount + 1);
        }
        else
        {
            delete *LambdaRef;
            delete LambdaRef;
        }
    };

    // Move the lambda to the heap.
    std::function<void(int)>* LambdaOnHeap = new std::function<void(int)>(
        std::move(LambdaImpl));

    // Update our double pointer's value to point at the lambda now on the heap.
    *LambdaRef = LambdaOnHeap;

    // Now call the lambda.
    (**LambdaRef)(0);

    return 0;
}

With the lambda referencing itself, this enables the lambda to re-invoke itself at some later point. This is useful if you need to perform retries or exponential back-off with some service, and you want to allow the developer to just pass in an std::function or lambda and call it a day.

On a sidenote...
The Discord integration for Unreal Engine 4 plugin uses this technique to call Discord Game SDK functions, since they can request retries with exponential back-off.

This solution only works if you’re cleaning up the lambda in exactly one place, and you can guarantee the code performing the clean up will only ever be called once.

For the problem we started with, we’ll need a solution that’s a little more involved. We’ll need to put our lambdas on the heap and reference count them.

Referencing counting our lambdas

Introducing FHeapLambda. This class allows you to create a “to be filled in” reference to a lambda; you can let this value be captured by other lambdas, and when it’s finally assigned, all of the copies of the FHeapLambda will point at the same place.

This is done by allocating a state object on the heap when FHeapLambda is first constructed (with the default constructor). Every copy and move increments the reference counter; while the destructor decrements the counter. When the reference counter reaches 0 in any of the copies, it knows it’s the last possible reference to the state and thus the heap-allocated state can be freed.

When you assign a lambda to an FHeapLambda, the lambda gets copied into that heap allocated state.

Since we’ll often be using FHeapLambda with clean up actions and delete statements, the implementation below optionally supports both “always fire at least once” and “fire at most once” semantics. Combined together you can guarantee that your clean up code runs regardless of the circumstances that causes the lambda to no longer be referenced.

Update:
The FHeapLambda implementation below has been modified to fix unsafe usage of this in the invocation handler and to support Clang (Android), since this blog post was originally published.
#include <functional>

enum class EHeapLambdaFlags
{
    // No additional behaviour.
    None = 0x0,

    // Always fire the lambda at least once (via the destructor) even if it's never explicitly invoked.
    AlwaysCleanup = 0x1,

    // Only fire the callback once, no matter how many time the () operator is used.
    OneShot = 0x2,

    // Always fire the callback only once.
    OneShotCleanup = OneShot | AlwaysCleanup,
};

template <EHeapLambdaFlags TFlagstypename... TArgsclass FHeapLambda
{
private:
    template <EHeapLambdaFlags Flagsstruct FHeapLambdaBehaviour;
    template <> struct FHeapLambdaBehaviour<EHeapLambdaFlags::None>
    {
    public:
        static void DoCleanup(FHeapLambda *Lambda)
        {
        }
        static bool ShouldOnlyFireOnce()
        {
            return false;
        }
    };
    template <> struct FHeapLambdaBehaviour<EHeapLambdaFlags::AlwaysCleanup>
    {
    public:
        static void DoCleanup(FHeapLambda *Lambda)
        {
            (*Lambda)();
        }
        static bool ShouldOnlyFireOnce()
        {
            return false;
        }
    };
    template <> struct FHeapLambdaBehaviour<EHeapLambdaFlags::OneShot>
    {
    public:
        static void DoCleanup(FHeapLambda *Lambda)
        {
        }
        static bool ShouldOnlyFireOnce()
        {
            return true;
        }
    };
    template <> struct FHeapLambdaBehaviour<EHeapLambdaFlags::OneShotCleanup>
    {
    public:
        static void DoCleanup(FHeapLambda *Lambda)
        {
            (*Lambda)();
        }
        static bool ShouldOnlyFireOnce()
        {
            return true;
        }
    };

    class HeapLambdaState
    {
    public:
        std::function<void(TArgs...)> *CallbackOnHeap;
        int RefCount;
        int FireCount;

        HeapLambdaState()
        {
            this->CallbackOnHeap = nullptr;
            this->RefCount = 1;
            this->FireCount = 0;
        }
        HeapLambdaState(const HeapLambdaState &) = delete;
        HeapLambdaState(HeapLambdaState &&) = delete;
        HeapLambdaState &operator=(const HeapLambdaState &Other) = delete;
        HeapLambdaState &operator=(HeapLambdaState &&Other) = delete;

        void ClearCallback()
        {
            if (this->CallbackOnHeap != nullptr)
            {
                delete this->CallbackOnHeap;
                this->CallbackOnHeap = nullptr;
            }
        }

        ~HeapLambdaState()
        {
            this->ClearCallback();
        }
    };
    HeapLambdaState *State;
    bool *DestructionTracker;

public:
    FHeapLambda()
    {
        this->State = new HeapLambdaState();
        this->DestructionTracker = nullptr;
    }
    FHeapLambda(const FHeapLambda &Other)
    {
        this->State = Other.State;
        this->State->RefCount++;
        this->DestructionTracker = nullptr;
    };
    FHeapLambda(FHeapLambda &&Other)
    {
        this->State = Other.State;
        this->State->RefCount++;
        this->DestructionTracker = nullptr;
    }
    template <class FFHeapLambda(F f) : FHeapLambda(std::function<void(TArgs...)>(f))
    {
    }
    FHeapLambda(const std::function<void(TArgs...)> &Callback)
    {
        this->State = new HeapLambdaState();
        this->Assign(Callback);
        this->DestructionTracker = nullptr;
    };
    ~FHeapLambda()
    {
        this->Reset();
        this->State = nullptr;

        // The destruction tracker lets the invocation code know if
        // the this pointer is no longer safe to use.
        if (this->DestructionTracker != nullptr)
        {
            *this->DestructionTracker = true;
        }
        this->DestructionTracker = nullptr;
    };
    // Do not implement this operator.
    FHeapLambda &operator=(const FHeapLambda &Other) = delete;
    FHeapLambda &operator=(FHeapLambda &&Other)
    {
        this->Reset();
        delete this->State;
        this->State = Other.State;
        this->State->RefCount++;
        return *this;
    }
    void Reset()
    {
        check(this->State->RefCount > 0);
        this->State->RefCount--;
        if (this->State->RefCount == 0)
        {
            if (this->State->FireCount == 0)
            {
                FHeapLambdaBehaviour<TFlags>::DoCleanup(this);
            }

            delete this->State;
            this->State = new HeapLambdaState();
        }
    }
    template <class Fvoid Assign(F Callback)
    {
        check(this->State->CallbackOnHeap == nullptr /* Lambda must not already be assigned! */);
        this->State->CallbackOnHeap = new std::function<void(TArgs...)>(Callback);
    }
    void Assign(std::function<void(TArgs...)> Callbackconst
    {
        check(this->State->CallbackOnHeap == nullptr /* Lambda must not already be assigned! */);
        this->State->CallbackOnHeap = new std::function<void(TArgs...)>(Callback);
    }
    template <class FFHeapLambda &operator=(F Callback)
    {
        this->Assign(std::function<void(TArgs...)>(Callback));
        return *this;
    }
    FHeapLambda &operator=(std::function<void(TArgs...)> Callback)
    {
        this->Assign(Callback);
        return *this;
    }
    void operator()(TArgs... argsconst
    {
        bool bSetupDestructionTracker = this->DestructionTracker == nullptr;

        if (this->State != nullptr && this->State->CallbackOnHeap != nullptr)
        {
            if (!FHeapLambdaBehaviour<TFlags>::ShouldOnlyFireOnce() || this->State->FireCount == 0)
            {
                // Keep a local for the state on the call stack, because once we invoke the
                // callback, our FHeapLambda this pointer might be freed, and we need to continue
                // to access the state to do cleanup.
                auto StateLocal = this->State;

                // Set up the destruction tracker so that we know after the callback
                // whether we should reset our state via the this pointer, or if we should
                // just clean up the state indirectly because we've already been freed.
                if (bSetupDestructionTracker)
                {
                    const_cast<FHeapLambda *>(this)->DestructionTracker = new bool(false);
                }
                auto DestructionTrackerLocal = const_cast<FHeapLambda *>(this)->DestructionTracker;

                // Obtain a reference while we run the callback, to prevent the state from being
                // freed while we still have a local reference to it.
                StateLocal->RefCount++;

                // Increment the fire count to track invocations.
                StateLocal->FireCount++;

                // Invoke the callback.
                (*StateLocal->CallbackOnHeap)(args...);

                if (FHeapLambdaBehaviour<TFlags>::ShouldOnlyFireOnce())
                {
                    // Ensures that any references held by the lambda are immediately freed after the first
                    // fire.
                    StateLocal->ClearCallback();
                }

                // If we're still alive, use the normal reset path. Otherwise, manually release the
                // reference off StateLocal.
                if (!bSetupDestructionTracker || *DestructionTrackerLocal == false)
                {
                    const_cast<FHeapLambda *>(this)->Reset();
                    if (bSetupDestructionTracker)
                    {
                        delete this->DestructionTracker;
                        const_cast<FHeapLambda *>(this)->DestructionTracker = nullptr;
                    }
                }
                else
                {
                    check(StateLocal->RefCount > 0);
                    StateLocal->RefCount--;
                    if (StateLocal->RefCount == 0)
                    {
                        delete StateLocal;
                    }
                    delete DestructionTrackerLocal;
                }
            }
        }
    }
};

Using FHeapLambda

Now we can use FHeapLambda to solve our original problem quite nicely.

void SomeTest(std::function<void()> OnDone)
{
    auto SomeInterface = ...;

    bool* bDidCallEventListener = new bool(false);
    FHeapLambda<EHeapLambdaFlags::OneShotCleanup> Cleanup;

    FDelegateHandle CallbackHandle = SomeInterface->AddEventListener(
        FDelegate::CreateLambda([bDidCallEventListenerCleanup](){
            bDidCallEventListener = true;
            Cleanup();
        }));

    FDelegateHandle TimerHandle = FTicker::GetCoreTicker().AddTicker(
        FTimerDelegate::CreateLambda([Cleanup](){
            Cleanup();
            return false;
        }), 5.0f /* 5 second timeout */);

    Cleanup = [bDidCallEventListenerCallbackHandleTimerHandleOnDone](){
        SomeInterface->RemoveEventListener(CallbackHandle);
        FTicker::GetCoreTicker().RemoveTicker(TimerHandle);
        this->TestTrue("Event listener gets called", *bDidCallEventListener);
        delete bDidCallEventListener;
        OnDone();
    };

    SomeInterface->DoAsyncCallThatShouldFireEventListener(...);
}
All code examples are MIT licensed unless otherwise specified.