root/trunk/alloc.d

/**Stuff having to do with memory management.  Mostly TempAlloc and some data
 * structure implementations that go with it.
 *
 * Author:  David Simcha*/
 /*
 * License:
 * Boost Software License - Version 1.0 - August 17th, 2003
 *
 * Permission is hereby granted, free of charge, to any person or organization
 * obtaining a copy of the software and accompanying documentation covered by
 * this license (the "Software") to use, reproduce, display, distribute,
 * execute, and transmit the Software, and to prepare derivative works of the
 * Software, and to permit third-parties to whom the Software is furnished to
 * do so, all subject to the following:
 *
 * The copyright notices in the Software and this entire statement, including
 * the above license grant, this restriction and the following disclaimer,
 * must be included in all copies of the Software, in whole or in part, and
 * all derivative works of the Software, unless such copies or derivative
 * works are solely in the form of machine-executable object code generated by
 * a source language processor.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE, TITLE AND NON-INFRINGEMENT. IN NO EVENT
 * SHALL THE COPYRIGHT HOLDERS OR ANYONE DISTRIBUTING THE SOFTWARE BE LIABLE
 * FOR ANY DAMAGES OR OTHER LIABILITY, WHETHER IN CONTRACT, TORT OR OTHERWISE,
 * ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 * DEALINGS IN THE SOFTWARE.
 */

module dstats.alloc;

import std.traits, core.memory, std.array, std.range, core.exception,
    std.functional, std.math, std.algorithm : max;

import core.stdc.stdio;

static import core.stdc.stdlib;

import dstats.base;

version(unittest) {
    import std.stdio, std.conv, std.random, dstats.sort;
    void main() {}
}

// This is just for convenience/code readability/saving typing.
enum ptrSize = (void*).sizeof;

// This was accidentally assumed in a few places and I'm too lazy to fix it
// until I see proof that it needs to be fixed.
static assert(bool.sizeof == 1);

template IsType(T, Types...) {
    // Original idea by Burton Radons, modified
    static if (Types.length == 0)
        const bool IsType = false;
    else
        const bool IsType = is(T == Types[0]) || IsType!(T, Types[1 .. $]);
}

template ArrayType1(T: T[]) {
    alias T ArrayType1;
}

template isReferenceType(Types...) {  //Thanks to Bearophile.
    static if (Types.length == 0) {
        const bool isReferenceType = false;
    } else static if (Types.length == 1) {
        static if (IsType!(Unqual!(Types[0]), bool, byte, ubyte, short, ushort,
                           int, uint, long, ulong, float, double, real, ifloat,
                           idouble, ireal, cfloat, cdouble, creal, char, dchar,
                           wchar) ) {
            const bool isReferenceType = false;
        } else static if ( is(Types[0] == struct) ) {
            const bool isReferenceType =
            isReferenceType!(FieldTypeTuple!(Types[0]));
        } else static if (isStaticArray!(Types[0])) {
            const bool isReferenceType = isReferenceType!(ArrayType1!(Types[0]));
        } else
            const bool isReferenceType = true;
    } else
        const bool isReferenceType = isReferenceType!(Types[0]) |
        isReferenceType!(Types[1 .. $]);
} // end isReferenceType!()

unittest {
    static assert(!isReferenceType!(typeof("Foo"[0])));
    static assert(isReferenceType!(uint*));
    static assert(!isReferenceType!(int[3]));
    struct noPtrs {
        uint f;
        uint b;
    }
    struct ptrs {
        uint* f;
        uint b;
    }
    static assert(!isReferenceType!(noPtrs));
    static assert(isReferenceType!(ptrs));
}

template blockAttribute(T) {
    static if (isReferenceType!(T))
        enum blockAttribute = 0;
    else enum blockAttribute = GC.BlkAttr.NO_SCAN;
}

///Returns a new array of type T w/o initializing elements.
T[] newVoid(T)(size_t length) {
    T* ptr = cast(T*) GC.malloc(length * T.sizeof, blockAttribute!(T));
    return ptr[0..length];
}

void lengthVoid(T)(ref T[] input, int newLength) {
    input.lengthVoid(cast(size_t) newLength);
}

///Lengthens an array w/o initializing new elements.
void lengthVoid(T)(ref T[] input, size_t newLength) {
    if (newLength <= input.length ||
            GC.sizeOf(input.ptr) >= newLength * T.sizeof) {
        input = input.ptr[0..newLength];  //Don't realloc if I don't have to.
    } else {
        T* newPtr = cast(T*) GC.realloc(input.ptr,
                                        T.sizeof * newLength, blockAttribute!(T));
        input = newPtr[0..newLength];
    }
}

private template Appends(T, U) {
    enum bool Appends = AppendsImpl!(T, U).ret;
}

private template AppendsImpl(T, U) {
    T[] a;
    U b;
    enum bool ret = is(typeof(a ~= b));
}

///Appends to an array, deleting the old array if it has to be realloced.
void appendDelOld(T, U)(ref T[] to, U from)
if(Appends!(T, U)) {
    auto old = to;
    to ~= from;
    if (old.ptr !is to.ptr && old.ptr !is null) delete old;
}

unittest {
    uint[] foo;
    foo.appendDelOld(5);
    foo.appendDelOld(4);
    foo.appendDelOld(3);
    foo.appendDelOld(2);
    foo.appendDelOld(1);
    assert(foo == cast(uint[]) [5,4,3,2,1]);
}

// Memory allocation routines.  These wrap malloc(), free() and realloc(),
// and guarantee alignment.
private enum size_t alignBytes = 16;


static void outOfMemory()  {
    throw new OutOfMemoryError("Out of memory in TempAlloc.");
}

private void* alignedMalloc(size_t size, bool shouldAddRange = false) {
    // We need (alignBytes - 1) extra bytes to guarantee alignment, 1 byte
    // to store the shouldAddRange flag, and ptrSize bytes to store
    // the pointer to the beginning of the block.
    void* toFree = core.stdc.stdlib.malloc(
        alignBytes + ptrSize + size
    );

    if(toFree is null) outOfMemory();

    // Add the offset for the flag and the base pointer.
    auto intPtr = cast(size_t) toFree + ptrSize + 1;

    // Align it.
    intPtr = (intPtr + alignBytes - 1) & (~(alignBytes - 1));
    auto ret = cast(void**) intPtr;

    // Store base pointer.
    (cast(void**) ret)[-1] = toFree;

    // Store flag.
    (cast(bool*) ret)[-1 - ptrSize] = shouldAddRange;

    if(shouldAddRange) {
        GC.addRange(ret, size);
    }

    return ret;
}

private void alignedFree(void* ptr) {
    // If it was allocated with alignedMalloc() then the pointer to the
    // beginning is at ptr[-1].
    auto addedRange = (cast(bool*) ptr)[-1 - ptrSize];

    if(addedRange) {
        GC.removeRange(ptr);
    }

    core.stdc.stdlib.free( (cast(void**) ptr)[-1]);
}

private void* alignedRealloc(void* ptr, size_t newLen, size_t oldLen) {
    auto storedRange = (cast(bool*) ptr)[-1 - ptrSize];
    auto newPtr = alignedMalloc(newLen, storedRange);
    memcpy(newPtr, ptr, oldLen);

    alignedFree(ptr);
    return newPtr;
}

/**A struct to allocate memory in a strictly first-in last-out order for
 * things like scratch space.  Technically, memory can safely escape the
 * scope in which it was allocated.  However, this is a very bad idea
 * unless being done within the private API of a class, struct or nested
 * function, where it can be guaranteed that LIFO will not be violated.
 *
 * Under the hood, this works by allocating large blocks (currently 4 MB)
 * from the GC, and sub-allocating these as a stack.  Very large allocations
 * (currently > 4MB) are simply performed on the heap.  There are two ways to
 * free memory:  Calling TempAlloc.free() frees the last allocated block.
 * Calling TempAlloc.frameFree() frees all memory allocated since the last
 * call to TempAlloc.frameInit().
 *
 * All allocations are aligned on 16-byte boundaries using padding, since on x86,
 * 16-byte alignment is necessary to make SSE2 work.  Note, however, that this
 * is implemented based on the assumption that the GC allocates using 16-byte
 * alignment (which appears to be true in druntime.)
 */
struct TempAlloc {
private:
    static struct Stack(T) {  // Simple, fast stack w/o error checking.
        private size_t capacity;
        private size_t index;
        private T* data;
        private enum sz = T.sizeof;

        private static size_t max(size_t lhs, size_t rhs) pure {
            return (rhs > lhs) ? rhs : lhs;
        }

