From baf0cfce8e13e91aedfebe8e0ebfd3432718097d Mon Sep 17 00:00:00 2001
From: Wojtek Figat <wojtek@figat.pl>
Date: Tue, 5 Aug 2025 22:13:21 +0200
Subject: [PATCH] Add support for using custom memory allocator in lambda bind
 to `Function`

---
 Source/Engine/Core/Delegate.h | 47 ++++++++++++++++++++++++++++++++---
 1 file changed, 43 insertions(+), 4 deletions(-)

diff --git a/Source/Engine/Core/Delegate.h b/Source/Engine/Core/Delegate.h
index 00614a9ff..f71fcb94b 100644
--- a/Source/Engine/Core/Delegate.h
+++ b/Source/Engine/Core/Delegate.h
@@ -62,7 +62,7 @@ private:
     struct Lambda
     {
         int64 Refs;
-        void (*Dtor)(void*);
+        void (*Dtor)(void* callee, Lambda* lambda);
     };
 
     void* _callee;
@@ -78,8 +78,7 @@ private:
     {
         if (Platform::InterlockedDecrement(&_lambda->Refs) == 0)
         {
-            ((Lambda*)_lambda)->Dtor(_callee);
-            Allocator::Free(_lambda);
+            _lambda->Dtor(_callee, _lambda);
         }
     }
 
@@ -189,9 +188,10 @@ public:
             LambdaDtor();
         _lambda = (Lambda*)Allocator::Allocate(sizeof(Lambda) + sizeof(T));
         _lambda->Refs = 1;
-        _lambda->Dtor = [](void* callee) -> void
+        _lambda->Dtor = [](void* callee, Lambda* lambda) -> void
         {
             static_cast<T*>(callee)->~T();
+            Allocator::Free(lambda);
         };
         _function = [](void* callee, Params... params) -> ReturnType
         {
@@ -201,6 +201,45 @@ public:
         new(_callee) T(lambda);
     }
 
+    /// <summary>
+    /// Binds a lambda with a custom memory allocator.
+    /// </summary>
+    /// <param name="tag">The custom allocation tag.</param>
+    /// <param name="lambda">The lambda.</param>
+    template<typename AllocationType, typename T>
+    void Bind(typename AllocationType::Tag tag, const T& lambda)
+    {
+        if (_lambda)
+            LambdaDtor();
+        using AllocationData = typename AllocationType::template Data<byte>;
+        static_assert(AllocationType::HasSwap, "Function lambda binding supports only custom allocators with swap operation.");
+
+        // Allocate lambda (using temp data)
+        AllocationData tempAlloc(tag);
+        tempAlloc.Allocate(sizeof(Lambda) + sizeof(AllocationData) + sizeof(T));
+
+        // Move temp data into the one allocated
+        AllocationData* dataAlloc = (AllocationData*)(tempAlloc.Get() + sizeof(Lambda));
+        new(dataAlloc) AllocationData();
+        dataAlloc->Swap(tempAlloc);
+
+        // Initialize lambda
+        _lambda = (Lambda*)dataAlloc->Get();
+        _lambda->Refs = 1;
+        _lambda->Dtor = [](void* callee, Lambda* lambda) -> void
+        {
+            static_cast<T*>(callee)->~T();
+            AllocationData* dataAlloc = (AllocationData*)((byte*)lambda + sizeof(Lambda));
+            dataAlloc->Free();
+        };
+        _function = [](void* callee, Params... params) -> ReturnType
+        {
+            return (*static_cast<T*>(callee))(Forward<Params>(params)...);
+        };
+        _callee = (byte*)_lambda + sizeof(AllocationData) + sizeof(Lambda);
+        new(_callee) T(lambda);
+    }
+
     /// <summary>
     /// Unbinds the function.
     /// </summary>