/**
 * Copyright (C) 2023-2026 Linaro Limited (or its affiliates). All rights reserved.
 * Copyright (C) 2012-2023 Arm Limited (or its affiliates).
 * All rights reserved.
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 */

#include "sectionmap.h"

#ifdef SECTIONMAP_DEBUG
#include <iostream>
#endif

#include <QtGlobal>

/**
 * Typedefs
 */
typedef SectionMapNode Node;
typedef SectionMapNode *NodePtr;

/*! Rotate node x left (see picture below). */
void SectionMapPrivateBase::rotateLeft( NodePtr x )
{
    /*   x          y
     *  / \  ->    / \
     * a   y      x   c
     *    / \    / \
     *   b   c  a   b
     *
     * BEFORE:
     * x->leftCount = a->leftCount + a->rightCount + 1;
     * x->rightCount = y->leftCount + y->rightCount + 1 = b->leftCount + b->rightCount + 1 + c->leftCount + c->rightCount + 1 + 1;
     * y->leftCount = b->leftCount + b->rightCount + 1;
     * y->rightCount = c->leftCount + c->rightCount + 1;
     *
     * AFTER:
     * x->leftCount = a->leftCount + a->rightCount + 1 (NO CHANGE)
     * x->rightCount = b->leftCount + b->rightCount + 1 (y->leftCount)
     * y->leftCount = x->leftCount + x->rightCount + 1 = a->leftCount + a->rightCount + 1 + b->leftCount + b->rightCount + 1 + 1; (+= x->leftCount + 1)
     * y->rightCount = c->leftCount + c->rightCount + 1 (NO CHANGE)
     */

    NodePtr y = x->right;
    x->right = y->left;
    if (y->left !=0)
        y->left->parent = x;
    y->parent = x->parent;
    if (x == root)
        root = y;
    else if (x == x->parent->left)
        x->parent->left = y;
    else
        x->parent->right = y;
    y->left = x;
    x->parent = y;
    x->rightCount = y->leftCount;
    y->leftCount += x->leftCount + x->count;
}


/*! Rotate node x right (see picture below). */
void SectionMapPrivateBase::rotateRight( NodePtr x  )
{
    /*     x        y
     *    / \      / \
     *   y   c -> a   x
     *  / \          / \
     * a   b        b   c
     *
     * BEFORE:
     * y->leftCount = a->leftCount + a->rightCount + 1
     * y->rightCount = b->leftCount + b->rightCount + 1
     * x->leftCount = y->leftCount + y->rightCount + 1 = a->leftCount + a->rightCount + 1 + b->leftCount + b->rightCount + 1 + 1;
     * x->rightCount = c->leftCount + c->rightCount + 1
     *
     * AFTER:
     * y->leftCount = a->leftCount + a->rightCount + 1; (NO CHANGE)
     * y->rightCount = x->leftCount + x->rightCount + 1 = b->leftCount + b->rightCount + 1 + c->leftCount + c->rightCount + 1 + 1; (+= x->rightCount + 1)
     * x->leftCount = b->leftCount + b->rightCount + 1; (y->rightCount)
     * x->rightCount = c->leftCount + c->rightCount + 1; (NO CHANGE)
     */

    NodePtr y = x->left;
    x->left = y->right;
    if (y->right != 0)
        y->right->parent = x;
    y->parent = x->parent;
    if (x == root)
        root = y;
    else if (x == x->parent->right)
        x->parent->right = y;
    else
        x->parent->left = y;
    y->right = x;
    x->parent = y;
    x->leftCount = y->rightCount;
    y->rightCount += x->rightCount + x->count;
}


/*! Rebalance the red-black tree (after an insertion).
 *  rebalance only manipulates the tree through rotateLeft and rotateRight so
 *  the left and right branch counts are maintained. */
void SectionMapPrivateBase::rebalance( NodePtr x )
{
    x->color = Node::Red;
    while ( x != root && x->parent->color == Node::Red ) {
        if ( x->parent == x->parent->parent->left ) {
            NodePtr y = x->parent->parent->right;
            if (y && y->color == Node::Red) {
                x->parent->color = Node::Black;
                y->color = Node::Black;
                x->parent->parent->color = Node::Red;
                x = x->parent->parent;
            } else {
                if (x == x->parent->right) {
                    x = x->parent;
                    rotateLeft( x );
                }
                x->parent->color = Node::Black;
                x->parent->parent->color = Node::Red;
                rotateRight (x->parent->parent);
            }
        } else {
            NodePtr y = x->parent->parent->left;
            if ( y && y->color == Node::Red ) {
                x->parent->color = Node::Black;
                y->color = Node::Black;
                x->parent->parent->color = Node::Red;
                x = x->parent->parent;
            } else {
                if (x == x->parent->left) {
                    x = x->parent;
                    rotateRight( x );
                }
                x->parent->color = Node::Black;
                x->parent->parent->color = Node::Red;
                rotateLeft( x->parent->parent );
            }
        }
    }
    root->color = Node::Black;
}

