/**
 * \file This holds all stufff related our memory managent.
 * I try the best as far as I can to reduce memory fragmentation
 * and unneccessary calls to alloc and free.
 *
 * To achive this I try an approach described here as "Quick Fit".
 * http://www.flounder.com/memory_allocation.htm
 *
 * The basic idea is to keep allocated memory segments and don't free
 * them again. Instead I will put them in a tree indexed by their size.
 * To get new memory I first have a look in the tree if there is
 * a fitting memory segment. Fitting mean, larger or exactly the size
 * I need. If there is one, use it. If not create a new one using 
 * usual malloc approach.
 * I won't split the reagions at all because most likely they will be
 * free soon again. This way I might waste some memory, so I have to
 * keep an eye on this.
 *
 * Right now I don't build an upper limit for allocation. The limit
 * still is the system memory itself.
 *
 * This is not implemented as a class because it will be used in the 
 * process of object creation.
 *
 * The data structure is a balanced tree with size as key.
 * Under the size key is a list of elements of the same size.
 *
 * \author	Georg Hopp
 *
 * \copyright
 * Copyright © 2012  Georg Hopp
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#define _GNU_SOURCE

#include <stdio.h>

#include <stdlib.h>
#include <string.h>
#include <search.h>
#include <unistd.h>

#include "tr/memory.h"
#include "tr/tree_macros.h"


struct memSegment
{
	size_t   ref_count;
    size_t   size;
    void   * ptr;

    TR_rbColor color;

	struct memSegment * data;

    struct memSegment * next;
    struct memSegment * last;

    struct memSegment * parent;
    struct memSegment * left;
    struct memSegment * right;
};

static
struct memSegment *
newElement(size_t size)
{
    struct memSegment * element = malloc(size);

	element->ref_count = 1;
    element->size      = size;
    element->ptr       = (void*)element + sizeof(struct memSegment);

	element->data      = element;

    element->next      = NULL;
    element->last      = NULL;

    element->color     = rbRed;
    element->parent    = NULL;
    element->left      = NULL;
    element->right     = NULL;

    return element;
}

static
int
_memSegmentFindCompare(const void * a, const void * b)
{
    struct memSegment * _a = (struct memSegment *)a;
    size_t              _b = *(size_t *)b;

    /*
     * find the smallest bigger or equal size segment
     */
    return _a->size < _b ? -1
        : _a->size > _b && _a->left && _a->left->size >= _b ? 1 : 0;
}

static
int
_memSegmentCompare(const void * a, const void * b)
{
    size_t _a = ((struct memSegment *)a)->size;
    size_t _b = ((struct memSegment *)b)->size;

    return _a < _b ? -1 : _a > _b ? 1 : 0;
}

/**
 * find element in tree
 */
static
struct memSegment *
findElement(struct memSegment * tree, size_t size)
{
	int found;

	TR_TREE_FIND(tree, &size, found, _memSegmentFindCompare);

	return found == 0 ? tree : NULL;
}

/**
 * insert element in tree
 */
