/*
 * Copyright 2003-2009, Ingo Weinhold <ingo_weinhold@gmx.de>.
 * Distributed under the terms of the MIT License.
 */


#include <util/AVLTreeBase.h>

#ifndef FS_SHELL
#       include <algorithm>
#       include <KernelExport.h>
#endif

#ifdef _KERNEL_MODE
#       define CHECK_FAILED(message...) panic(message)
#else
#       ifndef FS_SHELL
#               include <stdio.h>
#               include <OS.h>
#               define CHECK_FAILED(message...)                                 \
                        do {                                                                            \
                                fprintf(stderr, message);                               \
                                fprintf(stderr, "\n");                                  \
                                debugger("AVLTreeBase check failed");   \
                        } while (false)
#       else
#               define CHECK_FAILED(message...) dprintf(message)
#       endif
#endif


// maximal height of a tree
static const int kMaxAVLTreeHeight = 32;


// #pragma mark - AVLTreeCompare


AVLTreeCompare::~AVLTreeCompare()
{
}


// #pragma mark - AVLTreeBase


AVLTreeBase::AVLTreeBase(AVLTreeCompare* compare)
        : fRoot(NULL),
          fNodeCount(0),
          fCompare(compare)
{
}


AVLTreeBase::~AVLTreeBase()
{
}


void
AVLTreeBase::MakeEmpty()
{
        fRoot = NULL;
        fNodeCount = 0;
}


AVLTreeNode*
AVLTreeBase::LeftMost(AVLTreeNode* node) const
{
        if (node) {
                while (node->left)
                        node = node->left;
        }

        return node;
}


AVLTreeNode*
AVLTreeBase::RightMost(AVLTreeNode* node) const
{
        if (node) {
                while (node->right)
                        node = node->right;
        }

        return node;
}


AVLTreeNode*
AVLTreeBase::Previous(AVLTreeNode* node) const
{
        if (node) {
                // The previous node cannot be in the right subtree.
                if (node->left) {
                        // We have a left subtree, so go to the right-most node.
                        node = node->left;
                        while (node->right)
                                node = node->right;
                } else {
                        // No left subtree: Backtrack our path and stop, where we
                        // took the right branch.
                        AVLTreeNode* previous;
                        do {
                                previous = node;
                                node = node->parent;
                        } while (node && previous == node->left);
                }
        }

        return node;
}


AVLTreeNode*
AVLTreeBase::Next(AVLTreeNode* node) const
{
        if (node) {
                // The next node cannot be in the left subtree.
                if (node->right) {
                        // We have a right subtree, so go to the left-most node.
                        node = node->right;
                        while (node->left)
                                node = node->left;
                } else {
                        // No right subtree: Backtrack our path and stop, where we
                        // took the left branch.
                        AVLTreeNode* previous;
                        do {
                                previous = node;
                                node = node->parent;
                        } while (node && previous == node->right);
                }
        }

        return node;
}


AVLTreeNode*
AVLTreeBase::Find(const void* key) const
{
        AVLTreeNode* node = fRoot;

        while (node) {
                int cmp = fCompare->CompareKeyNode(key, node);
                if (cmp == 0)
                        return node;

                if (cmp < 0)
                        node = node->left;
                else
                        node = node->right;
        }

        return NULL;
}


AVLTreeNode*
AVLTreeBase::FindClosest(const void* key, bool less) const
{
        AVLTreeNode* node = fRoot;
        AVLTreeNode* parent = NULL;

        while (node) {
                int cmp = fCompare->CompareKeyNode(key, node);
                if (cmp == 0)
                        break;

                parent = node;
                if (cmp < 0)
                        node = node->left;
                else
                        node = node->right;
        }

        // not found: try to get close
        if (!node && parent) {
                node = parent;
                int expectedCmp = (less ? 1 : -1);
                int cmp = fCompare->CompareKeyNode(key, node);
                if (cmp != expectedCmp) {
                        // The node's value is less although we were asked for a greater
                        // value, or the other way around. We need to iterate to the next
                        // node in the respective direction. If there is no node, we fail.
                        if (less)
                                return Previous(node);
                        return Next(node);
                }
        }

        return node;
}


status_t
AVLTreeBase::Insert(AVLTreeNode* nodeToInsert)
{
        int result = _Insert(nodeToInsert);
        switch (result) {
                case OK:
                case HEIGHT_CHANGED:
                        return B_OK;
                case NO_MEMORY:
                        return B_NO_MEMORY;
                case DUPLICATE:
                default:
                        return B_BAD_VALUE;
        }
}


AVLTreeNode*
AVLTreeBase::Remove(const void* key)
{
        // find node
        AVLTreeNode* node = fRoot;
        while (node) {
                int cmp = fCompare->CompareKeyNode(key, node);
                if (cmp == 0)
                        break;
                else {
                        if (cmp < 0)
                                node = node->left;
                        else
                                node = node->right;
                }
        }

        // remove it
        if (node)
                _Remove(node);

        return node;
}


