NanoFunction.hpp

/**
 * @file NanoFunction.hpp
 * @brief A more efficient std::function replacement
 * @author Andrés Brugarolas
 * @version 1.0
 */

#pragma once

#include <functional>
#include <memory>
#include <type_traits>
#include <exception>
#include <cstddef>

namespace bruga::fn {

namespace internal {
    class BlockAllocator;
}

    template<class, size_t MaxSize = 4 * 1024 * 1024> // 4 Mb
    class Function;

    /**
     * @brief A more efficient std::function replacement
     * @tparam R The return type of the function
     * @tparam Args The argument types of the function
     * @tparam MaxSize The maximum size of the callable object that can be stored
     *
     * @tparam R
     * @tparam Args
     * @tparam MaxSize
     */
    template<class R, class... Args, size_t MaxSize>
    class Function<R(Args...), MaxSize> {
    private:
        enum class Operation {
            Clone, Destroy
        };

        using Invoker = R (*)(void *, Args &&...);
        using Manager = void (*)(void *, void *, Operation);

        union Storage {
            typename std::aligned_storage<MaxSize, alignof(max_align_t)>::type storage;
            void* ptr; // Used for oversized callables
        } data;

        template<typename F>
        static R invoke(void *data, Args &&... args) {
            F &f = *reinterpret_cast<F *>(&static_cast<Storage *>(data)->storage);
            return f(std::forward<Args>(args)...);
        }

        template<typename F>
        static void manage(void *dest, void *src, Operation op) {
            switch (op) {
                case Operation::Clone:
                    new(dest) F(*reinterpret_cast<F *>(src));
                    break;
                case Operation::Destroy:
                    reinterpret_cast<F *>(dest)->~F();
                    break;
            }
        }

        Invoker invoker = nullptr;
        Manager manager = nullptr;

    public:
        Function() noexcept {}

        Function(std::nullptr_t) noexcept {}

        Function(const Function& other) {
            if (other) {
                if (other.isStoredInternally()) {
                    // Directly copy the internal storage.
                    other.manager(data, other.data, Operation::Clone);
                } else {
                    // Allocate new memory and copy the external object.
                    void* ptr = internal::BlockAllocator::allocate(sizeof(f_type));
                    new(ptr) f_type(*static_cast<f_type*>(other.data.ptr));
                    data.ptr = ptr;
                }
                invoker = other.invoker;
                manager = other.manager;
            }
        }

        Function(Function&& other) noexcept : data(other.data), invoker(other.invoker), manager(other.manager) {
            other.invoker = nullptr;
            other.manager = nullptr;
            other.data.ptr = nullptr;
        }

        template<class F>
        Function(F &&f) : Function() {
            using f_type = std::decay_t<F>;
            if constexpr (sizeof(f_type) <= MaxSize) {
                new(&data.storage) f_type(std::forward<F>(f));
                invoker = [](void* data, Args&&... args) -> R {
                    return (*reinterpret_cast<f_type*>(&reinterpret_cast<Storage*>(data)->storage))(std::forward<Args>(args)...);
                };
                manager = [](void* dest, void* src, Operation op) {
                    auto srcPtr = reinterpret_cast<Storage*>(src);
                    if (op == Operation::Clone) {
                        new(&reinterpret_cast<Storage*>(dest)->storage) f_type(*reinterpret_cast<f_type*>(&srcPtr->storage));
                    } else if (op == Operation::Destroy) {
                        reinterpret_cast<f_type*>(&srcPtr->storage)->~f_type();
                    }
                };
            } else {
                void *ptr = internal::BlockAllocator::allocate(sizeof(f_type));
                new(ptr) f_type(std::forward<F>(f));
                data.ptr = ptr;
                invoker = [](void *data, Args &&... args) -> R {
                    return (*static_cast<f_type *>(static_cast<Storage *>(data)->ptr))(std::forward<Args>(args)...);
                };
                manager = [](void *dest, void *src, Operation op) {
                    auto srcPtr = static_cast<Storage *>(src)->ptr;
                    switch (op) {
                        case Operation::Clone: {
                            void *ptr = internal::BlockAllocator::allocate(sizeof(f_type));
                            new(ptr) f_type(*static_cast<f_type *>(srcPtr));
                            static_cast<Storage *>(dest)->ptr = ptr;
                            break;
                        }
                        case Operation::Destroy: {
                            f_type* obj = reinterpret_cast<f_type*>(&srcPtr->storage);
                            obj->~f_type();
                            internal::BlockAllocator::deallocate(obj, sizeof(f_type));
                            break;
                        }
                    }
                };
            }
        }

        ~Function() {
            if (manager) {
                manager(&data, nullptr, Operation::Destroy);
            }
        }

        R operator()(Args... args) const {
            return invoker(&data, std::forward<Args>(args)...);
        }

        explicit operator bool() const noexcept { return invoker != nullptr; }

        Function &operator=(const Function &other) {
            Function temp(other);
            temp.swap(*this);
            return *this;
        }

        void swap(Function &other) noexcept {
            std::swap(data, other.data);
            std::swap(invoker, other.invoker);
            std::swap(manager, other.manager);
        }
    };

namespace internal {
    class FreeListAllocator {
    public:
        FreeListAllocator(size_t totalSize) : totalSize(totalSize), used(0), peak(0) {
            memoryStart = ::operator new(totalSize); // Allocate memory block
            freeListHead = reinterpret_cast<FreeBlock*>(memoryStart);
            freeListHead->size = totalSize;
            freeListHead->next = nullptr;
        }