static
struct memSegment *
insertElement(struct memSegment ** tree, struct memSegment * element)
{
    struct memSegment * node     = *tree;
    struct memSegment * new_node = NULL;
    struct memSegment * u;
    struct memSegment * g;
	int                 found;

    element->next   = NULL;
    element->last   = NULL;

    element->color  = rbRed;
    element->parent = NULL;
    element->left   = NULL;
    element->right  = NULL;

	TR_TREE_FIND(node, element, found, _memSegmentCompare);

    // if tree is empty it's simple... :)
    if (NULL == node) {
        *tree = node = new_node = element;
	} else {
		// normal binary tree add....
		if (found == 0) {
			if (NULL == node->next) {
				node->next = element;
				node->last = element;
			} else {
                node->last->next = element;
                node->last       = element;
			}
			return node;
		} else {
			if (0 < found) {
				node->left         = element;
				node->left->parent = node;
				new_node = node = node->left;
			} else {
				node->right         = element;
				node->right->parent = node;
				new_node = node = node->right;

			}
		}
	}

    if (NULL != new_node) {
        /* 
         * handle reballancing rb style
         */
        while (1) {
            // case 1
            if (node->parent == NULL) {
                node->color = rbBlack;
                // we're done.... :)
                break;
            }

            // case 2
            if (node->parent->color == rbBlack) {
                // Tree is still valid ... wow, again we're done... :)
                break;
            }

            // case 3
            u = TR_TREE_UNCLE(node);
            g = TR_TREE_GRANDPARENT(node);

            if (u != NULL && u->color == rbRed) {
                node->parent->color = rbBlack;
                u->color            = rbBlack;
                g->color            = rbRed;

                node = g;
                continue;
            }

            // case 4
            if (node == node->parent->right && node->parent == g->left) {
                TR_TREE_ROTATE(left, tree, node->parent);
                node = node->left; 
            } else if (node == node->parent->left && node->parent == g->right) {

                TR_TREE_ROTATE(right, tree, node->parent);
                node = node->right; 
            }

            // case 5
            g = TR_TREE_GRANDPARENT(node);

            node->parent->color = rbBlack;
            g->color            = rbRed;

            if (node == node->parent->left) {
                TR_TREE_ROTATE(right, tree, g);
            } else {
                TR_TREE_ROTATE(left, tree, g);
            }

            // we're done..
            break;
        }
    }

    return new_node;
}