        void push(T elem) {
            if (capacity == index) {
                capacity = max(16, capacity * 2);
                data = cast(T*) core.stdc.stdlib.realloc(data, capacity * sz);
            }
            data[index++] = elem;
        }

        T pop() {
            index--;
            auto ret = data[index];
            data[index] = T.init;  // Prevent false ptrs.
            return ret;
        }

        void destroy() {
            if(data) {
                core.stdc.stdlib.free(data);
                data = null;
            }
        }
    }

    struct Block {
        size_t used = 0;
        void* space = null;
    }

    final class State {
        size_t used;
        void* space;
        size_t totalAllocs;
        void*[] lastAlloc;
        uint nblocks;
        uint nfree;
        size_t frameIndex;

        // inUse holds info for all blocks except the one currently being
        // allocated from.  freelist holds space ptrs for all free blocks.
        Stack!(Block) inUse;
        Stack!(void*) freelist;

        void putLast(void* last) {
            // Add an element to lastAlloc, checking length first.
            if (totalAllocs == lastAlloc.length)
                doubleSize(lastAlloc);
            lastAlloc[totalAllocs] = cast(void*) last;
            totalAllocs++;
        }

        void destroy() {
            if(space) {
                alignedFree(space);
                space = null;
            }

            if(lastAlloc) {
                core.stdc.stdlib.free(lastAlloc.ptr);
                lastAlloc = null;
            }

            while(inUse.index > 0) {
                auto toFree = inUse.pop();
                alignedFree(toFree.space);
            }

            while(freelist.index > 0) {
                auto toFree = freelist.pop();
                alignedFree(toFree);
            }

            inUse.destroy();
            freelist.destroy();
        }

        ~this() {
            destroy();
        }
    }

    static ~this() {
        if(state) {
            state.destroy();
            state = null;
        }
    }

    enum size_t alignBytes = 16U;
    enum size_t blockSize = 4U * 1024U * 1024U;
    enum size_t nBookKeep = blockSize / alignBytes * ptrSize;
    static State state;

    static void doubleSize(ref void*[] lastAlloc) {
        size_t newSize = lastAlloc.length * 2;
        void** ptr = cast(void**) core.stdc.stdlib.realloc(
            lastAlloc.ptr, newSize * ptrSize);
        lastAlloc = ptr[0..newSize];
    }

    static State stateInit() {
        State stateCopy;
        try { stateCopy = new State; } catch { outOfMemory(); }

        with(stateCopy) {
            space = alignedMalloc(blockSize);

            // We don't need 16-byte alignment for the bookkeeping array.
            lastAlloc = (cast(void**) core.stdc.stdlib.malloc(nBookKeep))
                        [0..nBookKeep / ptrSize];
            nblocks++;
        }

        state = stateCopy;
        return stateCopy;
    }

    static size_t getAligned(size_t nbytes) pure {
        // Only works if alignBytes is a power of two, but I think that's
        // a pretty safe assumption.
        return (nbytes + (alignBytes - 1)) & (~(alignBytes - 1));
    }

public:
    /**Allows caller to cache the state class on the stack and pass it in as a
     * parameter.  This is ugly, but results in a speed boost that can be
     * significant in some cases because it avoids a thread-local storage
     * lookup.  Also used internally.*/
    static State getState() {
        State stateCopy = state;
        return (stateCopy is null) ? stateInit : stateCopy;
    }

    /**Initializes a frame, i.e. marks the current allocation position.
     * Memory past the position at which this was last called will be
     * freed when frameFree() is called.  Returns a reference to the
     * State class in case the caller wants to cache it for speed.*/
    static State frameInit() {
        return frameInit(getState);
    }

    /**Same as frameInit() but uses stateCopy cached on stack by caller
     * to avoid a thread-local storage lookup.  Strictly a speed hack.*/
    static State frameInit(State stateCopy) {
        with(stateCopy) {
            putLast( cast(void*) frameIndex );
            frameIndex = totalAllocs;
        }
        return stateCopy;
    }

    /**Frees all memory allocated by TempAlloc since the last call to
     * frameInit().*/
    static void frameFree() {
        frameFree(getState);
    }

    /**Same as frameFree() but uses stateCopy cached on stack by caller
    * to avoid a thread-local storage lookup.  Strictly a speed hack.*/
    static void frameFree(State stateCopy) {
        with(stateCopy) {
            while (totalAllocs > frameIndex) {
                free(stateCopy);
            }
            frameIndex = cast(size_t) lastAlloc[--totalAllocs];
        }
    }

    /**Purely a convenience overload, forwards arguments to TempAlloc.malloc().*/
    static void* opCall(T...)(T args) {
        return TempAlloc.malloc(args);
    }

    /**Allocates nbytes bytes on the TempAlloc stack.  NOT safe for real-time
     * programming, since if there's not enough space on the current block,
     * a new one will automatically be created.  Also, very large objects
     * (currently over 4MB) will simply be heap-allocated.
     *
     * Bugs:  Memory allocated by TempAlloc is not scanned by the GC.
     * This is necessary for performance and to avoid false pointer issues.
     * Do not store the only reference to a GC-allocated object in
     * TempAlloc-allocated memory.*/
    static void* malloc(size_t nbytes) {
        return malloc(nbytes, getState);
    }

    /**Same as malloc() but uses stateCopy cached on stack by caller
    * to avoid a thread-local storage lookup.  Strictly a speed hack.*/
    static void* malloc(size_t nbytes, State stateCopy) {
        nbytes = getAligned(nbytes);
        with(stateCopy) {
            void* ret;
            if (blockSize - used >= nbytes) {
                ret = space + used;
                used += nbytes;
            } else if (nbytes > blockSize) {
                ret = alignedMalloc(nbytes);
            } else if (nfree > 0) {
                inUse.push(Block(used, space));
                space = freelist.pop;
                used = nbytes;
                nfree--;
                nblocks++;
                ret = space;
            } else { // Allocate more space.
                inUse.push(Block(used, space));
                space = alignedMalloc(blockSize);
                nblocks++;
                used = nbytes;
                ret = space;
            }
            putLast(ret);
            return ret;
        }
    }

    /**Frees the last piece of memory allocated by TempAlloc.  Since
     * all memory must be allocated and freed in strict LIFO order,
     * there's no need to pass a pointer in.  All bookkeeping for figuring
     * out what to free is done internally.*/
    static void free() {
        free(getState);
    }

    /**Same as free() but uses stateCopy cached on stack by caller
    * to avoid a thread-local storage lookup.  Strictly a speed hack.*/
    static void free(State stateCopy) {
        with(stateCopy) {
            void* lastPos = lastAlloc[--totalAllocs];

            // Handle large blocks.
            if (lastPos > space + blockSize || lastPos < space) {
                alignedFree(lastPos);
                return;
            }

            used = (cast(size_t) lastPos) - (cast(size_t) space);
            if (nblocks > 1 && used == 0) {
                freelist.push(space);
                Block newHead = inUse.pop;
                space = newHead.space;
                used = newHead.used;
                nblocks--;
                nfree++;

                if (nfree >= nblocks * 2) {
                    foreach(i; 0..nfree / 2) {
                        alignedFree(freelist.pop);
                        nfree--;
                    }
                }
            }
        }
    }

    /**Returns how many bytes are available in the current frame.*/
    static size_t slack() @property {
        return blockSize - getState().used;
    }

}

/**Allocates an array of type T and size size using TempAlloc.
 * Note that appending to this array using the ~= operator,
 * or enlarging it using the .length property, will result in
 * undefined behavior.  This is because, if the array is located
 * at the beginning of a TempAlloc block, the GC will think the
 * capacity is as large as a TempAlloc block, and will overwrite
 * adjacent TempAlloc-allocated data, instead of reallocating it.
 *
 * Bugs: Do not store the only reference to a GC-allocated reference object
 * in an array allocated by newStack because this memory is not
 * scanned by the GC.*/
T[] newStack(T)(size_t size, TempAlloc.State state = null) {
    if(state is null) {
        state = TempAlloc.getState();
    }

    size_t bytes = size * T.sizeof;
    T* ptr = cast(T*) TempAlloc.malloc(bytes, state);
    return ptr[0..size];
}

///**Same as newStack(size_t) but uses stateCopy cached on stack by caller
//* to avoid a thread-local storage lookup.  Strictly a speed hack.*/
//T[] newStack(T)(size_t size, TempAlloc.State state) nothrow {
//    size_t bytes = size * T.sizeof;
//    T* ptr = cast(T*) TempAlloc.malloc(bytes, state);
//    return ptr[0..size];
//}

/**Concatenate any number of arrays of the same type, placing results on
 * the TempAlloc stack.*/