        ~FreeListAllocator() {
            ::operator delete(memoryStart); // Release memory block
        }

        void* allocate(size_t size, size_t alignment = alignof(std::max_align_t)) {
            FreeBlock *prev = nullptr, *bestFitPrev = nullptr, *bestFit = nullptr;
            size_t bestFitAdjustment = std::numeric_limits<size_t>::max();

            for (FreeBlock* block = freeListHead; block != nullptr; block = block->next) {
                size_t adjustment = alignAdjustment(reinterpret_cast<void*>(block + 1), alignment);
                size_t totalSizeNeeded = size + adjustment;

                if (block->size >= totalSizeNeeded && adjustment < bestFitAdjustment) {
                    bestFit = block;
                    bestFitPrev = prev;
                    bestFitAdjustment = adjustment;
                }

                prev = block;
            }

            if (bestFit == nullptr) {
                bool grown = grow(size + alignment); // Ensure there's space for alignment
                if (!grown) return nullptr; // Unable to grow, out of memory
                return allocate(size, alignment); // Try allocation again after growing
            }

            // Calculate total size needed for the requested allocation, including alignment
            size_t totalSizeNeeded = size + bestFitAdjustment;

            // Split block if remaining space is sufficiently large
            size_t remainingSize = bestFit->size - totalSizeNeeded;
            if (remainingSize > sizeof(FreeBlock)) {
                FreeBlock* newBlock = reinterpret_cast<FreeBlock*>(reinterpret_cast<char*>(bestFit) + totalSizeNeeded);
                newBlock->size = remainingSize;
                newBlock->next = bestFit->next;

                if (bestFitPrev != nullptr) {
                    bestFitPrev->next = newBlock;
                } else {
                    freeListHead = newBlock;
                }

                bestFit->size = totalSizeNeeded; // Adjust size of allocated block
            } else {
                // No splitting, adjust the linked list
                if (bestFitPrev != nullptr) {
                    bestFitPrev->next = bestFit->next;
                } else {
                    freeListHead = bestFit->next;
                }
            }

            // Adjust used and peak usage
            used += size + bestFitAdjustment; // Include alignment adjustment in usage calculation
            peak = std::max(peak, used);

            // Return the aligned address within the bestFit block
            return reinterpret_cast<void*>(reinterpret_cast<char*>(bestFit) + bestFitAdjustment);
        }

        void deallocate(void* ptr) {
            FreeBlock* blockToInsert = static_cast<FreeBlock*>(ptr);
            FreeBlock* iterator = freeListHead;
            FreeBlock* prev = nullptr;

            while (iterator != nullptr && iterator < blockToInsert) {
                prev = iterator;
                iterator = iterator->next;
            }

            // Coalesce with next block if adjacent
            if (reinterpret_cast<char*>(blockToInsert) + blockToInsert->size == reinterpret_cast<char*>(iterator)) {
                blockToInsert->size += iterator->size + sizeof(FreeBlock);
                blockToInsert->next = iterator->next;
            } else {
                blockToInsert->next = iterator;
            }

            // Coalesce with previous block if adjacent
            if (prev != nullptr && reinterpret_cast<char*>(prev) + prev->size == reinterpret_cast<char*>(blockToInsert)) {
                prev->size += blockToInsert->size + sizeof(FreeBlock);
                prev->next = blockToInsert->next;
            } else if (prev != nullptr) {
                prev->next = blockToInsert;
            } else {
                freeListHead = blockToInsert;
            }

            used -= blockToInsert->size;
        }

    private:
        struct FreeBlock {
            size_t size;
            FreeBlock* next;
        };

        FreeBlock* freeListHead = nullptr;
        void* memoryStart = nullptr;
        size_t totalSize = 0;
        size_t used = 0;
        size_t peak = 0;

        FreeListAllocator(const FreeListAllocator&) = delete;
        FreeListAllocator& operator=(const FreeListAllocator&) = delete;

        // Method to grow the allocator's memory pool
        bool grow(size_t size, size_t alignment = alignof(std::max_align_t)) {
            size_t newSize = std::max(size, totalSize); // Double the size or match the requested size.
            void* newBlockMemory = ::operator new(newSize);
            if (!newBlockMemory) return false; // Failed to allocate new block

            // Create a new FreeBlock at the beginning of the new memory block
            FreeBlock* newBlock = reinterpret_cast<FreeBlock*>(newBlockMemory);
            newBlock->size = newSize;
            newBlock->next = freeListHead;
            freeListHead = newBlock;

            totalSize += newSize; // Increase total size managed by the allocator
            return true;
        }

        size_t alignAdjustment(const void* ptr, size_t alignment) const {
            size_t adjustment = alignment - (reinterpret_cast<size_t>(ptr) & (alignment - 1));
            if (adjustment == alignment) return 0; // Already aligned
            return adjustment;
        }
    };

    class BlockAllocator {
    private:
        inline static FreeListAllocator* freeListAllocator = new FreeListAllocator(6 * 1024 * 1024);
    public:
        static void* allocate(size_t size) {
            return freeListAllocator->allocate(size);
        }

        static void deallocate(void* ptr, size_t size) { 
            // Since FreeListAllocator does not use the size on deallocation, we ignore it here as well
            freeListAllocator->deallocate(ptr);
        }
    };
}

} // namespace bruga::fn