/*! Remove a node from the tree and possible promote the next node (in order,
 *  not in in the tree) in its place. */
NodePtr SectionMapPrivateBase::removeAndRebalance( NodePtr z )
{
    NodePtr y = z;
    NodePtr x;
    NodePtr x_parent;
    NodePtr tmp;

    if (y->left == 0) {
        x = y->right;
    } else {
        if (y->right == 0) {
            x = y->left;
        } else {
            y = y->right;
            while (y->left != 0)
                y = y->left;
            x = y->right;
        }
    }

    /* Promote y to be in z's place. */
    if (y != z) {
        /* This code is a little hairy. Try drawing the tree before and after
         * to understand what's going on. z is the node being removed. y is the
         * next node (in order, not in the tree) after z. x is the right branch
         * of y. */

        /* Move the left branch of z to be the left branch of y. */
        z->left->parent = y;
        y->left = z->left;
        z->rightCount -= y->leftCount;
        y->leftCount = z->leftCount;

        if (y != z->right) {
            /* Move the right branch of y (x) to be the left branch of y's
             * parent. Before this operation y is the left branch of y's
             * parent. We are replacing it with x so we need to adjust the
             * counts accordingly. */
            x_parent = y->parent;
            if (x)
                x->parent = y->parent;

            long long h = y->parent->leftCount - y->rightCount;
            tmp = y;
            while (tmp->parent != z) {
                tmp->parent->leftCount -= h;
                tmp = tmp->parent;
            }

            y->parent->left = x;
            y->parent->leftCount = y->rightCount;

            /* Move the right branch of z to be the left branch of y. z's
             * branches are completely unlinked now. y used to be in this
             * branch but has been promoted to the parent so we need to
             * remove its count from the total for the branch. */
            y->right = z->right;
            z->right->parent = y;
            y->rightCount = z->rightCount - y->count;
        } else {
            x_parent = y;
        }

        /* Reparent y so it becomes z->parent's child instead of z. */
        if (root == z)
            root = y;
        else if (z->parent->left == z)
            z->parent->left = y;
        else
            z->parent->right = y;
        y->parent = z->parent;

        // Swap the colors
        Node::Color c = y->color;
        y->color = z->color;
        z->color = c;
        y = z;
    } else {
        /* Reparent x so it becomes z->parent's child instead of z. */
        x_parent = y->parent;
        if (x)
            x->parent = y->parent;
        if (root == z)
            root = x;
        else if (z->parent->left == z)
            z->parent->left = x;
        else
            z->parent->right = x;
    }

    /* Rebalance the tree. Uses rotateLeft and rotateRight rather than
     * manipulating the tree directly so the counts are preserved. */
    if (y->color != Node::Red) {
        while (x != root && (x == 0 || x->color == Node::Black)) {
            if (x == x_parent->left) {
                NodePtr w = x_parent->right;
                if (w->color == Node::Red) {
                    w->color = Node::Black;
                    x_parent->color = Node::Red;
                    rotateLeft(x_parent);
                    w = x_parent->right;
                }
                if ((w->left == 0 || w->left->color == Node::Black) &&
                    (w->right == 0 || w->right->color == Node::Black)) {
                    w->color = Node::Red;
                    x = x_parent;
                    x_parent = x_parent->parent;
                } else {
                    if (w->right == 0 || w->right->color == Node::Black) {
                        if (w->left)
                            w->left->color = Node::Black;
                        w->color = Node::Red;
                        rotateRight(w);
                        w = x_parent->right;
                    }
                    w->color = x_parent->color;
                    x_parent->color = Node::Black;
                    if (w->right)
                        w->right->color = Node::Black;
                    rotateLeft(x_parent);
                    break;
                }
            } else {
                NodePtr w = x_parent->left;
                if (w->color == Node::Red) {
                    w->color = Node::Black;
                    x_parent->color = Node::Red;
                    rotateRight(x_parent);
                    w = x_parent->left;
                }
                if ((w->right == 0 || w->right->color == Node::Black) &&
                    (w->left == 0 || w->left->color == Node::Black)) {
                    w->color = Node::Red;
                    x = x_parent;
                    x_parent = x_parent->parent;
                } else {
                    if (w->left == 0 || w->left->color == Node::Black) {
                        if (w->right)
                            w->right->color = Node::Black;
                        w->color = Node::Red;
                        rotateLeft(w);
                        w = x_parent->left;
                    }
                    w->color = x_parent->color;
                    x_parent->color = Node::Black;
                    if (w->left)
                        w->left->color = Node::Black;
                    rotateRight(x_parent);
                    break;
                }
            }
        }
        if (x)
            x->color = Node::Black;
    }

    return y;
}


