#pragma once

#include <util/generic/ptr.h>
#include <util/memory/pool.h>

template <class T> class TOptionalAlloc;
class TOptionalAllocator {
public:
    typedef size_t size_type;
    typedef ptrdiff_t difference_type;
    typedef void value_type;

    template <class T>
    struct rebind {
        typedef TOptionalAlloc<T> other;
    };

    TOptionalAllocator(bool usePool = false)
        : MemPool(usePool ? new TMemoryPool(1024) : nullptr)
    {}

    inline void* AllocateMem(size_t n) {
        if (!!MemPool)
            return MemPool->Allocate(n);
        else
            return new char[n];
    }

    template<class T1>
    inline void deallocate(T1* p, size_t /*n*/) {
        if (!MemPool)
            delete [] p;
    }

    inline TMemoryPool* GetPool() {
        return MemPool.Get();
    }

    inline const TMemoryPool* GetPool() const {
        return MemPool.Get();
    }

    friend inline bool operator==(const TOptionalAllocator& l, const TOptionalAllocator& r) {
        return l.MemPool == r.MemPool;
    }

    friend inline bool operator!=(const TOptionalAllocator& l, const TOptionalAllocator& r) {
        return !(l == r);
    }

private:
    TAtomicSharedPtr<TMemoryPool> MemPool;
};

template <class T>
class TOptionalAlloc : public TOptionalAllocator {
public:

    typedef T* pointer;
    typedef const T* const_pointer;
    typedef T& reference;
    typedef const T& const_reference;
    typedef T value_type;

    TOptionalAlloc(const TOptionalAllocator& other)
        : TOptionalAllocator(other)
    {}

    inline T* allocate(size_t n) {
        return (T*)AllocateMem(n * sizeof(T));
    }
};

