MemMan 1.0.0.0 Released!

/**//**************************************************
 *
 * MemMan 1.0.0.0
 *
 * Copyright (C) 2007 - 2008 by Len3d
 * All rights reserved.
 *
 *************************************************/

#ifndef __MEM_MAN__
#define __MEM_MAN__



#define DEFAULT_MEMORY_ALIGN    16        //    default memory aligment set to this size
#define ALIGN_DEF                __declspec( align( DEFAULT_MEMORY_ALIGN ) )


#ifdef FORCEINLINE
#undef FORCEINLINE
#endif
#ifdef _DEBUG
#define FORCEINLINE        inline
#else
#define FORCEINLINE        __forceinline
#endif


typedef int    (*FreeContentFunc)( void *object, UINT requested_size );    //    can't be member function of a class


class MemoryAllocator {        //    memory allocator, takes charge of all system operations, also limits the usage of memory
private:
    //    elements of the memory priority queue, keep track of all memory blocks

    class MemoryQueueHeader {
    public:
        FORCEINLINE MemoryQueueHeader( UINT prior, void *obj, FreeContentFunc func, int al )
        {
            lchild = rchild = parent = NULL;
            priority = prior;
            object = obj;
            free_func = func;
            locked = TRUE;    //    we lock immediately after allocation
            align = static_cast<short>( al );
        }

        //    we don't need a deconstructor for this class

        FORCEINLINE void    lock()
        {
            locked = TRUE;
        }

        FORCEINLINE void    unlock()
        {
            locked = FALSE;
        }

        FORCEINLINE void    create_child( MemoryQueueHeader *key )
        {
            if( key->priority < priority )
            {
                if( lchild )
                    lchild->create_child( key );
                else
                {
                    lchild = key;
                    lchild->parent = this;
                }
            }
            else
            {
                if( rchild )
                    rchild->create_child( key );
                else
                {
                    rchild = key;
                    rchild->parent = this;
                }
            }
        }

        FORCEINLINE bool    search_memory( UINT search_size, int search_align, void * & search_result, 
                                            void *obj, FreeContentFunc func )
        {
            if( lchild && lchild->search_memory( search_size, search_align, search_result, obj, func ) )
                return true;

            if( align == search_align && free_content( search_size ) )
            {
                search_result = get_memory();

                object = obj;    //    update the attributes of the memory block
                free_func = func;

                return true;
            }

            if( rchild && rchild->search_memory( search_size, search_align, search_result, obj, func ) )
                return true;

            return false;
        }

        FORCEINLINE void    *get_memory()    //    the allocated memory block
        {
            return ((char *)this + sizeof(MemoryQueueHeader));
        }

        FORCEINLINE int        free_content( UINT requested_size )
        {
            if( !locked && free_func && object )
                return free_func( object, requested_size );
            else
                return 0;
        }

    public:
        ALIGN_DEF struct {
            MemoryQueueHeader    *lchild,    //    left child, right child and parent
                                *rchild,
                                *parent;
            UINT                priority;    //    priority for sorting the memory blocks
            void                *object;    //    the object for which the memory was allocated
            FreeContentFunc        free_func;    //    function to free the content of the object for requested size,
                                            //    memory blocks without this function will be restored to memory-mapped files.
            short                locked;        //    this memory block was locked by a thread
            short                align;        //    aligment of the allocated memory should match
            int                    pad;        //    padded to 32 Byte aligned
        };
    };

public:
    MemoryAllocator( UINT max_size )    //    max allowed memory usage
    {
        allocated_size = 0;
        available_size = max_size;

        queue = NULL;
        aligned_queue = NULL;
    }

    ~MemoryAllocator()
    {
        dealloc_header( queue );
        dealloc_aligned_header( aligned_queue );
    }