#ifdef SECTIONMAP_DEBUG

/*! Print the tree in order. */
void SectionMapPrivate::inorder( NodePtr x, int level )
{
    if ( !x ) {
        x = root;
        if ( !x )
            return;
    }
    if ( x->left )
        inorder( x->left, level + 1 );
    std::cout << level << " Key=" << key(x) << " Count=" << (x)->count << " LeftHidden=" << (x)->leftCount << " RightHidden=" << (x)->rightCount << std::endl;
    if ( x->right )
        inorder( x->right, level + 1 );
}

/*! Print the tree. */
void SectionMapPrivate::treeorder( NodePtr x, int level )
{
    if ( !x ) {
        x = root;
        if ( !x )
            return;
    }
    std::cout << level;
    for (int i=0;i<level;++i)
        std::cout << "  ";
    std::cout << " Key=" << key(x) << " Count=" << (x)->count << " LeftHidden=" << (x)->leftCount << " RightHidden=" << (x)->rightCount << std::endl;

    if ( x->left )
        treeorder( x->left, level + 1 );
    if ( x->right )
        treeorder( x->right, level + 1 );
}

/*! Fixup the branch totals. */
void SectionMapPrivate::fixup( NodePtr x )
{
    if ( !x ) {
        x = root;
        if ( !x )
            return;
    }
    x->leftCount = hiddenCount(x->left);
    x->rightCount = hiddenCount(x->right);

    if ( x->left )
        fixup( x->left );
    if ( x->right )
        fixup( x->right );
}

/*! Return the total for the given node and it's children. */
long long SectionMapPrivate::hiddenCount( NodePtr x )
{
    if ( !x )
        return 0;

    long long count = x->count;
    if ( x->left )
        count += hiddenCount( x->left );
    if ( x->right )
        count += hiddenCount( x->right );
    return count;
}

#endif


void SectionMapPrivate::clear()
{
    clear(root);
}


void SectionMapPrivate::clear(  NodePtr p )
{
    if (p == root)
        root = 0;
    while ( p != 0 ) {
        clear( p->right );
        NodePtr y = p->left;
        delete p;
        p = y;
    }
}

/*! Adjust the totals for child's parents by delta. */
void SectionMapPrivate::fixParents( NodePtr child, long long delta )
{
    while (child != root ) {
        NodePtr parent = child->parent;
        if (parent->left == child)
            parent->leftCount += delta;
        else
            parent->rightCount += delta;
        child = parent;
   }
}

/*! Insert a single node k into the tree. */
NodePtr SectionMapPrivate::insertSingle( const long long k, const long long count )
{
    NodePtr y;

    // Does the range to be inserted come immediately before an existing
    // range?
    if ( (y = find(k + count)) && k < y->key && y->key <= k + count ) {
        // Adjust the existing range to accomodate the new one.
        long long delta = y->key - k;
        y->key = k;
        // Can we consolidate this range with another one immediately before
        // it?
        NodePtr z = find(k - 1);
        if (z) {
            Q_ASSERT(y != z);
            y->key = z->key;
            delta += z->count;
            fixParents(z, -z->count);
            delete removeAndRebalance( z );
        }
        y->count += delta;
        fixParents(y, delta);
    // Does the range to be inserted come immediately after an existing
    // range?
    } else if ( (y = find(k - 1)) && k > y->key && y->key + y->count >= k ) {
        long long delta = (k + count) - (y->key + y->count);
        NodePtr z = find(k + count);
        // Can we consolidate this range with another one immediately before
        // it?
        if (z) {
            Q_ASSERT(y != z);
            delta += z->count;
            fixParents(z, -z->count);
            delete removeAndRebalance( z );
        }
        y->count += delta;
        fixParents(y, delta);
    } else {
        y = find(k, false);
        y = insert(y, k, count );
    }

    // Return the node covering the range just inserted.
    return y;
}


