regionallocator.d

module std.regionallocator;

import std.algorithm, std.conv, std.exception, std.traits, std.typecons;
import core.bitop, core.exception, core.memory;

interface Allocator {
    enum alignment = 1;
    enum isAutomatic = false;
    enum isScoped = false;
    enum freeIsChecked = false;
    // further traits

    void* alloc(size_t nbytes, Flag!"GCScan" scan);
    void free(void* p);
}

@trusted T[] newArray(T : T[], A)(A allocator, size_t sz)
{
    auto res = newUnitializedArray!(T[], A)(allocator, sz);
    res[] = T.init;
    return res;
}

// not allowed for manually freeing as there is no way to retrieve
// pointer from sliced array
@trusted T[] newUnitializedArray(T : T[], A)(A allocator, size_t sz)
if (A.isAutomatic || A.isScoped)
{
    return (cast(T*)allocator.alloc(sz * T.sizeof))[0 .. sz];
}

@trusted T* newT(T, A, Args...)(A allocator, Args args)
if (!is(T == class) && (A.isAutomatic || A.isScoped || A.freeIsChecked))
{
    auto chunk = newUnitializedArray!(void[])(allocator, T.sizeof);
    return emplace!T(chunk, args);
}

@trusted T newT(T, A, Args...)(A allocator, Args args)
if (is(T == class) && (A.isAutomatic || A.isScoped || A.freeIsChecked))
{
    enum classSize = __traits(classInstanceSize, T);
    auto chunk = newUnitializedArray!(void[])(allocator, classSize);
    return emplace!T(chunk, args);
}

final class RegionAllocator(SegmentAllocator = GCAllocator) : Allocator
{
    enum alignment = SegmentAllocator.alignment;
    enum isAutomatic = false;
    enum isScoped = true;
    enum freeIsChecked = false;

    enum defaultSegmentSize = 4 * 1024;

    // Allocators should be default constructible
    this(size_t segmentSize=defaultSegmentSize, SegmentAllocator segmentAllocator=null)
    {
        _segmentSize = nextPow2(segmentSize);
        if (segmentAllocator is null)
            segmentAllocator = new SegmentAllocator();
        _segmentAllocator = segmentAllocator;
        pushMark();
    }

    ~this()
    {
        popMark(0);
    }

    void* alloc(size_t nbytes, Flag!"GCScan" scan = Flag!"GCScan".no)
    {
        if (nbytes == 0)
            return null;

        nbytes = (nbytes + alignment - 1) & ~(alignment - 1);

        void* result;

        if (nbytes > _segmentSize)
        {
            _bigObjects.length = _marks.length;
            result = rawAlloc(nbytes);
            _bigObjects[_marks.length - 1] ~= result;
        }
        else
        {
            const top = _segments.length * _segmentSize;
            if (_fill + nbytes > top)
            {
                result = rawAlloc(_segmentSize);
                _segments ~= result;
                _fill = top + nbytes;
            }
            else
            {
                result = _segments[$-1] + (_fill & (_segmentSize - 1));
                _fill += nbytes;
            }
        }

        return enforceEx!OutOfMemoryError(result);
    }

    void free(void* p)
    {
        enforce(0, "Can't free individual pointers using RegionAllocator." ~
                "Instead use pushMark/popMark or autoFree.");
    }

    void clear()
    {
        popMark(0);
        pushMark();
    }

    size_t pushMark()
    {
        _marks ~= Mark(_fill);
        return _marks.length - 1;
    }