T[0] stackCat(T...)(T data) {
    foreach(array; data) {
        static assert(is(typeof(array) == typeof(data[0])));
    }

    size_t totalLen = 0;
    foreach(array; data) {
        totalLen += array.length;
    }
    auto ret = newStack!(Unqual!(typeof(T[0][0])))(totalLen);

    size_t offset = 0;
    foreach(array; data) {
        ret[offset..offset + array.length] = array[0..$];
        offset += array.length;
    }
    return cast(T[0]) ret;
}

void rangeCopy(T, U)(T to, U from) {
    static if(is(typeof(to[] = from[]))) {
        to[] = from[];
    } else static if(isRandomAccessRange!(T)) {
        size_t i = 0;
        foreach(elem; from) {
            to[i++] = elem;
        }
    }
}

/**Creates a duplicate of a range for temporary use within a function in the
 * best wsy that can be done safely.  If ElementType!(T) is a value type
 * or T is an array, the results can safely be placed in TempAlloc because
 * either it doesn't need to be scanned by the GC or there's guaranteed to be
 * another reference to the contents somewhere. Otherwise, the results
 * are placed on the GC heap.
 *
 * This function is much faster if T has a length, but works even if it doesn't.
 */
Unqual!(ElementType!(T))[] tempdup(T)(T data)
if(isInputRange!(T) && (isArray!(T) || !isReferenceType!(ElementType!(T)))) {
    alias ElementType!(T) E;
    alias Unqual!(E) U;
    static if(dstats.base.hasLength!(T)) {
        U[] ret = newStack!(U)(data.length);
        rangeCopy(ret, data);
        return ret;
    } else {
        auto state = TempAlloc.getState;
        auto startPtr = TempAlloc(0, state);
        size_t bytesCopied = 0;

        while(!data.empty) {  // Make sure range interface is being used.
            auto elem = data.front;
            if(state.used + U.sizeof <= TempAlloc.blockSize) {
                data.popFront;
                *(cast(U*) (startPtr + bytesCopied)) = elem;
                bytesCopied += U.sizeof;
                state.used += U.sizeof;
            } else {
                if(bytesCopied + U.sizeof >= TempAlloc.blockSize / 2) {
                    // Then just heap-allocate.
                    U[] result = (cast(U*) alignedMalloc(bytesCopied * 2))
                        [0..bytesCopied / U.sizeof * 2];

                    immutable elemsCopied = bytesCopied / U.sizeof;
                    result[0..elemsCopied] = (cast(U*) startPtr)[0..elemsCopied];
                    finishCopy(result, data, elemsCopied);
                    TempAlloc.free;
                    state.putLast(result.ptr);
                    return result;
                } else {
                    U[] oldData = (cast(U*) startPtr)[0..bytesCopied / U.sizeof];
                    state.used -= bytesCopied;
                    state.totalAllocs--;
                    U[] newArray = newStack!(U)(bytesCopied / U.sizeof + 1, state);
                    newArray[0..oldData.length] = oldData[];
                    startPtr = state.space;
                    newArray[$ - 1] = elem;
                    bytesCopied += U.sizeof;
                    data.popFront;
                }
            }
        }
        auto rem = bytesCopied % TempAlloc.alignBytes;
        if(rem != 0) {
            auto toAdd = 16 - rem;
            if(state.used + toAdd < TempAlloc.blockSize) {
                state.used += toAdd;
            } else {
                state.used = TempAlloc.blockSize;
            }
        }
        return (cast(U*) startPtr)[0..bytesCopied / U.sizeof];
    }
}

Unqual!(ElementType!(T))[] tempdup(T)(T data)
if(isInputRange!(T) && !(isArray!(T) || !isReferenceType!(ElementType!(T)))) {
    // Initial guess of how much space to allocate.  It's relatively large b/c
    // the object will be short lived, so speed is more important than space
    // efficiency.
    enum initialGuess = 128;

    alias Unqual!(ElementType!T) E;
    auto arr = (cast(E*) alignedMalloc(E.sizeof * initialGuess, true))
        [0..initialGuess];

    finishCopy(arr, data, 0);
    TempAlloc.getState.putLast(arr.ptr);
    return arr;
}

// Finishes copying a range to a C heap allocated array.  Assumes the first
// half of the input array is stuff already copied and the second half is
// free space.
private void finishCopy(T, U)(ref T[] result, U range, size_t alreadyCopied) {
    void doRealloc() {
        auto newPtr = cast(T*) alignedRealloc(
            result.ptr, result.length * T.sizeof * 2, result.length * T.sizeof
        );
        result = newPtr[0..result.length * 2];
    }

    auto index = alreadyCopied;
    foreach(elem; range) {
        if(index == result.length) doRealloc();
        result[index++] = elem;
    }

    result = result[0..index];
}

// See Bugzilla 2873.  This can be removed once that's fixed.
template hasLength(R) {
    enum bool hasLength = is(typeof(R.init.length) : ulong) ||
                      is(typeof(R.init.length()) : ulong);
}

// Now that Phobos does this well, this just forwards to Phobos.
Unqual!(IterType!(T))[] toArray(T)(T range) if(isIterable!(T)) {
    return std.array.array(range);
//    static if(isArray!(T)) {
//        // Allow fast copying by assuming that the input is an array.
//        return range.dup;
//    } else static if(hasLength!(T)) {
//        // Preallocate array, then copy.
//        auto ret = newVoid!(Unqual!(IterType!(T)))(range.length);
//        static if(is(typeof(ret[] = range[]))) {
//            ret[] = range[];
//        } else {
//            size_t pos = 0;
//            foreach(elem; range) {
//                ret[pos++] = elem;
//            }
//        }
//        return ret;
//    } else {
//        // Don't have length, have to use appending.
//        Unqual!(IterType!(T))[] ret;
//        auto app = appender(&ret);
//        foreach(elem; range) {
//            app.put(elem);
//        }
//        return ret;
//    }
}

unittest {
    // Create quick and dirty finite but lengthless range.
    static struct Count {
        uint num;
        uint upTo;
        @property size_t front() {
            return num;
        }
        void popFront() {
            num++;
        }
        @property bool empty() {
            return num >= upTo;
        }
    }

    TempAlloc(1024 * 1024 * 3);
    Count count;
    count.upTo = 1024 * 1025;
    auto asArray = tempdup(count);
    foreach(i, elem; asArray) {
        assert(i == elem, to!(string)(i) ~ "\t" ~ to!(string)(elem));
    }
    assert(asArray.length == 1024 * 1025);
    TempAlloc.free;
    TempAlloc.free;
    while(TempAlloc.getState.freelist.index > 0) {
        alignedFree(TempAlloc.getState.freelist.pop);
    }
}

/**A string to mixin at the beginning of a scope, purely for
 * convenience.  Initializes a TempAlloc frame using frameInit(),
 * and inserts a scope statement to delete this frame at the end
 * of the current scope.
 *
 * Slower than calling free() manually when only a few pieces
 * of memory will be allocated in the current scope, due to the
 * extra bookkeeping involved.  Can be faster, however, when
 * large amounts of allocations, such as arrays of arrays,
 * are allocated, due to caching of data stored in thread-local
 * storage.*/
immutable char[] newFrame =
    "TempAlloc.frameInit; scope(exit) TempAlloc.frameFree;";

unittest {
    /* Not a particularly good unittest in that it depends on knowing the
     * internals of TempAlloc, but it's the best I could come up w/.  This
     * is really more of a stress test/sanity check than a normal unittest.*/

    // Make sure state is completely reset.
    if(TempAlloc.state) TempAlloc.state.destroy();
    TempAlloc.state = null;

     // First test to make sure a large number of allocations does what it's
     // supposed to in terms of reallocing lastAlloc[], etc.
     enum nIter =  TempAlloc.blockSize * 5 / TempAlloc.alignBytes;
     foreach(i; 0..nIter) {
         TempAlloc(TempAlloc.alignBytes);
     }
     assert(TempAlloc.getState.nblocks == 5, to!string(TempAlloc.getState.nblocks));
     assert(TempAlloc.getState.nfree == 0);
     foreach(i; 0..nIter) {
        TempAlloc.free;
    }
    assert(TempAlloc.getState.nblocks == 1);
    assert(TempAlloc.getState.nfree == 2);

    // Make sure logic for freeing excess blocks works.  If it doesn't this
    // test will run out of memory.
    enum allocSize = TempAlloc.blockSize / 2;
    foreach(i; 0..50) {
        foreach(j; 0..50) {
            TempAlloc(allocSize);
        }
        foreach(j; 0..50) {
            TempAlloc.free;
        }
    }

    // Make sure data is stored properly.
    foreach(i; 0..10) {
        TempAlloc(allocSize);
    }
    foreach(i; 0..5) {
        TempAlloc.free;
    }
    void* space = TempAlloc.state.space;
    size_t used = TempAlloc.state.used;

    TempAlloc.frameInit;
    // This array of arrays should not be scanned by the GC because otherwise
    // bugs caused th not having the GC scan certain internal things in
    // TempAlloc that it should would not be exposed.
    uint[][] arrays = (cast(uint[]*) GC.malloc((uint[]).sizeof * 10,
                       GC.BlkAttr.NO_SCAN))[0..10];
    foreach(i; 0..10) {
        uint[] data = newStack!(uint)(250_000);
        foreach(j, ref e; data) {
            e = cast(uint) (j * (i + 1));  // Arbitrary values that can be read back later.
        }
        arrays[i] = data;
    }

    // Make stuff get overwrriten if blocks are getting GC'd when they're not
    // supposed to.
    GC.minimize;  // Free up all excess pools.
    uint[][] foo;
    foreach(i; 0..40) {
        foo ~= new uint[1_048_576];
    }
    foo = null;

    for(size_t i = 9; i != size_t.max; i--) {
        foreach(j, e; arrays[i]) {
            assert(e == j * (i + 1));
        }
    }
    TempAlloc.frameFree;
    assert(space == TempAlloc.state.space);
    assert(used == TempAlloc.state.used);
    while(TempAlloc.state.nblocks > 1 || TempAlloc.state.used > 0) {
        TempAlloc.free;
    }

    // Test that everything is really getting destroyed properly when
    // destroy() is called.  If not then this test will run out of memory.
    foreach(i; 0..1000) {
        TempAlloc.state.destroy();
        TempAlloc.state = null;

        foreach(j; 0..1_000) {
            auto ptr = TempAlloc.malloc(20_000);
            assert((cast(size_t) ptr) % TempAlloc.alignBytes == 0);
        }

        foreach(j; 0..500) {
            TempAlloc.free();
        }
    }
}

struct SHNode(K, V) {
    alias SHNode!(K, V) SomeType;
    SomeType* next;
    Unqual!(K) key;
    Unqual!(V) val;
}

/**Forward range struct for iterating over the keys or values of a
 * StackHash or StackSet.  The lifetime of this object must not exceed that
 * of the underlying StackHash or StackSet.*/
struct HashRange(K, S, bool vals = false) {
private:
    S* set;
    size_t index;
    S.Node* next;
    K* frontElem;
    size_t _length;

    this(S* set) {
        this.set = set;
        if(set.rNext[0] == set.usedSentinel) {
            this.popFront;
        } else {
            static if(vals) {
                frontElem = set.rVals.ptr;
            } else {
                frontElem = set.rKeys.ptr;
            }
            next = set.rNext[0];
        }
        this._length = set.length;
    }

public:
    ///
    void popFront()
    in {
        assert(!empty);
    } body {
        this._length--;
        if(next is null) {
            do {
                index++;
                if(index >= set.rNext.length) {
                    index = size_t.max;  // Sentinel for empty.
                    return;
                }
                next = set.rNext[index];
            } while(set.rNext[index] == set.usedSentinel);
            static if(vals) {
                frontElem = &(set.rVals[index]);
            } else {
                frontElem = &(set.rKeys[index]);
            }
        } else {
            static if(vals) {
                frontElem = &(next.val);
            } else {
                frontElem = &(next.key);
            }
            next = next.next;
        }
    }

    ///
    static if(vals) {
        @property ref Unqual!(K) front()
        in {
            assert(!empty);
        } body {
            return *frontElem;
        }
    } else {
       @property Unqual!(K) front()
       in {
            assert(!empty);
       } body {
            return *frontElem;
        }
    }

    ///
    @property bool empty() {
        return index == size_t.max;
    }

    ///
    @property size_t length() {
        return _length;
    }

    ///
    @property typeof(this) save() {
        return this;
    }
}

/**A hash table that allocates its memory on TempAlloc.  Good for building a
 * temporary hash tables that will not escape the current scope.
 *
 * To avoid TempAlloc memory leaks, use mixin(newFrame).
 *
 * Examples:
 * ---
 * mixin(newFrame);  // To make sure all memory gets freed at end of scope.
 * auto ss = StackHash!(uint)(5);
 * foreach(i; 0..5) {
 *     ss[i]++;
 * }
 * assert(ss[3] == 1);
 * ---
 *
 * Warning:
 * This implementation places removed nodes on an internal free list and
 * recycles them, since there is no way to delete TempAlloc-allocated data
 * in a non-LIFO order.  Therefore, you may not retain the address of a
 * variable stored in a StackHash after deleting it from the StachHash.
 * For example, DO NOT do this:
 * ---
 * SomeType* myPtr = &(myStackHash["foo"]);
 * myStackHash.remove("foo");
 * *myPtr = someValue;
 * ---
 */
struct StackHash(K, V) {
private:
    alias SHNode!(K, V) Node;

    // Using parallel arrays instead of structs to save on alignment overhead:
    Unqual!(K)[] rKeys;
    Unqual!(V)[] rVals;
    Unqual!(Node*)[] rNext;

    // Holds nodes that were deleted by remove().
    Node** freeList;

    TempAlloc.State TAState;
    size_t _length;

    // Tries to allocate off the free list.  Otherwise allocates off
    // TempAlloc.
    Node* allocNode() {
        if(*freeList is null) {
            return cast(Node*) TempAlloc(Node.sizeof, TAState);
        }
        auto ret = *freeList;
        *freeList = (*freeList).next;
        return ret;
    }

    // Add a removed node to the free list.
    void pushFreeList(Node* node) {
        if(*freeList is null) {
            node.next = null;  // Sentinel
            *freeList = node;
        } else {
            node.next = *freeList;
            *freeList = node;
        }
    }

    // rNext.ptr is stored in elements of rNext as a sentinel to indicate
    // that the corresponding slot is unused.
    Node* usedSentinel() @property {
        return cast(Node*) rNext.ptr;
    }

    Node* newNode(K key) {
        Node* ret = allocNode();
        ret.key =  key;
        ret.val =  V.init;
        ret.next = null;
        return ret;
    }

    Node* newNode(K key, V val) {
        Node* ret = allocNode();
        ret.key =  key;
        ret.val = val;
        ret.next = null;
        return ret;
    }

    hash_t getHash(K key) {
        static if(is(K : long) && K.sizeof <= hash_t.sizeof) {
            hash_t hash = cast(hash_t) key;
        } else static if(__traits(compiles, key.toHash())) {
            hash_t hash = key.toHash();
        } else {
            hash_t hash = typeid(K).getHash(&key);
        }
        hash %= rNext.length;
        return hash;
    }


public:
    /**Due to the nature of TempAlloc, you must specify on object creation
     * the approximate number of elements your table will have.  Too large a
     * number will waste space and incur poor cache performance.  Too low a
     * number will make this struct perform like a linked list.  Generally,
     * if you're building a table from some other range, some fraction of the
     * size of that range is a good guess.*/
    this(size_t nElem) {
        // Obviously, the caller can never mean zero, because this struct
        // can't work at all with nElem == 0, so assume it's a mistake and fix
        // it here.
        if(nElem == 0)
            nElem++;
        TAState = TempAlloc.getState;
        rKeys = newStack!(K)(nElem, TAState);
        rVals = newStack!(V)(nElem, TAState);

        // Allocate free list in same block with Node ptrs.  That's what the
        // + 1 is for.
        rNext = newStack!(Node*)(nElem + 1, TAState);
        freeList = &(rNext[$ - 1]);
        *freeList = null;
        rNext = rNext[0..$ - 1];

        foreach(ref rKey; rKeys) {
            rKey =  K.init;
        }
        foreach(ref rVal; rVals) {
            rVal = V.init;
        }
        foreach(ref r; rNext) {
            r = usedSentinel;
        }


    }

    /**Index an element of the range.  If it does not exist, it will be created
     * and initialized to V.init.*/
    ref V opIndex(K key) {
        hash_t hash = getHash(key);

        if(rNext[hash] == usedSentinel) {
            rKeys[hash] =  key;
            rNext[hash] = null;
            _length++;
            return rVals[hash];
        } else if(rKeys[hash] == key) {
            return rVals[hash];
        } else {  // Collision.  Start chaining.
            Node** next = &(rNext[hash]);
            while(*next !is null) {
                if((**next).key ==  key) {
                    return (**next).val;
                }
                next = &((**next).next);
            }
            *next = newNode(key);
            _length++;
            return (**next).val;
        }
    }

    ///
    V opIndexAssign(V val, K key) {
        hash_t hash = getHash(key);

        if(rNext[hash] == usedSentinel) {
            rKeys[hash] =  key;
            rVals[hash] = val;
            rNext[hash] = null;
            _length++;
            return val;
        } else if(rKeys[hash] ==  key) {
            rVals[hash] = val;
            return val;
        } else {  // Collision.  Start chaining.
            Node** next = &(rNext[hash]);
            while(*next !is null) {
                if((**next).key == key) {
                    (**next).val = val;
                    return val;
                }
                next = &((**next).next);
            }
            _length++;
            *next = newNode(key, val);
            return val;
        }
    }

    ///
    V* opIn_r(K key) {
        hash_t hash = getHash(key);

        if(rNext[hash] == usedSentinel) {
            return null;
        } else if(rKeys[hash] == key) {
            return &(rVals[hash]);
        } else {  // Collision.  Start chaining.
            Node* next = rNext[hash];
            while(next !is null) {
                if(next.key == key) {
                    return &(next.val);
                }
                next = next.next;
            }
            return null;
        }
   }

    ///
    void remove(K key) {
        hash_t hash = getHash(key);

        Node** next = &(rNext[hash]);
        if(rNext[hash] == usedSentinel) {
            return;
        } else if(rKeys[hash] == key) {
            _length--;
            if(rNext[hash] is null) {
                rKeys[hash] = K.init;
                rVals[hash] = V.init;
                rNext[hash] = usedSentinel;
                return;
            } else {
                Node* toPush = *next;

                rKeys[hash] = (**next).key;
                rVals[hash] = (**next).val;
                rNext[hash] = (**next).next;

                pushFreeList(toPush);
                return;
            }
        } else {  // Collision.  Start chaining.
            while(*next !is null) {
                if((**next).key == key) {
                    _length--;

                    Node* toPush = *next;
                    *next = (**next).next;

                    pushFreeList(toPush);
                    break;
                }
                next = &((**next).next);
            }
            return;
        }
   }

    /**Returns a forward range to iterate over the keys of this table.
     * The lifetime of the HashRange must not exceed the lifetime of this
     * StackHash.*/
    HashRange!(K, StackHash!(K, V)) keys() {
        return typeof(return)(&this);
    }

    /**Returns a forward range to iterate over the values of this table.
     * The lifetime of the HashRange must not exceed the lifetime of this
     * StackHash.*/
    HashRange!(V, StackHash!(K, V), true) values() {
       return typeof(return)(&this);
    }

    ///
    @property size_t length() const {
        return _length;
    }

    /**
    Attempt to look up a key and return a default value if the key is not
    present.
    */
    V get(K key, lazy V defaultValue) {
        auto ptr = key in this;
        if(ptr) return *ptr;
        return defaultValue;
    }

    int opApply(int delegate(ref K, ref V) dg) {
        auto k = this.keys;
        auto v = this.values;
        int res;

        while(!k.empty) {
            auto kFront = k.front;
            res = dg(kFront, v.front);
            k.popFront;
            v.popFront;
            if(res) {
                break;
            }
        }

        return res;
    }

    real efficiency() {
       uint used = 0;
       foreach(root; rNext) {
           if(root != usedSentinel) {
               used++;
           }
       }
       return cast(real) used / rNext.length;
    }
}

unittest {
    alias StackHash!(string, uint) mySh;

    {  // Basic sanity checks.
        mixin(newFrame);
        auto data = mySh(2);  // Make sure we get some collisions.
        data["foo"] = 1;
        data["bar"] = 2;
        data["baz"] = 3;
        data["waldo"] = 4;
        assert(!("foobar" in data));
        assert(*("foo" in data) == 1);
        assert(*("bar" in data) == 2);
        assert(*("baz" in data) == 3);
        assert(*("waldo" in data) == 4);
        assert(data["foo"] == 1);
        assert(data["bar"] == 2);
        assert(data["baz"] == 3);
        assert(data["waldo"] == 4);
        auto myKeys = toArray(data.keys);
        qsort(myKeys);
        assert(myKeys == cast(string[]) ["bar", "baz", "foo", "waldo"]);
        auto myValues = toArray(data.values);
        qsort(myValues);
        assert(myValues == [1U, 2, 3, 4]);
        {
            auto k = data.keys;
            auto v = data.values;
            while(!k.empty) {
                assert(data[k.front] == v.front);
                k.popFront;
                v.popFront;
            }
        }
        foreach(v; data.values) {
            assert(v > 0 && v < 5);
        }
    }

    alias StackHash!(uint, uint) mySh2;
    {   // Test remove.
        mixin(newFrame);

        auto foo = mySh2(7);
        for(uint i = 0; i < 200; i++) {
            foo[i] = i;
        }
        assert(foo.length == 200);
        for(uint i = 0; i < 200; i += 2) {
            foo.remove(i);
        }
        foreach(i; 20..200) {
            foo.remove(i);
        }
        for(uint i = 0; i < 20; i++) {
            if(i & 1) {
                assert(i in foo);
                assert(*(i in foo) == i);
            } else {
                assert(!(i in foo));
            }
        }
        auto vals = toArray(foo.values);
        assert(foo.length == 10);
        assert(vals.qsort == [1U, 3, 5, 7, 9, 11, 13, 15, 17, 19]);
    }

    { // Monte carlo unittesting against builtin hash table.
        mixin(newFrame);
        uint[uint] builtin;
        auto monteSh = mySh2(20_000);
        uint[] nums = newStack!uint(100_000);
        foreach(ref num; nums) {
            num = uniform(0U, uint.max);
        }

        foreach(i; 0..1_000_000) {
            auto index = uniform(0, cast(uint) nums.length);
            if(index in builtin) {
                assert(index in monteSh);
                assert(builtin[index] == nums[index]);
                assert(monteSh[index] == nums[index]);
                builtin.remove(index);
                monteSh.remove(index);
            } else {
                assert(!(index in monteSh));
                builtin[index] = nums[index];
                monteSh[index] = nums[index];
            }
        }

        assert(builtin.length == monteSh.length);
        foreach(k, v; builtin) {
            assert(k in monteSh);
            assert(*(k in builtin) == *(k in monteSh));
            assert(monteSh[k] == v);
        }

        // Make sure nothing is missed in iteration.  Since both keys and
        // values use the same struct, just with a few static if statements,
        // if it works for keys and simple tests work for values, it works.
        foreach(k; monteSh.keys) {
            builtin.remove(k);
        }
        assert(builtin.length == 0);

    }
}

/**A hash set that allocates its memory on TempAlloc.  Good for building a
 * temporary set that will not escape the current scope.
 *
 * To avoid TempAlloc memory leaks, use mixin(newFrame).
 *
 * Examples:
 * ---
 * mixin(newFrame);  // To make sure all memory gets freed at end of scope.
 * auto ss = StackSet!(uint)(5);
 * foreach(i; 0..5) {
 *     ss.insert(i);
 * }
 * assert(3 in ss);
 * ---
 */
struct StackSet(K) {
private:
    // Choose smallest representation of the data.
    struct Node1 {
        Node1* next;
        K key;
    }

    struct Node2 {
        K key;
        Node2* next;
    }

    static if(Node1.sizeof < Node2.sizeof) {
        alias Node1 Node;
    } else {
        alias Node2 Node;
    }

    Unqual!(K)[] rKeys;
    Node*[] rNext;

    Node** freeList;

    TempAlloc.State TAState;
    size_t _length;

    Node* usedSentinel() {
        return cast(Node*) rNext.ptr;
    }

    // Tries to allocate off the free list.  Otherwise allocates off
    // TempAlloc.
    Node* allocNode() {
        if(*freeList is null) {
            return cast(Node*) TempAlloc(Node.sizeof, TAState);
        }
        auto ret = *freeList;
        *freeList = (*freeList).next;
        return ret;
    }

    // Add a removed node to the free list.
    void pushFreeList(Node* node) {
        if(*freeList is null) {
            node.next = null;  // Sentinel
            *freeList = node;
        } else {
            node.next = *freeList;
            *freeList = node;
        }
    }

    Node* newNode(K key) {
        Node* ret = allocNode();
        ret.key = key;
        ret.next = null;
        return ret;
    }

    hash_t getHash(K key) {
        static if(is(K : long) && K.sizeof <= hash_t.sizeof) {
            hash_t hash = cast(hash_t) key;
        } else static if(__traits(compiles, key.toHash())) {
            hash_t hash = key.toHash();
        } else {
            hash_t hash = typeid(K).getHash(&key);
        }
        hash %= rNext.length;
        return hash;
    }

public:
    /**Due to the nature of TempAlloc, you must specify on object creation
     * the approximate number of elements your set will have.  Too large a
     * number will waste space and incur poor cache performance.  Too low a
     * number will make this struct perform like a linked list.  Generally,
     * if you're building a set from some other range, some fraction of the
     * size of that range is a good guess.*/
    this(size_t nElem) {
        // Obviously, the caller can never mean zero, because this struct
        // can't work at all with nElem == 0, so assume it's a mistake and fix
        // it here.
        if(nElem == 0)
            nElem++;
        TAState = TempAlloc.getState;

        // Allocate the free list as the last element of rNext.
        rNext = newStack!(Node*)(nElem + 1, TAState);
        freeList = &(rNext[$ - 1]);
        *freeList = null;
        rNext = rNext[0..$ - 1];

        foreach(ref root; rNext) {
            root = usedSentinel;
        }

        rKeys = newStack!(Unqual!(K))(nElem, TAState);
        foreach(ref root; rKeys) {
            root = K.init;
        }
    }

    ///
    void insert(K key) {
        hash_t hash = getHash(key);

        if(rNext[hash] == usedSentinel) {
            rKeys[hash] = key;
            rNext[hash] = null;
            _length++;
            return;
        } else if(rKeys[hash] == key) {
            return;
        } else {  // Collision.  Start chaining.
            Node** next = &(rNext[hash]);
            while(*next !is null) {
                if((**next).key == key) {
                    return;
                }
                next = &((**next).next);
            }
            *next = newNode(key);
            _length++;
            return;
        }
    }

    /**Returns a forward range of the elements of this struct.  The range's
     * lifetime must not exceed the lifetime of this object.*/
    HashRange!(K, typeof(this)) elems() {
        auto ret = typeof(return)(&this);
        return ret;
    }

    ///
    bool opIn_r(K key) {
        hash_t hash = getHash(key);

        if(rNext[hash] == usedSentinel) {
            return false;
        } else if(rKeys[hash] == key) {
            return true;
        } else {  // Collision.  Start chaining.
            Node* next = rNext[hash];
            while(next !is null) {
                if(next.key == key) {
                    return true;
                }
                next = next.next;
            }
            return false;
        }
   }

    ///
    void remove(K key) {
        hash_t hash = getHash(key);

        Node** next = &(rNext[hash]);
        if(rNext[hash] == usedSentinel) {
            return;
        } else if(rKeys[hash] == key) {
            _length--;
            if(rNext[hash] is null) {
                rKeys[hash] = K.init;
                rNext[hash] = usedSentinel;
                return;
            } else {
                Node* toPush = *next;

                rKeys[hash] = (**next).key;
                rNext[hash] = (**next).next;

                pushFreeList(toPush);
                return;
            }
        } else {  // Collision.  Start chaining.
            while(*next !is null) {
                if((**next).key == key) {
                    _length--;
                    Node* toPush = *next;

                    *next = (**next).next;
                    pushFreeList(toPush);
                    break;
                }
                next = &((**next).next);
            }
            return;
        }
   }

    ///
    @property size_t length() {
       return _length;
    }
}

unittest {
    { // "Normal" unittesting.
        mixin(newFrame);
        alias StackSet!(uint) mySS;
        mySS set = mySS(12);
        foreach(i; 0..20) {
            set.insert(i);
        }
        assert(toArray(set.elems).qsort == seq(0U, 20U));

        for(uint i = 0; i < 20; i += 2) {
            set.remove(i);
        }

        foreach(i; 0..20) {
            if(i & 1) {
                assert(i in set);
            } else {
                assert(!(i in set));
            }
        }
        uint[] contents;

        foreach(elem; set.elems) {
            contents ~= elem;
        }
        assert(contents.qsort == [1U,3,5,7,9,11,13,15,17,19]);
    }

    { // Monte carlo unittesting against builtin hash table.
        mixin(newFrame);
        bool[uint] builtin;
        auto monteSh = StackSet!uint(20_000);

        foreach(i; 0..1_000_000) {
            auto index = uniform(0, 100_000);
            if(index in builtin) {
                assert(index in monteSh);
                builtin.remove(index);
                monteSh.remove(index);
            } else {
                assert(!(index in monteSh));
                builtin[index] = 1;
                monteSh.insert(index);
            }
        }

        assert(builtin.length == monteSh.length);
        foreach(k, v; builtin) {
            assert(k in monteSh);
        }

        foreach(k; monteSh.elems) {
            builtin.remove(k);
        }
        assert(builtin.length == 0);
    }
}

private int height(T)(const T node) nothrow {
    return (node is null) ? 0 : node.height;
}

struct AVLNodeRealHeight(T) {
    T payload;
    typeof(this)* left;
    typeof(this)* right;
    int height;

    int balance() const nothrow @property {
        return .height(left) - .height(right);
    }

    void fixHeight() nothrow {
        auto leftHeight = .height(left);
        auto rightHeight = .height(right);

        height = ((leftHeight > rightHeight) ? leftHeight : rightHeight) + 1;
    }

    bool isLeaf() nothrow @property {
        return left is null && right is null;
    }
}

/* Store the height in the low order bits of the pointers to save space,
 * since TempAlloc allocates 16-byte aligned memory anyhow, but only if
 * this would be smaller after considering alignment.
 */
struct AVLNodeBitwise(T) {
    T payload;
    size_t _left;
    size_t _right;

    enum size_t mask = 0b1111;
    enum size_t notMask = ~mask;

    typeof(this)* left() nothrow @property {
        return cast(typeof(return)) (_left & notMask);
    }

    const(typeof(this))* left() const nothrow @property {
        return cast(typeof(return)) (_left & notMask);
    }

    void left(typeof(this)* newLeft) nothrow @property
    in {
        assert((cast(size_t) newLeft & mask) == 0);
    } body {
        _left &= mask;
        _left |= cast(size_t) newLeft;
        assert(left is newLeft);
    }

    typeof(this)* right() nothrow @property {
        return cast(typeof(return)) (_right & notMask);
    }

    const(typeof(this))* right() const nothrow @property {
        return cast(typeof(return)) (_right & notMask);
    }

    void right(typeof(this)* newRight) nothrow @property
    in {
        assert((cast(size_t) newRight & mask) == 0);
    } body {
        _right &= mask;
        _right |= cast(size_t) newRight;
        assert(right is newRight);
    }

    int height() const nothrow @property {
        return (((_left & mask) << 4) |
                    (_right & mask));
    }

    void height(int newHeight) nothrow @property {
        _right &= notMask;
        _right |= (newHeight & mask);
        newHeight >>= 4;
        _left &= notMask;
        _left |= (newHeight & mask);
    }

    int balance() const nothrow @property {
        return .height(left) - .height(right);
    }

    void fixHeight() nothrow {
        auto leftHeight = .height(left);
        auto rightHeight = .height(right);

        height = ((leftHeight > rightHeight) ? leftHeight : rightHeight) + 1;
    }

    bool isLeaf() const nothrow @property {
        return left is null && right is null;
    }
}

private template GetAligned(uint size) {
    static if(size % TempAlloc.alignBytes == 0) {
        enum GetAligned = 0;
    } else {
        enum GetAligned =
            size - size % TempAlloc.alignBytes + TempAlloc.alignBytes;
    }
}

/**An AVL tree implementation on top of TempAlloc.  If elements are removed,
 * they are stored on an internal free list and recycled when new elements
 * are added to the tree.
 *
 * Template paramters:
 *
 * T = The type to be stored in the tree.
 *
 * key = Function to access the key that what you're storing is to be compared
 *       on.
 *
 * compFun = The function for comparing keys.
 *
 * Examples:
 * ---
 * struct StringNum {
 *     string someString;
 *     uint num;
 * }
 *
 * // Create a StackTree of StringNums, sorted in descending order, using
 * // someString for comparison.
 * auto myTree = StackTree!(StringNum, "a.someString", "a > b")();
 *
 * // Add some elements.
 * myTree.insert( StringNum("foo", 1));
 * myTree.insert( StringNum("bar", 2));
 * myTree.insert( StringNum("foo", 3));
 *
 * assert(myTree.find("foo") == StringNum("foo", 3));
 * assert(myTree.find("bar") == StringNum("bar", 2));
 * ---
 *
 * Note:  This tree supports a compile-time interface similar to StackSet
 * and can be used as a finite set implementation.
 *
 * Warning:
 * This implementation places removed nodes on an internal free list and
 * recycles them, since there is no way to delete TempAlloc-allocated data
 * in a non-LIFO order.  Therefore, you may not retain the address of a
 * variable stored in a StackTree after deleting it from the StackTree.
 * For example, DO NOT do this:
 * ---
 * SomeType* myPtr = "foo" in myTree;
 * myTree.remove("foo");
 * *myPtr = someValue;
 * ---
 */