    FORCEINLINE void    *alloc( UINT size, UINT priority, 
                                void *object, FreeContentFunc free_func )
    {
        if( size == 0 )
        {
            return NULL;
        }
        else if( size > available_size )        //    searching has the complexity of O(N)
        {
            void    *ptr = NULL;

            if( queue && queue->search_memory( size, 0, ptr, object, free_func ) )
                return ptr;
            else
                return NULL;    //    the system has run out of memory
        }
        else    //    the complexity is O(logN)
        {
            allocated_size += ( size + sizeof(MemoryQueueHeader) );
            available_size -= ( size + sizeof(MemoryQueueHeader) );

            MemoryQueueHeader    *elem;

            //    allocate a block
            elem = new (sys_alloc( sizeof(MemoryQueueHeader) + size )) MemoryQueueHeader( priority, object, free_func, 0 );

            if( queue )
                queue->create_child( elem );    //    insert the node
            else
                queue = elem;    //    be the root

            return elem->get_memory();
        }
    }

    FORCEINLINE void    *aligned_alloc( UINT size, UINT priority, 
                                        void *object, FreeContentFunc free_func, 
                                        UINT align = DEFAULT_MEMORY_ALIGN )
    {
        if( size == 0 )
        {
            return NULL;
        }
        else if( size > available_size )        //    searching has the complexity of O(N)
        {
            void    *ptr = NULL;

            if( aligned_queue && aligned_queue->search_memory( size, align, ptr, object, free_func ) )
                return ptr;
            else
                return NULL;    //    the system has run out of memory
        }
        else    //    the complexity is O(logN)
        {
            allocated_size += ( size + sizeof(MemoryQueueHeader) );
            available_size -= ( size + sizeof(MemoryQueueHeader) );

            MemoryQueueHeader    *elem;

            //    allocate an aligned block
            elem = new (sys_aligned_alloc( sizeof(MemoryQueueHeader) + size, align )) MemoryQueueHeader( priority, object, free_func, align );

            if( aligned_queue )
                aligned_queue->create_child( elem );    //    insert the node
            else
                aligned_queue = elem;    //    be the root

            return elem->get_memory();
        }
    }

    //    a lock must be used before the object being deallocated, the complexity is O(1)
    
    FORCEINLINE void    lock( void *ptr )
    {
        if( ptr )
        {
            MemoryQueueHeader    *header = (MemoryQueueHeader *)((char *) ptr - sizeof(MemoryQueueHeader));

            header->lock();
        }
    }

    FORCEINLINE void    unlock( void *ptr )
    {
        if( ptr )
        {
            MemoryQueueHeader    *header = (MemoryQueueHeader *)((char *) ptr - sizeof(MemoryQueueHeader));

            header->unlock();
        }
    }

    //    deallocating has the complexity of O(logN)

    FORCEINLINE void    dealloc( void *ptr )
    {
        if( ptr )
        {
            MemoryQueueHeader    *header, *node, *parent;

            header = (MemoryQueueHeader *)((char *) ptr - sizeof(MemoryQueueHeader));
            parent = header->parent;

            if( header->lchild )    //    has left child
            {
                node = find_rightmost_child( header->lchild );

                //    rebuild the link
                if( node != header->lchild )
                {
                    if( node->parent )
                        node->parent->rchild = NULL;    //    clear the link

                    node->lchild = header->lchild;
                    if( node->lchild )
                        node->lchild->parent = node;
                }
                node->rchild = header->rchild;
                if( node->rchild )
                    node->rchild->parent = node;
                node->parent = parent;

                if( parent )    //    has parent
                {
                    if( parent->lchild == header )
                        parent->lchild = node;
                    else if( parent->rchild == header )
                        parent->rchild = node;
                }
                else    //    it's the root
                {
                    queue = node;
                }
            }
            else if( header->rchild )    //    has right child
            {
                node = find_leftmost_child( header->rchild );

                //    rebuild the link
                if( node != header->rchild )
                {
                    if( node->parent )
                        node->parent->lchild = NULL;    //    clear the link

                    node->rchild = header->rchild;
                    if( node->rchild )
                        node->rchild->parent = node;
                }
                node->lchild = header->lchild;
                if( node->lchild )
                    node->lchild->parent = node;
                node->parent = parent;

                if( parent )    //    has parent
                {
                    if( parent->lchild == header )
                        parent->lchild = node;
                    else if( parent->rchild == header )
                        parent->rchild = node;
                }
                else    //    it's the root
                {
                    queue = node;
                }
            }
            else    //    has no children
            {
                if( parent )    //    has parent
                {
                    if( parent->lchild == header )
                        parent->lchild = NULL;
                    else if( parent->rchild == header )
                        parent->rchild = NULL;
                }
                else    //    it's the root, clear it
                {
                    queue = NULL;
                }
            }

            sys_dealloc( header );    //    deallocate the block
        }
    }