    void popMark(size_t mark)
    {
        enforce(mark < _marks.length);
        enforce(_marks[mark]._refcount == 0, "Stall autoFree references to stack.");

        assert(isSorted!q{a._fill < b._fill}(_marks));
        assert(!canFind!q{a._fill > b}(_marks, _fill));
        // to expensive for enforcement ?
        assert(!canFind!q{a._refcount != 0}(_marks[mark .. $]), "Stall autoFree references to stack.");

        _fill = _marks[mark]._fill;
        _marks = _marks[0 .. mark];

        const keep = (_fill + _segmentSize - 1) >> bsr(_segmentSize);
        foreach(seg; _segments[keep .. $])
            rawFree(seg);
        _segments = _segments[0 .. keep];

        if (mark < _bigObjects.length)
        {
            foreach(objs; _bigObjects[mark .. $])
                foreach(obj; objs)
                    rawFree(obj);
            _bigObjects = _bigObjects[0 .. mark];
        }

        assert(_marks.length == mark);
        assert(_bigObjects.length <= mark);
    }

    /* Will call pushMark and return a ref counted struct using RAII to call popMark.
     */
    auto autoFree()
    {
        static struct AutoFree
        {
            this(RegionAllocator rallocator)
            {
                _rallocator = rallocator;
                if (_rallocator !is null)
                {
                    _markidx = rallocator.pushMark();
                    incref();
                }
            }

            this(this)
            {
                if (_rallocator !is null)
                {
                    incref();
                }
            }

            ~this()
            {
                if (_rallocator !is null)
                {
                    if (decref())
                    {
                        _rallocator.popMark(_markidx);
                        // @@ BUG 6630 @@ this will silently delete the alias this class in RegionAllocator
                        version (none)
                            _rallocator = null;
                        else
                            *(cast(void**)&this._rallocator) = null;
                    }
                }
            }

            @property size_t mark() const { return _markidx; }

        private:
            void incref()
            {
                ++_rallocator._marks[_markidx]._refcount;
            }

            bool decref()
            {
                assert(_rallocator._marks[_markidx]._refcount);
                return --_rallocator._marks[_markidx]._refcount == 0;
            }

            RegionAllocator _rallocator;
            size_t _markidx;
        }

        return AutoFree(this);
    }

 private:

    void* rawAlloc(size_t nbytes)
    {
        return _segmentAllocator.alloc(nbytes, Flag!"GCScan".no);
    }

    void rawFree(void* p)
    {
        _segmentAllocator.free(p);
    }

    static struct Mark {
        size_t _fill;
        size_t _refcount;
    }

    Mark[] _marks;
    void*[][] _bigObjects;
    void*[] _segments;
    size_t _fill;
    immutable size_t _segmentSize;
    SegmentAllocator _segmentAllocator;
    // should alias to nested allocators to allow access to specific functions
    alias _segmentAllocator this;
}

final class GCAllocator : Allocator
{
    enum alignment = 8;
    enum isAutomatic = true;
    enum isScoped = false;
    enum freeIsChecked = true; // ?

    void* alloc(size_t nbytes, Flag!"GCScan" scan)
    {
        return GC.malloc(nbytes, scan ? 0 : GC.BlkAttr.NO_SCAN);
    }

    void free(void* p)
    {
        GC.free(p);
    }
}

final class StdCAllocator : Allocator
{
    version (OSX)
      enum alignment = 16;
    else
      enum alignment = 8;
    enum isAutomatic = false;
    enum isScoped = false;
    enum freeIsChecked = false;

    import core.stdc.stdlib;

    void* alloc(size_t nbytes, Flag!"GCScan" scan)
    {
        auto p = core.stdc.stdlib.malloc(nbytes);
        if (scan)
            GC.addRange(p, nbytes);
        return p;
    }

    void free(void* p)
    {
        // let GC decide if we've added it earlier
        GC.removeRange(p);
        core.stdc.stdlib.free(p);
    }
}

final class SimpleFreeListAllocator(Wrapped = StdCAllocator) : Allocator
{
    enum alignment = min(size_t.sizeof, Wrapped.alignment);
    enum isAutomatic = false;
    enum isScoped = false;
    enum freeIsChecked = false;

    this(Wrapped wrapped=null)
    {
        if (wrapped is null)
            wrapped = new Wrapped;
        _wrapped = wrapped;
    }

    static if (!(Wrapped.isAutomatic || Wrapped.isScoped))
    {
        ~this()
        {
            clear();
        }
    }

