自己动手实现allocator

最近在追旧番《STL代码剖析》。真的是很旧很旧的番了，STL在94年开始走入STL，这本书则是2002年出版的，C++03和C++11还不知何在的年代。看完第二章之后合上书，想自己写一个allocator。发现看书过程中自认为“所言极是”的地方，居然根本写不出来。虽然从前也写过内存池 (mempool)，但是STL的手法和我等闲辈果然不是一个层次。于是只好边写边翻书，也算顺利得写出来了一个不支持fill和copy的版本。毕竟，stl_uninitialized相对于stl_construct和stl_alloc简单多了，我个人也不常用它来初始化容器。

为了方便，就写在了一个文件里，就请假装他们是分开写的吧。。。（PS：需要-std_c++11哦，因为我是在受不了C98里面煞笔的尖括号问题！）

（顺便回顾回顾auto_ptr、placement new、#pragma pack(n)这些最近才长的知识。）

/* alloc.h */
#include <cstdlib> //for malloc and free

#if 0
#   include <new>
#   define __THROW_BAD_ALLOC throw bad_alloc;
#else
#   include <iostream>
    using namespace std;
#   define __THROW_BAD_ALLOC {cerr<<"out of memory"<<endl; exit(1);};
#endif

/* first level allocator */
/* __malloc_alloc_template */
template<int inst>
class __malloc_alloc_template
{
private:
    //out of memory
    //keep trying alloc until succeed
    static void* oom_alloc(size_t bytes)
    {
        set_malloc_handler(0);
        void * res = 0;
        while(1)
        {
            if(!__malloc_alloc_oom_handler) {__THROW_BAD_ALLOC;}
            (*__malloc_alloc_oom_handler)();
            if(res = malloc(bytes)) return res;
        }
    }
    static void* oom_realloc(void* p, size_t bytes)
    {
        set_malloc_handler(0);
        void * res = 0;
        while(1)
        {
            if(!__malloc_alloc_oom_handler) __THROW_BAD_ALLOC;
            (*__malloc_alloc_oom_handler)();
            if(res = realloc(p, bytes)) return res;
        }
    }
    static void (* __malloc_alloc_oom_handler)();

    //set the malloc handler function as f
    static void (* set_malloc_handler( void (*f)() )) ()
    {
        void (*old)() = __malloc_alloc_oom_handler;
        __malloc_alloc_oom_handler = f;
        return old;
    }
public:
    static void* allocate(size_t bytes)
    {
        void * res = malloc(bytes);
        if(!res) res = oom_alloc(bytes);
        return res;
    }
    static void deallocate(void* p, size_t bytes)
    {
        /*
        char * q = p;
        while(bytes--)
        {
            free(q);
            q++;
        }
        */
        free(p);
    }
    static void* reallocate(void* p, size_t bytes)
    {
        void * res = realloc(p, bytes);
        if(!res) res = oom_realloc(p, bytes);
        return res;
    }
};
template<int inst>
void (* __malloc_alloc_template<inst>::__malloc_alloc_oom_handler)() = 0;

typedef __malloc_alloc_template<0> malloc_alloc;

/* end of __malloc_alloc_template */

//memory pool
/* __default_alloc_template */

enum {MIN_BLOCK = 8, MAX_BLOCK = 128, NFREELISTS = MAX_BLOCK/MIN_BLOCK};

union obj
{
    obj* free_list_link;
    char client_data[1];
};

template <bool threads, int inst>
class __default_alloc_template
{

    static obj* freelist[NFREELISTS];
    //start of memory pool
    static char* start_free;
    //end of memory pool
    static char* end_free;
    //size of the whole memory pool
    static size_t heap_size;

private:
    //n -> 8*m
    //eg. 7->8, 8->8, 9->16
    static size_t ROUND_UP(size_t n)
    {
    //n -> freelist_index
    //eg. 7->0, 8->0, 9->1, 16->1
        return (n + MIN_BLOCK - 1) & ~(MIN_BLOCK-1);
    }
    static size_t FREELIST_INDEX(size_t n)
    {
        return (n/MIN_BLOCK - 1);
    }
    static char* chunk_alloc(size_t size, int & nobjs)
    {
        char* res = 0;
        size_t bytes_needed = nobjs * size;
        size_t bytes_left = end_free - start_free;
        //if the left space meets the demand, allocate and return
        if(bytes_needed <= bytes_left)
        {
            res = start_free;
            start_free += bytes_needed;
            return res;
        }
        //if the left space partially meets the demand, allocate
        if(size <= bytes_left)
        {
            res = start_free;
            nobjs = bytes_left/size;
            start_free += nobjs * size; //note that nobjs*size != bytes_left
            return res;
        }
        //if the left space cannot meets even one block demand,
        size_t bytes_to_get = 2 * bytes_needed + ROUND_UP(heap_size >> 4); //bytes which are to be gained.
        //put the left part into the correct freelist
        if(bytes_left > 0)
        {
            obj** my_freelist = freelist + FREELIST_INDEX(bytes_left);
            ((obj*)(start_free))->free_list_link = *my_freelist;
            *my_freelist = (obj*)(start_free);
        }
        //and try to gain more space from the heap.
        start_free = static_cast<char*>(malloc(bytes_to_get));
        if(0 == start_free)
        {
            int i;
            obj ** my_freelist;
            //search unused blocks in the freelist whose size le "size", and reuse it
            //otherwise the space in the freelist might grow huge.
            for(i=size; i<MAX_BLOCK; i+=MIN_BLOCK)
            {
                my_freelist = freelist + FREELIST_INDEX(i);
                obj* p = *my_freelist;
                if(p)
                {
                    *my_freelist = p->free_list_link;
                    start_free = (char*)(p);
                    end_free = start_free + i;
                    //now end_free - start_free == i >= size
                    return (chunk_alloc(size, nobjs));
                }
            }
            //end_free = 0; /?
            //turn to malloc_alloc::oom_alloc if failed.
            start_free = static_cast<char*>(malloc_alloc::allocate(bytes_to_get));
        }
        end_free = start_free + bytes_to_get;
        heap_size += bytes_to_get;
        return chunk_alloc(size, nobjs);
    }