    FORCEINLINE void    aligned_dealloc( void *ptr )
    {
        if( ptr )
        {
            MemoryQueueHeader    *header, *node, *parent;

            header = (MemoryQueueHeader *)((char *) ptr - sizeof(MemoryQueueHeader));
            parent = header->parent;

            if( header->lchild )    //    has left child
            {
                node = find_rightmost_child( header->lchild );

                //    rebuild the link
                if( node != header->lchild )
                {
                    if( node->parent )
                        node->parent->rchild = NULL;    //    clear the link

                    node->lchild = header->lchild;
                    if( node->lchild )
                        node->lchild->parent = node;
                }
                node->rchild = header->rchild;
                if( node->rchild )
                    node->rchild->parent = node;
                node->parent = parent;

                if( parent )    //    has parent
                {
                    if( parent->lchild == header )
                        parent->lchild = node;
                    else if( parent->rchild == header )
                        parent->rchild = node;
                }
                else    //    it's the root
                {
                    queue = node;
                }
            }
            else if( header->rchild )    //    has right child
            {
                node = find_leftmost_child( header->rchild );

                //    rebuild the link
                if( node != header->rchild )
                {
                    if( node->parent )
                        node->parent->lchild = NULL;    //    clear the link

                    node->rchild = header->rchild;
                    if( node->rchild )
                        node->rchild->parent = node;
                }
                node->lchild = header->lchild;
                if( node->lchild )
                    node->lchild->parent = node;
                node->parent = parent;

                if( parent )    //    has parent
                {
                    if( parent->lchild == header )
                        parent->lchild = node;
                    else if( parent->rchild == header )
                        parent->rchild = node;
                }
                else    //    it's the root
                {
                    queue = node;
                }
            }
            else    //    has no children
            {
                if( parent )    //    has parent
                {
                    if( parent->lchild == header )
                        parent->lchild = NULL;
                    else if( parent->rchild == header )
                        parent->rchild = NULL;
                }
                else    //    it's the root, clear it
                {
                    queue = NULL;
                }
            }

            sys_aligned_dealloc( header );    //    deallocate the block
        }
    }

private:
    //    help functions

    FORCEINLINE void    dealloc_header( MemoryQueueHeader *node )
    {
        if( node )
        {
            if( node->lchild )
                dealloc_header( node->lchild );

            if( node->rchild )
                dealloc_header( node->rchild );

            sys_dealloc( node );
        }
    }

    FORCEINLINE void    dealloc_aligned_header( MemoryQueueHeader *node )
    {
        if( node )
        {
            if( node->lchild )
                dealloc_aligned_header( node->lchild );

            if( node->rchild )
                dealloc_aligned_header( node->rchild );

            sys_aligned_dealloc( node );
        }
    }

    FORCEINLINE MemoryQueueHeader    *find_leftmost_child( MemoryQueueHeader *node )
    {
        if( node && node->lchild )
            return find_leftmost_child( node->lchild );
        else
            return node;
    }

    FORCEINLINE MemoryQueueHeader    *find_rightmost_child( MemoryQueueHeader *node )
    {
        if( node && node->rchild )
            return find_rightmost_child( node->rchild );
        else
            return node;
    }

private:
    //    encapsulate system operations

    FORCEINLINE void    *sys_alloc( UINT size )
    {
        return malloc( size );
    }

    FORCEINLINE void    *sys_aligned_alloc( UINT size, UINT align )
    {
        return _mm_malloc( size, align );
    }

    FORCEINLINE void    sys_dealloc( void *ptr )
    {
        free( ptr );
    }

    FORCEINLINE void    sys_aligned_dealloc( void *ptr )
    {
        _mm_free( ptr );
    }

private:
    //    memory statistics
    UINT                            allocated_size, 
                                    available_size;

    //    implement priority queues to record all the allocations
    MemoryQueueHeader                *queue;
    MemoryQueueHeader                *aligned_queue;
};



#endif    //    __MEM_MAN__