    void* alloc(size_t nbytes, Flag!"GCScan" scan)
    {
        if (_freeLists.length)
            enforce(_scan == scan,
                    "Mixing scan an no-scan allocations for non-empty FreeListAllocator is unimplemented.");
        else
            _scan = scan;

        nbytes = nextPow2(nbytes);
        const idx = bsr(nbytes);
        void* res;
        if (idx < _freeLists.length && _freeLists[idx].length)
        {
            res = _freeLists[idx][$-1];
            _freeLists[idx] = _freeLists[idx][0 .. $-1];
        }
        else
        {
            // Add size_t after rounding up to power of 2 so that allocating pow2 has optimal fit.
            res = _wrapped.alloc(nbytes + size_t.sizeof, scan);
            *(cast(size_t*)res) = nbytes;
            res += size_t.sizeof;
        }
        return res;
    }

    void free(void* p)
    {
        const size = *(cast(size_t*)p - 1);
        enforce((size & (size - 1)) == 0, "Deallocating foreign pointer or memory corruption.");
        const idx = bsr(size);
        _freeLists.length = max(_freeLists.length, idx + 1);
        _freeLists[idx] ~= p;
    }

    void clear()
    {
        foreach(blocks; _freeLists)
            foreach(block; blocks)
                _wrapped.free(block - size_t.sizeof);
    }

    void*[][] _freeLists;
    Wrapped _wrapped;
    Flag!"GCScan" _scan;
    alias _wrapped this;
}

private size_t nextPow2(size_t n)
{
    enforce(n);
    auto msbPow2 = 1 << bsr(n);
    if (n & (msbPow2 - 1))
        msbPow2 <<= 1;
    return msbPow2;
}

unittest {
    assert(nextPow2(1) == 1);
    assert(nextPow2(2) == 2);
    assert(nextPow2(3) == 4);
    assert(nextPow2(4) == 4);
    assert(nextPow2(63) == 64);
    assert(nextPow2(64) == 64);
    assert(nextPow2(255) == 256);
    assert(nextPow2(256) == 256);
    assert(nextPow2(257) == 512);
    assert(nextPow2(4 * 1024 * 1024 - 1) == 4 * 1024 * 1024);
}

unittest {
    scope auto ralloc = new RegionAllocator!(StdCAllocator)(256);
    auto outer = ralloc.autoFree();
    foreach(i; 0 .. 300)
        ralloc.alloc(i);

    auto inner = ralloc.autoFree();
    foreach(i; 0 .. 20)
    {
        // Support an idiom that is hardly supported by RAII other
        // than allowing uninitialized Guards.
        size_t mark;
        foreach(j; 0 .. 20)
        {
            import std.random;
            if (std.random.dice(0.5, 0.5))
            {
                if (!mark) mark = ralloc.pushMark();
                ralloc.alloc(i * j);
            }
        }
        if (mark)
        {
            // using push marks creates a simple logical relation
            // one can even use free(max(mark1, mark2))
            assert(mark > inner.mark);
            ralloc.popMark(mark);
        }
    }
}

unittest {
    static struct S {
        @property int a() { return _a; }
        int _a = 12;
    }

    static class C {
        @property S s() { return _s; }
        S _s;
        alias s this;
    }

    scope auto ralloc = new RegionAllocator!(SimpleFreeListAllocator!());

    size_t cnt = 2;
    while (cnt--)
    {
        auto ary = newArray!(int[])(ralloc, 4);
        static assert(is(typeof(ary) == int[]));
        assert(ary == [0, 0, 0, 0]);

        auto sary = newT!(int[4])(ralloc);
        static assert(is(typeof(sary) == int[4]*));
        assert(*sary == [0, 0, 0, 0]);

        auto s = newT!S(ralloc);
        assert(s.a == 12);

        auto c = newT!C(ralloc);
        assert(c.a == 12);

        ralloc.clear();
    }
}

void main() {
}