static
struct memSegment *
deleteElement(struct memSegment ** tree, struct memSegment * element)
{
    struct memSegment * node = *tree;
    struct memSegment * del_node;
    struct memSegment * child;
    struct memSegment * s;
	int                 found;

    // find the relevant node and it's parent
	TR_TREE_FIND(node, element, found, _memSegmentCompare);

	//while (node) {
    if (found != 0) {
		return NULL;
	} else {
		if (NULL != node->next) {
			if (NULL != node->parent) {
				if (node == node->parent->left) {
					node->parent->left = node->next;
				} else {
					node->parent->right = node->next;
				}
			} else {
				*tree = node->next;
			}

			if (NULL != node->left) {
				node->left->parent = node->next;
			}

			if (NULL != node->right) {
				node->right->parent = node->next;
			}
			
			node->next->last   = node->last;
			node->next->color  = node->color;
			node->next->parent = node->parent;
			node->next->left   = node->left;
			node->next->right  = node->right;

			return node;
		}
	}

    del_node = node;

    // now our cases follows...the first one is the same as with
    // simple binary search trees. Two non null children.

    // case 1: two children
    if (NULL != node->left && NULL != node->right) {
        struct memSegment * successor;
        struct memSegment * tmpparent;
        struct memSegment * tmpleft;
        struct memSegment * tmpright;
        TR_rbColor tmpcolor;

        TR_TREE_INORDER_SUCC(node, successor);
        tmpparent = successor->parent;
        tmpleft   = successor->left;
        tmpright  = successor->right;
        tmpcolor  = successor->color;

        TR_TREE_REPLACE_NODE(tree, node, successor);

        successor->color        = node->color;
        successor->left         = node->left;
        successor->left->parent = successor;
        // the right one might be successor...
        if (node->right == successor) {
            successor->right = node;
            node->parent     = successor;
        } else {
            successor->right    = node->right;
            node->right->parent = successor;
            node->parent        = tmpparent;
            tmpparent->left     = node;
        }

        node->color  = tmpcolor;
        node->left   = tmpleft;
        node->right  = tmpright;
    }

    // Precondition: n has at most one non-null child.
    child = (NULL == node->right) ? node->left : node->right;
    TR_TREE_REPLACE_NODE(tree, node, child);

    // delete one child case
    // TODO this is overly complex as simply derived from the function...
    // maybe this can be simplified. Maybe not...check.
    if (node->color == rbBlack) {
        if (NULL != child && child->color == rbRed) {
            child->color = rbBlack;
            // done despite modifying tree itself if neccessary..
            return del_node;
        } else {
            if (NULL != child) {
                node = child;
            } else {
                node->color = rbBlack;
                node->left  = NULL;
                node->right = NULL;
            }
        }
    } else {
        return del_node;
    }

    // delete and rb rebalance...
    while(1) {
        // case 1
        if (NULL == node->parent) {
            // done again
            break;
        }

        // case 2
        s = TR_TREE_SIBLING(node);

        if (NULL != s && s->color == rbRed) {
            node->parent->color = rbRed;
            s->color            = rbBlack;

            /*
             * detect which child we are...assumption
             * if we are not parent->right and parent->right is not
             * null we must be left, even if its set to NULL previously
             */
            if (NULL != node->parent->right && node != node->parent->right) {
                TR_TREE_ROTATE(left, tree, node->parent);
            } else {
                TR_TREE_ROTATE(right, tree, node->parent);
            }
        }

        s = TR_TREE_SIBLING(node);
        // case 3 / 4
        if (NULL == s || ((s->color == rbBlack) &&
                    (NULL == s->left || s->left->color == rbBlack) &&
                    (NULL == s->right || s->right->color == rbBlack))) {

            if (NULL != s) {
                s->color = rbRed;
            }

            if (node->parent->color == rbBlack) {
                // case 3
                node = node->parent;
                continue;
            } else {
                // case 4
                node->parent->color = rbBlack;
                // and done again...
                break;
            }
        }

        // case 5
        if  (NULL != s && s->color == rbBlack) {
            // this if statement is trivial,
            // due to case 2 (even though case 2 changed the sibling to a
            // sibling's child,
            // the sibling's child can't be red, since no red parent can
            // have a red child).
            //
            // the following statements just force the red to be on the
            // left of the left of the parent,
            // or right of the right, so case 6 will rotate correctly.
            if ((node == node->parent->left) &&
                    (NULL == s->right || s->right->color == rbBlack) &&
                    (NULL != s->left && s->left->color == rbRed)) {

                // this last test is trivial too due to cases 2-4.
                s->color       = rbRed;
                s->left->color = rbBlack;

                TR_TREE_ROTATE(right, tree, s);
            } else if ((node == node->parent->right) &&
                    (NULL == s->left || s->left->color == rbBlack) &&
                    (NULL != s->right && s->right->color == rbRed)) {
                // this last test is trivial too due to cases 2-4.
                s->color        = rbRed;
                s->right->color = rbBlack;

                TR_TREE_ROTATE(left, tree, s);
            }
        }

        s = TR_TREE_SIBLING(node);
        // case 6
        if (NULL != s) {
            s->color = node->parent->color;
        }

        if (NULL != node && NULL != node->parent) {
            node->parent->color = rbBlack;

            /*
             * detect which child we are...assumption
             * if we are not parent->right and parent->right is not
             * null we must be left, even if its set to NULL previously
             */
            if (NULL != node->parent->right && node != node->parent->right) {
                if (NULL != s->right) {
                    s->right->color = rbBlack;
                }
                TR_TREE_ROTATE(left, tree, node->parent);
            } else {
                if (NULL != s->left) {
                    s->left->color = rbBlack;
                }
                TR_TREE_ROTATE(right, tree, node->parent);
            }
        }

        // done...
        break;
    }
 
    return del_node;
}

static
void
post(struct memSegment * tree, void (*cb)(struct memSegment *, int))
{
    struct memSegment * previous = tree;
    struct memSegment * node     = tree;
    int                 depth    = 1;

    /*
     * I think this has something like O(n+log(n)) on a ballanced
     * tree because I have to traverse back the rightmost leaf to
     * the root to get a break condition.
     */
    while (node) {
        /*
         * If we come from the right so nothing and go to our
         * next parent.
         */
        if (((NULL == node->left || previous == node->left)
					&& NULL == node->right)
                || previous == node->right) {

            struct memSegment * parent = node->parent;

            cb(node, depth);

            previous = node;
            node     = parent;
            depth--;
            continue;
        }

        if ((NULL == node->left || previous == node->left)) {
            /*
             * If there are no more elements to the left or we
             * came from the left, process data.
             */
            previous = node;

            if (NULL != node->right) {
                node = node->right;
                depth++;
            } else {
                node = node->parent;
                depth--;
            }
        } else {
            /*
             * if there are more elements to the left go there.
             */
            previous = node;
            node     = node->left;
            depth++;
        }
    }
}