struct StackTree(T, alias key = "a", alias compFun = "a < b") {
private:

    alias AVLNodeBitwise!(T) BitwiseNode;
    alias AVLNodeRealHeight!(T) RealNode;

    enum size_t bitSize = GetAligned!(BitwiseNode.sizeof);
    enum size_t realHeightSize = GetAligned!(RealNode.sizeof);

    static if(bitSize < realHeightSize ) {
        alias AVLNodeBitwise!(T) Node;
    } else {
        alias AVLNodeRealHeight!(T) Node;
    }

    alias binaryFun!(compFun) comp;
    alias unaryFun!(key) getKey;

    Node* head;
    Node** freeList;
    size_t _length;
    TempAlloc.State TAState;

    static bool insertComp(T lhs, T rhs) {
        return comp( getKey(lhs), getKey(rhs));
    }

    static Node* rotateRight(Node* node)
    in {
        assert(node.left !is null);
        assert( abs(node.balance) <= 2);

    } body {
        Node* newHead = node.left;
        node.left = newHead.right;
        newHead.right = node;

        node.fixHeight();
        newHead.fixHeight();

        assert( abs(node.balance) < 2);
        return newHead;
    }

    static Node* rotateLeft(Node* node)
    in {
        assert(node.right !is null);
        assert( abs(node.balance) <= 2);
    } body {
        Node* newHead = node.right;
        node.right = newHead.left;
        newHead.left = node;

        node.fixHeight();
        newHead.fixHeight();

        assert( abs(node.balance) < 2);
        return newHead;
    }

    static Node* rebalance(Node* node)
    in {
        assert(node is null || abs(node.balance) <= 2);
    } out(ret) {
        assert( abs(ret.balance) < 2);
    } body {
        if(node is null) {
            return null;
        }

        immutable balance = node.balance;
        if(abs(balance) <= 1) {
            return node;
        }

        if(balance == -2) {

            // It should be impossible to have a balance factor of -2 if
            // node.right is null.
            assert(node.right !is null);
            immutable rightBalance = node.right.balance;
            assert( abs(rightBalance) < 2);

            if(rightBalance == 1) {
                node.right = rotateRight(node.right);
                node.fixHeight();
            }

            assert(node.balance == -2);
            return rotateLeft(node);

        } else if(balance == 2) {
            // It should be impossible to have a balance factor of 2 if
            // node.left is null.
            assert(node.left !is null);
            immutable leftBalance = node.left.balance;
            assert( abs(leftBalance) < 2);

            if(leftBalance == -1) {
                node.left = rotateLeft(node.left);
                node.fixHeight();
            }

            assert(node.balance == 2);
            return rotateRight(node);
        }

        // AVL tree invariant is that abs(balance) <= 2 even during
        // insertion/deletion.
        assert(0);
    }

    void pushFreeList(Node* node) {
        node.left = null;
        node.right = *freeList;
        *freeList = node;
    }

    Node* popFreeList()
    in {
        assert(freeList);
        assert(*freeList);
    } body {
        auto ret = *freeList;
        *freeList = ret.right;
        return ret;
    }

    Node* newNode(T payload)
    in {
        assert(freeList, "Uninitialized StackTree!(" ~ T.stringof ~ ")");
        assert(TAState, "Uninitialized StackTree!(" ~ T.stringof ~ ")");
    } body {
        Node* ret;
        if(*freeList !is null) {
            ret = popFreeList();
        } else {
            ret = cast(Node*) TempAlloc.malloc(Node.sizeof, TAState);
        }

        ret.payload = payload;
        ret.left = null;
        ret.right = null;
        ret.height = 1;
        return ret;
    }

public:
    /**De facto constructor.  Not using a "real" c'tor only because structs
     * don't support default c'tors yet.  This must be called, or else you will
     * get an access violation when you try to insert an element.
     */
    static typeof(this) opCall() {
        typeof(this) ret;
        ret.TAState = TempAlloc.getState();
        ret.freeList = newStack!(Node*)(1).ptr;
        *(ret.freeList) = null;
        return ret;
    }

    /**Insert an element.*/
    void insert(T toInsert) {
        if(head is null) {
            head = newNode(toInsert);
            _length++;
        } else {
            head = insertImpl(toInsert, head);
        }
    }

    Node* insertImpl(T toInsert, Node* insertInto) {
        if( insertComp(toInsert, insertInto.payload) ) {
            if(insertInto.left is null) {
                insertInto.left = newNode(toInsert);
                _length++;
            } else {
                insertInto.left = insertImpl(toInsert, insertInto.left);
            }
        } else if( insertComp(insertInto.payload, toInsert) ) {
            if(insertInto.right is null) {
                insertInto.right = newNode(toInsert);
                _length++;
            } else {
                insertInto.right = insertImpl(toInsert, insertInto.right);
            }
        } else {
            // This is correct:  If the comparison key is only part of the
            // payload, the old payload may not be equal to the new payload,
            // even if the comparison keys are equal.
            insertInto.payload = toInsert;
            return insertInto;
        }

        insertInto.fixHeight();
        return rebalance(insertInto);
    }

    /**Remove an element from this tree.  The type of U is expected to be the
     * type of the key that this tree is sorted on.
     */
    void remove(U)(U whatToRemove) {
        Node* removedNode;
        Node* leftMost;

        Node* removeLeftMost(Node* node) {
            if(node.left is null) {
                auto ret = node.right;
                node.right = null;
                leftMost = node;
                return ret;
            }

            node.left = removeLeftMost(node.left);
            node.fixHeight();
            return rebalance(node);
        }

        Node* removeSuccessor(Node* node) {
            if(node.right is null) {
                assert(node.left.isLeaf);
                leftMost = node.left;

                node.left = null;
                return node;
            }

            node.right = removeLeftMost(node.right);
            node.fixHeight();
            return node;
        }

        Node* removeImpl(U whatToRemove, Node* whereToRemove) {
            static bool findComp(V, W)(V lhs, W rhs) {
                static if(is(V == T)) {
                    static assert(is(W == U));
                    return comp( getKey(lhs), rhs);
                } else {
                    static assert(is(V == U));
                    static assert(is(W == T));
                    return comp(lhs, getKey(rhs) );
                }
            }

            if(whereToRemove is null) {
                return null;
            }

            if( findComp(whatToRemove, whereToRemove.payload) ){
                whereToRemove.left = removeImpl(whatToRemove, whereToRemove.left);
                whereToRemove.fixHeight();
                return rebalance(whereToRemove);
            } else if( findComp(whereToRemove.payload, whatToRemove) ) {
                whereToRemove.right = removeImpl(whatToRemove, whereToRemove.right);
                whereToRemove.fixHeight();
                return rebalance(whereToRemove);
            } else {
                // We've found it.
                _length--;
                removedNode = whereToRemove;
                if(whereToRemove.isLeaf) {
                    return null;
                }

                whereToRemove = removeSuccessor(whereToRemove);
                if(leftMost is null) {
                    return null;
                }

                leftMost.left = whereToRemove.left;
                leftMost.right = whereToRemove.right;
                leftMost.fixHeight();
                return rebalance(leftMost);
            }
        }

        head = removeImpl(whatToRemove, head);

        debug(EXPENSIVE) assertAvl(head);

        if(removedNode !is null) {
            pushFreeList(removedNode);
        }
    }

    /**Find an element and return it.  Throw an exception if it is not
     * present.  U is expected to be the type of the key that this tree is
     * sorted on.*/
    T find(U)(U whatToFind) {
        T* ptr = dstatsEnforce( opIn_r!(U)(whatToFind),
            "Item not found:  " ~ to!string(whatToFind));
        return *ptr;
    }

    /**Find an element and return a pointer to it, or null if not present.*/
    T* opIn_r(U)(U whatToFind) {
        auto ret = findImpl!(U)(whatToFind, head);
        if(ret is null) {
            return null;
        }
        return &(ret.payload);
    }

    Node* findImpl(U)(U whatToFind, Node* whereToFind) {
        static bool findComp(V, W)(V lhs, W rhs) {
            static if(is(V == T)) {
                static assert(is(W == U));
                return comp( getKey(lhs), rhs );
            } else {
                static assert(is(V == U));
                static assert(is(W == T));
                return comp( lhs, getKey(rhs) );
            }
        }

        if(whereToFind is null) {
            return null;
        }

        if( findComp(whatToFind, whereToFind.payload) ){
            return findImpl!(U)(whatToFind, whereToFind.left);
        } else if( findComp(whereToFind.payload, whatToFind) ) {
            return findImpl!(U)(whatToFind, whereToFind.right);
        } else {
            // We've found it.
            return whereToFind;
        }

        assert(0);
    }

