/*
=============================================================================
    File:   Memory.h
    Desc:   Memory management.
    ToDo:   Memory limits and watchdogs, thread-safety,
            allow users to define their own OutOfMemory handlers, etc.
=============================================================================
*/
#pragma once

//@todo: remove it from here
#include <Base/Object/EnumType.h>   // for MemoryHeapT
//#include <Base/Template/SmartPtr/TPtr.h>
#include <Base/Template/Containers/Array/TStaticArray.h>

//---------------------------------------------------------------------------
//      Defines.
//---------------------------------------------------------------------------

// 1 - Track each memory allocation/deallocation, record memory usage statistics and detect memory leaks and so on.
#define MX_DEBUG_MEMORY         MX_DEBUG

// 1 - Redirect the global 'new' and 'delete' to our memory manager.
// this is dangerous, it may be better to define them on per-type basis or use special macros.
#define MX_OVERRIDE_GLOBAL_NEWDELETE    (0)

// 1 - prevent the programmer from using malloc/free.
#define MX_HIDE_MALLOC_AND_FREE     (0)

//---------------------------------------------------------------------------
//      Declarations.
//---------------------------------------------------------------------------

class mxClass;

// memory heaps/arenas/areas
// split by size/lifetime
enum EMemoryHeap
{
#define DECLARE_MEMORY_HEAP( NAME )     Heap##NAME
#include <Base/Memory/MemoryHeaps.inl>
#undef DECLARE_MEMORY_HEAP

    HeapCount,  //!<= Marker. Don't use.
};
mxDECLARE_ENUM( EMemoryHeap, UINT8, MemoryHeapT );

mxSTOLEN("based on Bitsquid foundation library");
/// Base class for memory allocators.
///
/// Note: Regardless of which allocator is used, prefer to allocate memory in larger chunks
/// instead of in many small allocations. This helps with data locality, fragmentation,
/// memory usage tracking, etc.
///
/// NOTE: ideally, 'dynamically-resizing' containers and the associated memory manager
/// should know about each other to efficiently work together.
/// We could pass the size of allocated memory to the Free() function (like STL allocators),
/// but this optimization is usually avoided in practice (a new source of bugs).
///
/// Memory allocations/deallocations are expensive operations
/// so use of virtual member functions should be ok.
///
struct AHeap : NonCopyable
{
    /// Allocates the specified amount of memory aligned to the specified alignment.
    virtual void* Allocate( U32 _bytes, U32 _alignment ) = 0;

    /// Frees an allocation previously made with Allocate().
    ///@todo: pass the allocated size as an optimization?
    virtual void Deallocate( void* _memory ) = 0;

    /// Returns the amount of usable memory allocated at _memory. _memory must be a pointer
    /// returned by Allocate() that has not yet been deallocated. The value returned
    /// will be at least the size specified to Allocate(), but it can be bigger.
    /// (The allocator may round up the allocation to fit into a set of predefined
    /// slot sizes.)
    ///
    /// Not all allocators support tracking the size of individual allocations.
    /// An allocator that doesn't support it will return SIZE_NOT_TRACKED.
    virtual U32 allocated_size( const void* _memory ) const
    {
        return SIZE_NOT_TRACKED;
    }

public:
    /// Returns the total amount of memory allocated by this allocator. Note that the
    /// size returned can be bigger than the size of all individual allocations made,
    /// because the allocator may keep additional structures.
    ///
    /// If the allocator doesn't track memory, this function returns SIZE_NOT_TRACKED.
    virtual U32 total_allocated() const
    {
        return SIZE_NOT_TRACKED;
    }

    virtual void Validate() {}; // verify heap (debug mode only)
    virtual void Optimize() {}; // compact heap

    /// writes memory statistics, usage & leak info to the specified file
    virtual void Dump( ALog& log ) {};

    /// \return The name of the allocator implementing this interface
    //virtual const char* GetName() const { return mxSTRING_Unknown; }

    /// A return value indicating that 'the size is not tracked'.
    static const U32 SIZE_NOT_TRACKED = 0;

protected:
    AHeap() {}
    virtual ~AHeap() {}

    /// Allocators cannot be copied.
    PREVENT_COPY(AHeap);
};

extern TStaticArray< AHeap*, HeapCount >    g_heaps;

ERet Memory_Initialize();
void Memory_Shutdown(); // shutdown and assert on memory leaks

namespace Heaps
{

/// thread-safe ring buffer allocator, can be used for passing data between threads
inline AHeap& scratch() { return *g_heaps[HeapTemporary]; }

/// default, generic heap
inline AHeap& global() { return *g_heaps[HeapGeneric]; }

void DumpStats();

}//namespace Heaps

//---------------------------------------------------------------------------
//      HELPER MACROS
//---------------------------------------------------------------------------

#define mxALLOC( size, align )          (g_heaps[HeapGeneric]->Allocate( (size), (align) ))
#define mxFREE( memory, size )          (g_heaps[HeapGeneric]->Deallocate( (memory)/*, (size)*/ ))