bool
AVLTreeBase::Remove(AVLTreeNode* node)
{
        switch (_Remove(node)) {
                case OK:
                case HEIGHT_CHANGED:
                        return true;
                case NOT_FOUND:
                default:
                        return false;
        }
}


void
AVLTreeBase::CheckTree() const
{
        int nodeCount = 0;
        _CheckTree(NULL, fRoot, nodeCount);
        if (nodeCount != fNodeCount) {
                CHECK_FAILED("AVLTreeBase::CheckTree(): node count mismatch: %d vs %d",
                        nodeCount, fNodeCount);
        }
}


void
AVLTreeBase::_RotateRight(AVLTreeNode** nodeP)
{
        // rotate the nodes
        AVLTreeNode* node = *nodeP;
        AVLTreeNode* left = node->left;

        *nodeP = left;

        left->parent = node->parent;
        node->left = left->right;
        if (left->right)
                left->right->parent = node;
        left->right = node;
        node->parent = left;

        // adjust the balance factors
        // former pivot
        if (left->balance_factor >= 0)
                node->balance_factor++;
        else
                node->balance_factor += 1 - left->balance_factor;

        // former left
        if (node->balance_factor <= 0)
                left->balance_factor++;
        else
                left->balance_factor += node->balance_factor + 1;
}


void
AVLTreeBase::_RotateLeft(AVLTreeNode** nodeP)
{
        // rotate the nodes
        AVLTreeNode* node = *nodeP;
        AVLTreeNode* right = node->right;

        *nodeP = right;

        right->parent = node->parent;
        node->right = right->left;
        if (right->left)
                right->left->parent = node;
        right->left = node;
        node->parent = right;

        // adjust the balance factors
        // former pivot
        if (right->balance_factor <= 0)
                node->balance_factor--;
        else
                node->balance_factor -= right->balance_factor + 1;

        // former right
        if (node->balance_factor >= 0)
                right->balance_factor--;
        else
                right->balance_factor += node->balance_factor - 1;
}


int
AVLTreeBase::_BalanceInsertLeft(AVLTreeNode** node)
{
        if ((*node)->balance_factor < LEFT) {
                // tree is left heavy
                AVLTreeNode** left = &(*node)->left;
                if ((*left)->balance_factor == LEFT) {
                        // left left heavy
                        _RotateRight(node);
                } else {
                        // left right heavy
                        _RotateLeft(left);
                        _RotateRight(node);
                }
                return OK;
        }

        if ((*node)->balance_factor == BALANCED)
                return OK;

        return HEIGHT_CHANGED;
}


int
AVLTreeBase::_BalanceInsertRight(AVLTreeNode** node)
{
        if ((*node)->balance_factor > RIGHT) {
                // tree is right heavy
                AVLTreeNode** right = &(*node)->right;
                if ((*right)->balance_factor == RIGHT) {
                        // right right heavy
                        _RotateLeft(node);
                } else {
                        // right left heavy
                        _RotateRight(right);
                        _RotateLeft(node);
                }
                return OK;
        }

        if ((*node)->balance_factor == BALANCED)
                return OK;

        return HEIGHT_CHANGED;
}


int
AVLTreeBase::_Insert(AVLTreeNode* nodeToInsert)
{
        struct node_info {
                AVLTreeNode**   node;
                bool                    left;
        };

        node_info stack[kMaxAVLTreeHeight];
        node_info* top = stack;
        const node_info* const bottom = stack;
        AVLTreeNode** node = &fRoot;

        // find insertion point
        while (*node) {
                int cmp = fCompare->CompareNodes(nodeToInsert, *node);
                if (cmp == 0)   // duplicate node
                        return DUPLICATE;
                else {
                        top->node = node;
                        if (cmp < 0) {
                                top->left = true;
                                node = &(*node)->left;
                        } else {
                                top->left = false;
                                node = &(*node)->right;
                        }
                        top++;
                }
        }

        // insert node
        *node = nodeToInsert;
        nodeToInsert->left = NULL;
        nodeToInsert->right = NULL;
        nodeToInsert->balance_factor = BALANCED;
        fNodeCount++;

        if (top == bottom)
                nodeToInsert->parent = NULL;
        else
                (*node)->parent = *top[-1].node;

        // do the balancing
        int result = HEIGHT_CHANGED;
        while (result == HEIGHT_CHANGED && top != bottom) {
                top--;
                node = top->node;
                if (top->left) {
                        // left
                        (*node)->balance_factor--;
                        result = _BalanceInsertLeft(node);
                } else {
                        // right
                        (*node)->balance_factor++;
                        result = _BalanceInsertRight(node);
                }
        }

        return result;
}


int
AVLTreeBase::_BalanceRemoveLeft(AVLTreeNode** node)
{
        int result = HEIGHT_CHANGED;

        if ((*node)->balance_factor > RIGHT) {
                // tree is right heavy
                AVLTreeNode** right = &(*node)->right;
                if ((*right)->balance_factor == RIGHT) {
                        // right right heavy
                        _RotateLeft(node);
                } else if ((*right)->balance_factor == BALANCED) {
                        // right none heavy
                        _RotateLeft(node);
                        result = OK;
                } else {
                        // right left heavy
                        _RotateRight(right);
                        _RotateLeft(node);
                }
        } else if ((*node)->balance_factor == RIGHT)
                result = OK;

        return result;
}


int
AVLTreeBase::_BalanceRemoveRight(AVLTreeNode** node)
{
        int result = HEIGHT_CHANGED;

        if ((*node)->balance_factor < LEFT) {
                // tree is left heavy
                AVLTreeNode** left = &(*node)->left;
                if ((*left)->balance_factor == LEFT) {
                        // left left heavy
                        _RotateRight(node);
                } else if ((*left)->balance_factor == BALANCED) {
                        // left none heavy
                        _RotateRight(node);
                        result = OK;
                } else {
                        // left right heavy
                        _RotateLeft(left);
                        _RotateRight(node);
                }
        } else if ((*node)->balance_factor == LEFT)
                result = OK;

        return result;
}


int
AVLTreeBase::_RemoveRightMostChild(AVLTreeNode** node, AVLTreeNode** foundNode)
{
        AVLTreeNode** stack[kMaxAVLTreeHeight];
        AVLTreeNode*** top = stack;
        const AVLTreeNode* const* const* const bottom = stack;

        // find the child
        while ((*node)->right) {
                *top = node;
                top++;
                node = &(*node)->right;
        }

        // found the rightmost child: remove it
        // the found node might have a left child: replace the node with the
        // child
        *foundNode = *node;
        AVLTreeNode* left = (*node)->left;
        if (left)
                left->parent = (*node)->parent;
        *node = left;
        (*foundNode)->left = NULL;
        (*foundNode)->parent = NULL;

        // balancing
        int result = HEIGHT_CHANGED;
        while (result == HEIGHT_CHANGED && top != bottom) {
                top--;
                node = *top;
                (*node)->balance_factor--;
                result = _BalanceRemoveRight(node);
        }

        return result;
}


int
AVLTreeBase::_Remove(AVLTreeNode* node)
{
        if (!node)
                return NOT_FOUND;

        // remove node
        AVLTreeNode* parent = node->parent;
        bool isLeft = (parent && parent->left == node);
        AVLTreeNode** nodeP
                = (parent ? (isLeft ? &parent->left : &parent->right) : &fRoot);
        int result = HEIGHT_CHANGED;
        AVLTreeNode* replace = NULL;

        if (node->left && node->right) {
                // node has two children
                result = _RemoveRightMostChild(&node->left, &replace);
                replace->parent = parent;
                replace->left = node->left;
                replace->right = node->right;
                if (node->left) // check necessary, if node->left == replace
                        node->left->parent = replace;
                node->right->parent = replace;
                replace->balance_factor = node->balance_factor;
                *nodeP = replace;

                if (result == HEIGHT_CHANGED) {
                        replace->balance_factor++;
                        result = _BalanceRemoveLeft(nodeP);
                }
        } else if (node->left) {
                // node has only left child
                replace = node->left;
                replace->parent = parent;
                replace->balance_factor = node->balance_factor + 1;
                *nodeP = replace;
        } else if (node->right) {
                // node has only right child
                replace = node->right;
                replace->parent = node->parent;
                replace->balance_factor = node->balance_factor - 1;
                *nodeP = replace;
        } else {
                // node has no child
                *nodeP = NULL;
        }

        node->parent = node->left = node->right = NULL;
        fNodeCount--;

        // do the balancing
        while (result == HEIGHT_CHANGED && parent) {
                node = parent;
                parent = node->parent;
                bool oldIsLeft = isLeft;
                isLeft = (parent && parent->left == node);
                nodeP = (parent ? (isLeft ? &parent->left : &parent->right) : &fRoot);
                if (oldIsLeft) {
                        // left
                        node->balance_factor++;
                        result = _BalanceRemoveLeft(nodeP);
                } else {
                        // right
                        node->balance_factor--;
                        result = _BalanceRemoveRight(nodeP);
                }
        }

        return result;
}


int
AVLTreeBase::_CheckTree(AVLTreeNode* parent, AVLTreeNode* node,
        int& _nodeCount) const
{
        if (node == NULL) {
                _nodeCount = 0;
                return 0;
        }

        if (parent != node->parent) {
                CHECK_FAILED("AVLTreeBase::_CheckTree(): node %p parent mismatch: "
                        "%p vs %p", node, parent, node->parent);
        }

        int leftNodeCount;
        int leftDepth = _CheckTree(node, node->left, leftNodeCount);

        int rightNodeCount;
        int rightDepth = _CheckTree(node, node->right, rightNodeCount);

        int balance = rightDepth - leftDepth;
        if (balance < LEFT || balance > RIGHT) {
                CHECK_FAILED("AVLTreeBase::_CheckTree(): unbalanced subtree: %p", node);
        } else if (balance != node->balance_factor) {
                CHECK_FAILED("AVLTreeBase::_CheckTree(): subtree %p balance mismatch: "
                        "%d vs %d", node, balance, node->balance_factor);
        }

        _nodeCount = leftNodeCount + rightNodeCount + 1;
        return max_c(leftDepth, rightDepth) + 1;
}