    /**Iterate over the elements of this tree in sorted order.*/
    int opApply( int delegate(ref T) dg) {
        int res;
        int opApplyImpl(Node* node) {
            if(node is null) {
                return 0;
            }
            res = opApplyImpl(node.left);
            if(res) {
                return res;
            }
            res = dg(node.payload);
            if(res) {
                return res;
            }
            res = opApplyImpl(node.right);
            return res;
        }

        return opApplyImpl(head);
    }

    /**Number of elements in the tree.*/
    @property size_t length() const pure nothrow @property {
        return _length;
    }
}

private int assertAvl(T)(T node) {
    if(node is null) {
        return 0;
    }

    int leftHeight = assertAvl(node.left);
    int rightHeight = assertAvl(node.right);
    assert(node.height == max(leftHeight, rightHeight) + 1);
    assert( abs(node.balance) < 2,
        text( height(node.left), '\t', height(node.right)));

    if(node.left) {
        assert(node.left.payload < node.payload);
    }

    if(node.right) {
        assert(node.right.payload > node.payload,
            text(node.payload, ' ', node.right.payload));
    }

    return node.height;
}


unittest {
    // Test against StackSet on random data.
    mixin(newFrame);
    StackTree!(uint) myTree = StackTree!(uint)();
    StackSet!(uint) ss = StackSet!(uint)(500);
    foreach(i; 0..1_000_000) {
        uint num = uniform(0, 1_000);
        if(num in ss) {
            assert(num in myTree);
            assert(*(num in myTree) == num);
            ss.remove(num);
            myTree.remove(num);
        } else {
            assert(!(num in myTree));
            ss.insert(num);
            myTree.insert(num);
        }
    }
    assertAvl(myTree.head);
}

/**Struct that iterates over keys or values of a StackTreeAA.
 *
 * Bugs:  Uses opApply instead of the more flexible ranges, because I
 * haven't figured out how to iterate efficiently and in sorted order over a
 * tree without control of the stack.
 */
struct TreeAaIter(T, alias mapFun) {
    alias unaryFun!(mapFun) mFun;
    T tree;
    alias typeof(*(tree.head)) Node;

//    TreeRange!(T, mapFun) asRange() {
//        dstatsEnforce(0, "Not implemented yet.");
//    }

    alias typeof( mFun( typeof(tree.head.payload).init) ) IterType;

    ///
    int opApply( int delegate(ref IterType) dg) {
        int res;
        int opApplyImpl(Node* node) {
            if(node is null) {
                return 0;
            }
            res = opApplyImpl(node.left);
            if(res) {
                return res;
            }

            static if(__traits(compiles, dg(mFun(node.payload)))) {
                res = dg(mFun(node.payload));
            } else {
                auto toDg = mFun(node.payload);
                res = dg(toDg);
            }

            if(res) {
                return res;
            }
            res = opApplyImpl(node.right);
            return res;
        }

        return opApplyImpl(tree.head);
    }

    ///
    @property size_t length() const pure nothrow {
        return tree.length;
    }
}

private struct StackTreeAANode(K, V) {
    Unqual!(K) key;
    Unqual!(V) value;
}

/**An associative array implementation based on StackTree.  Lookups and
 * insertions are O(log N).  This is significantly slower in both theory and
 * practice than StackHash, but you may want to use it if:
 *
 * 1.  You don't know the approximate size of the table you will be creating
 *     in advance.  Unlike StackHash, this AA implementation does not need
 *     to pre-allocate anything.
 *
 * 2.  You care more about worst-case performance than average-case
 *     performance.
 *
 * 3.  You have a good comparison function for your type, but not a good hash
 *     function.
 *
 */
struct StackTreeAA(K, V) {
    alias StackTreeAANode!(K, V) Node;
    StackTree!(Node, "a.key") tree;

    static typeof(this) opCall() {
        typeof(this) ret;
        ret.tree = typeof(tree)();
        return ret;
    }

    /**Looks up key in the table, returns it by reference.  If it does not
     * exist, it will be created and initialized to V.init.  This is handy,
     * for example, when counting things with integer types.
     */
    ref V opIndex(K key) {
        Node* result = key in tree;
        if(result is null) {
            tree.insert( Node(key, V.init));
            result = key in tree;
        }

        return result.value;
    }

    ///
    V opIndexAssign(V val, K key) {
        tree.insert( Node(key, val));
        return val;
    }

    ///
    V* opIn_r(K key) {
        auto nodePtr = key in tree;
        if(nodePtr is null) {
            return null;
        }

        return &(nodePtr.value);
    }

    ///
    void remove(K key) {
        tree.remove(key);
    }

    ///
    @property size_t length() const pure nothrow {
        return tree.length;
    }

    ///
    TreeAaIter!( typeof(tree), "a.key") keys() @property {
        typeof(return) ret;
        ret.tree = tree;
        return ret;
    }

    private static ref Unqual!(V) getVal(ref Node node) {
        return node.value;
    }

    ///
    TreeAaIter!( typeof(tree), getVal) values() @property {
        return typeof(return)(tree);
    }

    /**Iterate over both the keys and values of this associative array.*/
    int opApply( int delegate(ref Unqual!(K), ref Unqual!(V)) dg) {
        alias typeof(*(tree.head)) TreeNode;
        int res;
        int opApplyImpl(TreeNode* node) {
            if(node is null) {
                return 0;
            }
            res = opApplyImpl(node.left);
            if(res) {
                return res;
            }
            res = dg(node.payload.key, node.payload.value);
            if(res) {
                return res;
            }
            res = opApplyImpl(node.right);
            return res;
        }

        return opApplyImpl(tree.head);
    }

}

unittest {
    // Test against builtin AA on random data.
    {
        mixin(newFrame);
        alias StackTreeAA!(string, uint) mySh;
        auto data = mySh();
        data["foo"] = 1;
        data["bar"] = 2;
        data["baz"] = 3;
        data["waldo"] = 4;
        assert(!("foobar" in data));
        assert(*("foo" in data) == 1);
        assert(*("bar" in data) == 2);
        assert(*("baz" in data) == 3);
        assert(*("waldo" in data) == 4);
        assert(data["foo"] == 1);
        assert(data["bar"] == 2);
        assert(data["baz"] == 3);
        assert(data["waldo"] == 4);

        assert(data.length == 4);
        auto myKeys = toArray(data.keys);
        qsort(myKeys);
        assert(myKeys == cast(string[]) ["bar", "baz", "foo", "waldo"]);
        auto myValues = toArray(data.values);
        qsort(myValues);
        assert(myValues == [1U, 2, 3, 4]);

        foreach(v; data.values) {
            assert(v > 0 && v < 5);
        }
    }

    alias StackTreeAA!(uint, uint) mySh2;
    {   // Test remove.
        mixin(newFrame);

        auto foo = mySh2();
        for(uint i = 0; i < 200; i++) {
            foo[i] = i;
        }
        assert(foo.length == 200);
        for(uint i = 0; i < 200; i += 2) {
            foo.remove(i);
        }
        foreach(i; 20..200) {
            foo.remove(i);
        }
        for(uint i = 0; i < 20; i++) {
            if(i & 1) {
                assert(i in foo);
                assert(*(i in foo) == i);
            } else {
                assert(!(i in foo));
            }
        }
        auto vals = toArray(foo.values);
        assert(foo.length == 10);
        assert(vals.qsort == [1U, 3, 5, 7, 9, 11, 13, 15, 17, 19]);
    }

    { // Monte carlo unittesting against builtin hash table.
        mixin(newFrame);
        uint[uint] builtin;
        auto monteSh = mySh2();
        uint[] nums = newStack!uint(100_000);
        foreach(ref num; nums) {
            num = uniform(0U, uint.max);
        }

        foreach(i; 0..10_000) {
            auto index = uniform(0, cast(uint) nums.length);
            if(index in builtin) {
                assert(index in monteSh);
                assert(builtin[index] == nums[index]);
                assert(monteSh[index] == nums[index]);
                builtin.remove(index);
                monteSh.remove(index);
            } else {
                assert(!(index in monteSh));
                builtin[index] = nums[index];
                monteSh[index] = nums[index];
            }
        }

        assert(builtin.length == monteSh.length);
        foreach(k, v; builtin) {
            assert(k in monteSh);
            assert(*(k in builtin) == *(k in monteSh));
            assert(monteSh[k] == v);
        }

        // Make sure nothing is missed in iteration.  Since both keys and
        // values use the same struct, just with a few static if statements,
        // if it works for keys and simple tests work for values, it works.
        foreach(k; monteSh.keys) {
            builtin.remove(k);
        }
        assert(builtin.length == 0);
    }
}