/// declare our own 'new' and 'delete' so that these can be found via 'find All References'

#if defined(new_one) || defined(free_one)
#   error 'new_one' and 'free_one' have already been  defined - shouldn't happen!
#endif

#define new_one( x )    new x
#define free_one( x )   delete x

// Array operators

#if defined(new_array) || defined(free_array)
#   error 'new_array' and 'free_array' have already been defined - shouldn't happen!
#endif

#define new_array( x, num )     new x [num]
#define free_array( x )         delete[] x

#if MX_OVERRIDE_GLOBAL_NEWDELETE
    #error Incompatible options: overriden global 'new' and 'delete' and per-class memory heaps.
#endif

//
//  mxDECLARE_CLASS_ALLOCATOR
//
//  'className' can be used to track instances of the class.
//
#define mxDECLARE_CLASS_ALLOCATOR( memClass, className )\
public:\
    typedef className THIS_TYPE;    \
    static EMemoryHeap GetHeap() { return memClass; }   \
    mxFORCEINLINE void* operator new      ( size_t sizeInBytes ){ return mxAllocX( memClass, sizeInBytes ); }   \
    mxFORCEINLINE void  operator delete   ( void* ptr )         { mxFreeX( memClass, ptr ); }                   \
    mxFORCEINLINE void* operator new      ( size_t, void* ptr ) { return ptr; }                                 \
    mxFORCEINLINE void  operator delete   ( void*, void* )      { }                                             \
    mxFORCEINLINE void* operator new[]    ( size_t sizeInBytes ){ return mxAllocX( memClass, sizeInBytes ); }   \
    mxFORCEINLINE void  operator delete[] ( void* ptr )         { mxFreeX( memClass, ptr ); }                   \
    mxFORCEINLINE void* operator new[]    ( size_t, void* ptr ) { return ptr; }                                 \
    mxFORCEINLINE void  operator delete[] ( void*, void* )      { }                                             \

//
//  mxDECLARE_VIRTUAL_CLASS_ALLOCATOR
//
#define mxDECLARE_VIRTUAL_CLASS_ALLOCATOR( memClass, className )\
public:\
    typedef className THIS_TYPE;    \
    mxFORCEINLINE void* operator new      ( size_t sizeInBytes ){ return mxAllocX( memClass, sizeInBytes ); }   \
    mxFORCEINLINE void  operator delete   ( void* ptr )         { mxFreeX( memClass, ptr ); }                   \
    mxFORCEINLINE void* operator new      ( size_t, void* ptr ) { return ptr; }                                 \
    mxFORCEINLINE void  operator delete   ( void*, void* )      { }                                             \
    mxFORCEINLINE void* operator new[]    ( size_t sizeInBytes ){ return mxAllocX( memClass, sizeInBytes ); }   \
    mxFORCEINLINE void  operator delete[] ( void* ptr )         { mxFreeX( memClass, ptr ); }                   \
    mxFORCEINLINE void* operator new[]    ( size_t, void* ptr ) { return ptr; }                                 \
    mxFORCEINLINE void  operator delete[] ( void*, void* )      { }

//
//  mxDECLARE_NONVIRTUAL_CLASS_ALLOCATOR
//
#define mxDECLARE_NONVIRTUAL_CLASS_ALLOCATOR( memClass, className )\
public:\
    typedef className THIS_TYPE;    \
    mxFORCEINLINE void* operator new      ( size_t sizeInBytes ){ return mxAllocX( memClass, sizeInBytes ); }   \
    mxFORCEINLINE void  operator delete   ( void* ptr )         { mxFreeX( memClass, ptr ); }                   \
    mxFORCEINLINE void* operator new      ( size_t, void* ptr ) { return ptr; }                                 \
    mxFORCEINLINE void  operator delete   ( void*, void* )      { }                                             \
    mxFORCEINLINE void* operator new[]    ( size_t sizeInBytes ){ return mxAllocX( memClass, sizeInBytes ); }   \
    mxFORCEINLINE void  operator delete[] ( void* ptr )         { mxFreeX( memClass, ptr ); }                   \
    mxFORCEINLINE void* operator new[]    ( size_t, void* ptr ) { return ptr; }                                 \
    mxFORCEINLINE void  operator delete[] ( void*, void* )      { }

//---------------------------------------------------------------------------
//      Prevent usage of plain old 'malloc' & 'free' if needed.
//---------------------------------------------------------------------------

#if MX_HIDE_MALLOC_AND_FREE

    #define malloc( size )          ptBREAK
    #define free( mem )             ptBREAK
    #define calloc( num, size )     ptBREAK
    #define realloc( mem, newsize ) ptBREAK

#endif // MX_HIDE_MALLOC_AND_FREE

//------------------------------------------------------------------------
//  Useful macros
//------------------------------------------------------------------------

#ifndef SAFE_DELETE
#define SAFE_DELETE( p )        { if( p != nil ) { delete (p);     (p) = nil; } }
#endif

#ifndef SAFE_DELETE_ARRAY
#define SAFE_DELETE_ARRAY( p )  { if( p != nil ) { delete[] (p);   (p) = nil; } }
#endif

#define mxNULLIFY_POINTER(p)    {(p) = nil;}



#if 0 // mxSUPPORTS_VARIADIC_TEMPLATES

    /// Creates a new object of the given type using the specified allocator.
    template< typename TYPE, typename... ARGs >
    TYPE* mxNEW( AHeap& _heap, ARGs&&... _args )
    {
        void * storage = _heap.Allocate( sizeof(TYPE), mxALIGNOF(TYPE) );
        return new(storage) TYPE( std::forward< ARGs >(_args)... );
    }

#else

    /// Creates a new object of the given type using the specified allocator.
    #define mxNEW( HEAP, TYPE, ... )\
        (new\
            ((HEAP).Allocate( sizeof(TYPE), GET_ALIGNMENT_FOR_CONTAINERS(TYPE) ))\
                TYPE(__VA_ARGS__))

#endif

    /// Frees an object allocated with mxNEW.
    template< typename TYPE, class HEAP >
    void mxDELETE( HEAP& _heap, TYPE* _o )
    {
        if( _o != nullptr ) {
            _o->~TYPE();
            _heap.Deallocate( _o );
        }
    }

///
#define mxTRY_ALLOC_SCOPED( _VAR, _ARRAY_COUNT, _ALLOCATOR )\
    mxDO( __TryAllocate( _VAR, _ARRAY_COUNT, _ALLOCATOR ) );\
    AutoFree __autofree##_VAR##__LINE__( _ALLOCATOR, _VAR );

/// internal function used by the mxTRY_ALLOC_SCOPED macro
template< typename TYPE >
ERet __TryAllocate( TYPE *& _pointer, U32 _count, AHeap &_allocator )
{
    TYPE* p = static_cast< TYPE* >( _allocator.Allocate( sizeof(TYPE)*_count, GET_ALIGNMENT_FOR_CONTAINERS(TYPE) ) );
    if( !p ) {
        return ERR_OUT_OF_MEMORY;
    }
    _pointer = p;
    return ALL_OK;
}

/// automatically frees pointed memory; used internally by the mxTRY_ALLOC_SCOPED macro
class AutoFree {
    AHeap & m_heap;
    void *  m_memory;
public:
    AutoFree( AHeap & _heap, void* _memory )
        : m_heap( _heap ), m_memory( _memory )
    {}
    ~AutoFree() {
        m_heap.Deallocate( m_memory );
    }
};

/// renders any heap thread-safe
class ThreadSafeProxyHeap : public AHeap
{
    mutable SpinWait    m_lock;
    AHeap & m_heap; //!< the actual allocator we're using under the hood
public:
    ThreadSafeProxyHeap( AHeap &_heap )
        : m_heap( _heap )
    {
    }
    ERet Initialize()
    {
        m_lock.Initialize();
        return ALL_OK;
    }
    void Shutdown()
    {
        m_lock.Shutdown();
    }
    virtual void* Allocate( U32 _bytes, U32 _alignment ) override
    {
        mxASSERT(_bytes > 0);
        SpinWait::Lock  scopedLock( m_lock );
        void* result = m_heap.Allocate( _bytes, _alignment );
        return result;
    }
    virtual void Deallocate( void* _memory ) override
    {
        if( _memory ) {
            SpinWait::Lock  scopedLock( m_lock );
            m_heap.Deallocate( (void*)_memory );
        }
    }
    virtual U32 allocated_size( const void* _memory ) const
    {
        SpinWait::Lock  scopedLock( m_lock );
        return m_heap.allocated_size( _memory );
    }
};

/// _aligned_malloc/_aligned_free
class DefaultHeap : public AHeap
{
public:
    DefaultHeap();
    virtual void* Allocate( U32 _bytes, U32 _alignment ) override;
    virtual void Deallocate( void* _memory ) override;
    virtual U32 allocated_size( const void* _memory ) const override;
};

  /// A proxy allocator that simply draws upon another allocator.
  /// http://bitsquid.blogspot.ru/2010/09/custom-memory-allocation-in-c.html
  class TrackingHeap : public AHeap {
      AHeap & m_allocator;
  public:
      TrackingHeap( const char* _name, AHeap & _allocator )
      : m_allocator( _allocator )
    {}
      virtual void* Allocate( U32 _bytes, U32 _alignment ) override {
      return m_allocator.Allocate( _bytes, _alignment );
    }
      virtual void Deallocate( void* _memory ) override {
      m_allocator.Deallocate( _memory );
    }
      virtual U32   allocated_size( const void* _memory ) const override {
      return m_allocator.allocated_size( _memory );
    }
  };

//--------------------------------------------------------------//
//              End Of File.                                    //
//--------------------------------------------------------------//