/*! Creates a new node covering the range k => k + count and inserts it as the
 *  child of y. If y is 0 the node becomes the new root. */
NodePtr SectionMapPrivate::insert( NodePtr y, const long long k, const long long count )
{
    NodePtr z = new Node( k );
    if (y == 0) {
        root = z;
    } else if (key(y) > k ) {
        y->left = z;
    } else {
        y->right = z;
    }
    z->parent = y;
    z->left = 0;
    z->right = 0;
    z->count = count;
    // Adjust the totals for the parents of z to account for the inserted node.
    fixParents(z, count);

    rebalance( z );
    return z;
}


void SectionMap::hideSection(long long s)
{
    if (!find(s))
        insertSingle(s);
}


void SectionMap::hideSections(long long s, long long e)
{
    showSections(s, e);
    insertSingle(s, e - s + 1);
}


/*! Finds the node that contains k. If exact is false and there is not an exact
 *  match then it will return an adjacent node if there is one, else the next or
 *  previous node in the order. */
NodePtr SectionMapPrivate::find( const long long k, bool exact ) const
{
    NodePtr y = 0;
    NodePtr x = root;
    NodePtr z = 0;
    while ( x != 0 ) {
        y = x;
        // Remember adjacent node.
        if (key(x) > k) {
            if (key(x) == k + 1)
                z = x;
            x = x->left;
        // Is it an exact match?
        } else if ( k >= key(x) && k < key(x) + x->count ) {
            return y;
        } else {
            // Remember adjacent node.
            if (key(x) + x->count == k)
                z = x;
            x = x->right;
        }
    }
    if (!exact) {
        // Prefer an adjacent node if we found one.
        if (z)
            y = z;
        return y;
    }
    return 0;
}

void SectionMap::showSection(long long s)
{
    NodePtr y = find( s );
    if (y) {
        // Is this section at the start of the node?
        if (s == key(y)) {
            fixParents(y, -1);
            if (y->count > 1) {
                ++y->key;
                --y->count;
            } else {
                delete removeAndRebalance( y );
            }
        // Is this section at the end of the node?
        } else if (s == key(y) + y->count - 1) {
            --y->count;
            fixParents(y, -1);
        // If it's not at the start or end we need to split the node.
        } else {
            long long oldcount = y->count;
            y->count = s - key(y);
            long long delta = y->count - oldcount;
            // Fixup the totals for the parents of y to account for the removed
            // section.
            fixParents(y, delta);
            // Insert a new node after the section being removed.
            NodePtr z = find(s + 1, false);
            insert(z, s + 1, key(y) + oldcount - s - 1);
        }
    }
}

void SectionMap::showSections(long long s, long long e)
{
    NodePtr y = find( s, false);
    NodePtr z = find( e, false);
    if ( y == z &&
         (!y || y->key + y->count <= s || y->key > e) ) {
        return;
    }
    // A less than ideal implementation!
    for (long long i=s;i<=e;++i)
        showSection(i);
}


bool SectionMap::isSectionHidden(long long s) const
{
    return find( s ) != 0;
}


bool SectionMap::isSectionHidden(long long s, long long e) const
{
    NodePtr x = find( s );
    if (x == 0)
        return false;
    return e < x->key + x->count;
}


/*! Returns the physical position of section s */
long long SectionMap::sectionPos(long long s) const
{
    NodePtr x = root; // Root node.

    long long leftCount = 0;
    while ( x != 0 ) {
        if ( key(x) > s ) {
            x = x->left;
        } else {
            leftCount += x->leftCount;
            if ( s >= key(x) && s < key(x) + x->count ) {
                leftCount += s - key(x);
                break;
            }
            leftCount += x->count;
            x = x->right;
        }
    }
    return s - leftCount;
}


/*! Returns the logical position of section s */
long long SectionMap::sectionAt(long long s) const
{
    NodePtr x = root; // Root node.

    long long leftCount = 0;
    while ( x != 0 ) {
        // Calculate the physical position of this node.
        long long at = key(x) - leftCount - x->leftCount;
        if ( at > s ) {
            x = x->left;
        } else {
            leftCount += x->leftCount + x->count;
            x = x->right;
        }
    }

    return s + leftCount;
}