    //n should be 8*m
    static void* refill(size_t n)
    {
        //the default refill number of the block of size n
        int nobjs = 20;
        char* chunk = chunk_alloc(n, nobjs);
        obj** my_freelist = freelist + FREELIST_INDEX(n);
        obj* res;
        obj* cur;
        int i;

        //if chunk_alloc only return 1 block;
        if(1 == nobjs) return (chunk);
        //else if many blocks are chunked
        res = (obj*)(chunk);
        *my_freelist = cur = (obj*)((char*)(res)+n);
        for(i=2; i<nobjs; i++)
        {
            cur->free_list_link = (obj*)((char*)(cur)+n);
            cur = cur->free_list_link;
        }
        cur->free_list_link = 0;
        return res;
    }

public:
    static void * allocate(size_t bytes)
    {
        if(MAX_BLOCK < bytes) return malloc_alloc::allocate(bytes);
        //gain a space from freelist
        obj** my_freelist = freelist + FREELIST_INDEX(bytes);
        obj* res = *my_freelist;
        if(0 == res)
        {
            return refill(ROUND_UP(bytes));
        }
        *my_freelist = res->free_list_link;
        return res;
    }
    static void deallocate(void* p, size_t bytes)
    {
        if(MAX_BLOCK < bytes){ malloc_alloc::deallocate(p, bytes);return; }
        //return the deallocated space to freelist;
        obj** my_freelist = freelist + FREELIST_INDEX(bytes);
        static_cast<obj*>(p)->free_list_link = *my_freelist;
        *my_freelist = static_cast<obj*>(p);
    }
};

template <bool threads, int inst>
char* __default_alloc_template<threads, inst>::start_free = 0;
template <bool threads, int inst>
char* __default_alloc_template<threads, inst>::end_free = 0;
template <bool threads, int inst>
size_t __default_alloc_template<threads, inst>::heap_size = 0;
template <bool threads, int inst>
obj * __default_alloc_template<threads, inst>::freelist[NFREELISTS]={0};

/* end of __default_alloc_template */

#define __NODE_ALLOCATOR_THREADS 0

#ifdef __USE_MALLOC_ALLOC
typedef malloc_alloc alloc;
#else
typedef __default_alloc_template<__NODE_ALLOCATOR_THREADS, 0> alloc;
#endif // __USE_MALLOC

/* simple_alloc */
template<class T, class Alloc=alloc>
class simple_alloc
{
public:
    typedef T           value_type;
    typedef T*          pointer;
    typedef const T*    const_pointer;
    typedef T&          reference;
    typedef const T&    const_reference;
    typedef size_t      size_type;
    typedef ptrdiff_t   difference_type;
public:
    static T* allocate(size_t n)
    {
        if(n==0) return 0;
        return static_cast<T*>(Alloc::allocate(n*sizeof(T)));
    }
    static const T* allocate(void)
    {
        return Alloc::allocate(sizeof(T));
    }
    static void deallocate(T* p, size_t n)
    {
        if(0!=n) Alloc::deallocate(p, n*sizeof(T));
    }
    static void deallocate(T* p)
    {
        Alloc::deallocate(p, sizeof(T));
    }

    inline void construct(T* p, const T& value)
    {
        new (p) T(value);
    }
    inline void destroy(T* p)
    {
        return destroy(p, has_trivial_destructor<T>::value);
    }
    inline void destroy(T* p, const bool htd)
    {
        if(htd)
            cerr << "I am a trivial destructor"<<endl;
        else
            p->~T();
    }
};
/* end of simple_alloc */

/* end of alloc.h */

/* main.cpp */

#include <vector>
#include <auto_ptr.h>

#define __USE_MALLOC_ALLOC 1

#pragma pack(4)
struct s
{
    int i;
    int j;
    char c;
    s(int _i, int _j=0):i(_i), j(_j){}
    ~s()
    {
        cerr <<"I am not a trivial destructor."<<endl;
    }
};
#pragma pack()

ostream & operator << (ostream & o, const s& _s)
{
    o << "(" <<_s.i<< ", " <<_s.j<<")";
    return o;
}

main()
{
    typedef s T;
    cout <<"sizeof(T) "<<sizeof(T)<<endl;
    auto_ptr<vector<T, simple_alloc<int>>> pv(new vector<T, simple_alloc<int>> ());
    vector<T, simple_alloc<int>> &v = (*pv);
    //vector<T, simple_alloc<int>> v;
    for(int i=0; i<10; i++)
        v.push_back(T(i));
    vector<T, simple_alloc<int>>::iterator it;
    for(it=v.begin(); it!=v.end(); it++)
        cout <<(*it)<<endl;
    cout <<endl;
}
/* end of main.cpp */

看过jjhou翻译的Effective C++系列之后，深深地变成了他的脑残粉。本来本着闲着没事追追星的心态去看STL代码剖析，不到1/5的内容里，我已经被STL大师深深折服。不愧是大师，代码写得面面俱到，宏定义用得如火纯青，光是allocator就秒杀只会写new和delete的5星渣渣。建议和我一样的C++新手们也去看看，看完就不是新手啦！