MemMan 1.0.0.2 Released!

/**//**************************************************
 *
 * MemMan 1.0.0.2
 *
 * 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 ) )


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



#pragma pack(push,1)



class MemoryAllocator {        //    memory allocator, takes charge of all system operations, also limits the usage of memory
private:
    typedef UINT    MemoryHeaderBias;    //    biased value type of memory header

    //    elements of the memory priority queue, keep track of all memory blocks

    class MemoryHeader {
    public:
        FORCEINLINE MemoryHeader( UINT prior, void *obj, FreeContentFunc func, UINT al )
        {
            lchild = rchild = parent = NULL;
            priority = prior;
            object = obj;
            free_func = func;
            locked = TRUE;    //    we lock immediately after allocation
            align = al;
            ref_count = 1;    //    assume one block references this header at beginning
            bias = sizeof(MemoryHeader);
            bfact = 0;    //    initialize balance factor with 0
        }

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

        FORCEINLINE void    lock()
        {
            locked = TRUE;
        }

        FORCEINLINE void    unlock()
        {
            locked = FALSE;
        }

        FORCEINLINE MemoryHeader    *insert( MemoryHeader *key )
        {
            MemoryHeader    *root = this;
            bool            taller = true;

            create_child( key, root, taller );

            return root;    //    AVL may change the root, we return the new root here
        }

        FORCEINLINE void    create_child( MemoryHeader *key, MemoryHeader * & root, bool & taller )
        {
            if( key->priority < priority )
            {
                if( lchild )
                    lchild->create_child( key, root, taller );
                else
                {
                    lchild = key;
                    lchild->parent = this;
                    taller = true;
                }

                if( taller )
                {
                    switch( bfact )
                    {
                    case 1:
                        left_balance( root );
                        taller = false;
                        break;
                    case 0:
                        bfact = 1;
                        taller = true;
                        break;
                    case -1:
                        bfact = 0;
                        taller = false;
                        break;
                    }
                }
            }
            else
            {
                if( rchild )
                    rchild->create_child( key, root, taller );
                else
                {
                    rchild = key;
                    rchild->parent = this;
                    taller = true;
                }

                if( taller )
                {
                    switch( bfact )
                    {
                    case 1:
                        bfact = 0;
                        taller = false;
                        break;
                    case 0:
                        bfact = -1;
                        taller = true;
                        break;
                    case -1:
                        right_balance( root );
                        taller = false;
                        break;
                    }
                }
            }
        }

        FORCEINLINE void    left_rotate( MemoryHeader * & root )
        {
            MemoryHeader    *rc;

            rc = rchild;
            
            rchild = rc->lchild;
            if( rchild )
                rchild->parent = this;

            rc->lchild = this;
            rc->parent = parent;

            if( parent )
            {
                if( parent->lchild == this )
                    parent->lchild = rc;
                else if( parent->rchild == this )
                    parent->rchild = rc;
            }
            else
            {
                root = rc;
            }

            parent = rc;
        }

        FORCEINLINE void    right_rotate( MemoryHeader * & root )
        {
            MemoryHeader    *lc;

            lc = lchild;
            
            lchild = lc->rchild;
            if( lchild )
                lchild->parent = this;

            lc->rchild = this;
            lc->parent = parent;

            if( parent )
            {
                if( parent->lchild == this )
                    parent->lchild = lc;
                else if( parent->rchild == this )
                    parent->rchild = lc;
            }
            else
            {
                root = lc;
            }

            parent = lc;
        }

        FORCEINLINE void    left_balance( MemoryHeader * & root )
        {
            MemoryHeader    *lc, *rd;
            
            lc = lchild;

            switch( lc->bfact )
            {
            case 1:
                bfact = lc->bfact = 0;
                
                right_rotate( root );
                
                break;

            case 0:
                rd = lc->rchild;

                switch( rd->bfact )
                {
                case 1:
                    bfact = -1;
                    lc->bfact = 0;
                    break;
                case 0:
                    bfact = lc->bfact = 0;
                    break;
                case -1:
                    bfact = 0;
                    lc->bfact = 1;
                    break;
                }

                rd->bfact = 0;

                lchild->left_rotate( root );
                right_rotate( root );

                break;
            }
        }

        FORCEINLINE void    right_balance( MemoryHeader * & root )
        {
            MemoryHeader    *rc, *ld;
            
            rc = rchild;

            switch( rc->bfact )
            {
            case 1:
                bfact = rc->bfact = 0;
                
                left_rotate( root );
                
                break;

            case 0:
                ld = rc->lchild;

                switch( ld->bfact )
                {
                case 1:
                    bfact = -1;
                    rc->bfact = 0;
                    break;
                case 0:
                    bfact = rc->bfact = 0;
                    break;
                case -1:
                    bfact = 0;
                    rc->bfact = 1;
                    break;
                }

                ld->bfact = 0;

                rchild->right_rotate( root );
                left_rotate( root );

                break;
            }
        }

        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(MemoryHeader));
        }

        FORCEINLINE void    dec_ref_count()    //    decrease the reference count
        {
            -- ref_count;
        }

        FORCEINLINE UINT    get_ref_count()    //    the reference count
        {
            return ref_count;
        }

        //    try to free some content of the object for requested size,
        //    update the reference count and return the size of the freed memory
        FORCEINLINE int        free_content( UINT requested_size )
        {
            if( !locked && free_func && object )
                return free_func( object, requested_size, ref_count );
            else
                return 0;
        }

    public:
        ALIGN_DEF struct {    //    48 Byte aligned including the header bias
            MemoryHeader        *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.
            int                    locked;        //    this memory block was locked by a thread
            int                    bfact;        //    balance factor of the AVL tree
            UINT                align;        //    aligment of the allocated memory should match
            UINT                ref_count;    //    how many blocks reference this header
            UINT                pad;        //    padded to 48 Byte, no use
            MemoryHeaderBias    bias;        //    the header bias
        };
    };

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 += ( sizeof(MemoryHeader) + size );
            available_size -= ( sizeof(MemoryHeader) + size );

            MemoryHeader    *elem;

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

            if( queue )
                queue = queue->insert( 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 += ( sizeof(MemoryHeader) + size );
            available_size -= ( sizeof(MemoryHeader) + size );

            MemoryHeader    *elem;

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

            if( aligned_queue )
                aligned_queue = aligned_queue->insert( 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 )
        {
            MemoryHeader    *header = get_memory_header( ptr );

            header->lock();
        }
    }

    FORCEINLINE void    unlock( void *ptr )
    {
        if( ptr )
        {
            MemoryHeader    *header = get_memory_header( ptr );

            header->unlock();
        }
    }

    //    deallocating has the complexity of O(logN)

    FORCEINLINE void    dealloc( void *ptr, UINT size )
    {
        if( ptr )
        {
            MemoryHeader    *header, *node, *parent;

            header = get_memory_header( ptr );

            header->dec_ref_count();

            if( header->get_ref_count() != 0 )    //    still have objects reference this
                return;

            parent = header->parent;

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

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

                    node->lchild = header->lchild;
                    node->lchild->parent = node;
                }

                node->rchild = header->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->lchild )    //    has only left child
            {
                //    rebuild the links
                node = header->lchild;
                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 only right child
            {
                //    rebuild the links
                node = header->rchild;
                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 child
            {
                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;
                }
            }

            allocated_size -= ( sizeof(MemoryHeader) + size );
            available_size += ( sizeof(MemoryHeader) + size );

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

    FORCEINLINE void    aligned_dealloc( void *ptr, UINT size )
    {
        if( ptr )
        {
            MemoryHeader    *header, *node, *parent;

            header = get_memory_header( ptr );

            header->dec_ref_count();

            if( header->get_ref_count() != 0 )    //    still have objects reference this
                return;

            parent = header->parent;

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

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

                    node->lchild = header->lchild;
                    node->lchild->parent = node;
                }

                node->rchild = header->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
                {
                    aligned_queue = node;
                }
            }
            else if( header->lchild )    //    has only left child
            {
                //    rebuild the links
                node = header->lchild;
                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
                {
                    aligned_queue = node;
                }
            }
            else if( header->rchild )    //    has only right child
            {
                //    rebuild the links
                node = header->rchild;
                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
                {
                    aligned_queue = node;
                }
            }
            else    //    has no child
            {
                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
                {
                    aligned_queue = NULL;
                }
            }

            allocated_size -= ( sizeof(MemoryHeader) + size );
            available_size += ( sizeof(MemoryHeader) + size );

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

private:
    //    help functions

    FORCEINLINE void    dealloc_header( MemoryHeader *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( MemoryHeader *node )
    {
        if( node )
        {
            if( node->lchild )
                dealloc_aligned_header( node->lchild );

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

            sys_aligned_dealloc( node );
        }
    }

    FORCEINLINE MemoryHeader    *find_rightmost_child( MemoryHeader *node )
    {
        while( node && node->rchild )
        {
            node = node->rchild;
        }

        return node;
    }

    FORCEINLINE MemoryHeader    *get_memory_header( void *ptr )
    {
        //    get the biased value first
        MemoryHeaderBias    *pb = (MemoryHeaderBias *)((char *) ptr - sizeof(MemoryHeaderBias));

        //    then get the header using the biased value
        MemoryHeader    *header = (MemoryHeader *)((char *) ptr - (*pb));

        return header;
    }

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
    MemoryHeader                *queue;
    MemoryHeader                *aligned_queue;
};


#pragma pack(pop)



#endif    //    __MEM_MAN__