static
struct memSegment * segments = NULL;

static
void
segmentFree(struct memSegment * segment, int depth)
{
    while (NULL != segment) {
        struct memSegment * next = segment->next;
        free(segment);
        segment = next;
    }
}

void *
TR_reference(void * mem)
{
	struct memSegment * seg = (mem - sizeof(struct memSegment));

	seg->ref_count++;

	return mem;
}

/*
 * This tries to reflect the memory management behaviour of the
 * GNU version of malloc. For other versions this might need
 * to be changed to be optimal.
 *
 * However, GNU malloc keeps separate pools for each power of
 * 2 memory size up to page size. So one page consists all of
 * memory blocks of the same sizei (a power of 2).
 *
 * Also as far as I understand the smallest allocatable block is
 * 8 bytes. At least the adresses are alwayse a multiple of 8.
 *
 * So lets say page size is 4096. There is nothing allocated
 * right now. We allocate a block of 8 bytes. This will request
 * a memory page from the OS. Then define it as a page containing
 * 8 byte blocks and return the address of the first one of these.
 * Any subsequent call to malloc for 8 bytes will return one of the
 * blocks within this page as long as there are some left.
 *
 * So what we do here is up to page size round the request size up
 * to the next power of 2 >= 8.
 * Sizes greater then pagesize will be round up to the next
 * multiple of pagesize. As far as I understand these are not
 * pooled anyway.
 *
 * For now this assumes we are on a little endian machine.
 */
void *
TR_malloc(size_t size)
{
	struct memSegment * seg   = NULL;
	long                psize = sysconf(_SC_PAGESIZE);

	size += sizeof(struct memSegment);

	if (size > psize) {
		if (0 != (size % psize)) {
			// size if not a multiple of pagesize so bring it to one.
			size = ((size / psize) + 1) * psize;
		}
	} else {
		if (size < 8) {
			size = 8;
		} else {
			size_t check = size >> 4;
			size_t mask  = 0x1F;

			while (check >>= 1) {
				mask = (mask << 1) | 1;
			}

			if (size != (size & ~(mask >> 1))) {
				size = (size << 1) & ~mask;
			}
		}
	}

#ifdef MEM_OPT
	seg = findElement(segments, size);
#endif

	if (NULL == seg) {
		seg = newElement(size);
	} else {
		// remove the found one from the tree as we use it now.
		seg = deleteElement(&segments, seg);
	}

	return seg->ptr;
}

/**
 * this is a really memory wasting solution....just to be able to
 * use calloc, which might be faster then malloc/memset solution.
 *
 * Maybe this is a bad idea, as we need to memset the buffer anyway
 * if it comes from our tree, which hopefully should be the majority
 * of cases.
 */
void *
TR_calloc(size_t nmemb, size_t size)
{
	size_t   _size = nmemb * size;
	void   * mem   = TR_malloc(_size);

	memset(mem, 0, _size);

	return mem;
}

void
TR_free(void ** mem)
{
	if (NULL != *mem) {
		struct memSegment * seg = (*mem - sizeof(struct memSegment));

		if (1 < seg->ref_count) {
			seg->ref_count--;
		} else {
#ifdef MEM_OPT
			insertElement(&segments, seg);
#else
			free(seg);
#endif
		}

		*mem = NULL;
	}
}

size_t
TR_getSize(void * mem)
{
	struct memSegment * segment;

	if (NULL == mem) {
		return 0;
	}

	segment = (struct memSegment *)(mem - sizeof(struct memSegment));
	return segment->size;
}

void
TR_cleanup()
{
#ifdef MEM_OPT
	post(segments, segmentFree);
#endif
}

// vim: set ts=4 sw=4: