/* Red-black tree support for Ping903
   Copyright (C) 2020-2023 Sergey Poznyakoff
  
   Ping903 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, or (at your option)
   any later version.
  
   Ping903 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 Ping903.  If not, see <http://www.gnu.org/licenses/>.
*/
#include <config.h>
#include <stdlib.h>
#include <string.h>
#include "ping903.h"

typedef struct rbt_node RBT_NODE;

typedef enum { RED, BLACK } RBT_COLOR;

struct rbt_node {
    RBT_NODE *left, *right, *parent;
    RBT_COLOR color;
    HOSTPING *host;
};

struct rbt_tree {
    int (*cmp)(HOSTPING*, HOSTPING*);
    RBT_NODE *root;	/* Root of the tree */
};

RBT_TREE *
rbt_tree_create(int (*cmp)(HOSTPING*, HOSTPING*))
{
    RBT_TREE *tree = malloc(sizeof(*tree));
    if (tree) {
	tree->cmp = cmp;
	tree->root = NULL;
    }
    return tree;
}

/*
 * Red-black tree properties:
 *  1. Each node is either red or black.
 *  2. The root node is black.
 *  3. All leaves are black and contain no data.
 *  4. Every red node has two children, and both are black. IOW, the
 *     parent of every red node is black.
 *  5. All paths from any given node to its leaf nodes contain the same
 *     number of black nodes. 
 */

/*
 * Auxiliary functions for accessing nodes. 
 */

/*
 * Return the grandparent node of N. Prerequisite: N may not be root. 
 */
static inline RBT_NODE *
grandparent(RBT_NODE *n)
{
    return n->parent->parent;
}

/*
 * Return the sibling node of N. Prerequisite: N may not be root. 
 */
static inline RBT_NODE *
sibling(RBT_NODE *n)
{
    return (n == n->parent->left) ? n->parent->right : n->parent->left;
}

/*
 * Return the uncle node of N. Prerequisite: N must be at least 2 nodes
 * away from root. 
 */
static inline RBT_NODE *
uncle(RBT_NODE *n)
{
    return sibling(n->parent);
}

/*
 * Returns the color of the node N. Empty leaves are represented by NULL,
 * therefore NULL is assumed to be black (see property 3). 
 */
static inline RBT_COLOR
node_color(RBT_NODE *n)
{
    return n == NULL ? BLACK : n->color;
}

/*
 * Replace the OLDN with NEWN. Does not modify OLDN. 
 */
static void
replace_node(RBT_TREE *tree, RBT_NODE *oldn, RBT_NODE *newn)
{
    if (oldn->parent == NULL)
        tree->root = newn;
    else if (oldn == oldn->parent->left)
        oldn->parent->left = newn;
    else
        oldn->parent->right = newn;

    if (newn != NULL)
        newn->parent = oldn->parent;
}

/*
 * Rotate the TREE left over the node N. 
 */
static void
rotate_left(RBT_TREE *tree, RBT_NODE *n)
{
    RBT_NODE *right = n->right;
    replace_node(tree, n, right);
    n->right = right->left;
    if (right->left != NULL)
        right->left->parent = n;
    right->left = n;
    n->parent = right;
}

/*
 * Rotate the TREE right over the node N. 
 */
static void
rotate_right(RBT_TREE *tree, RBT_NODE *n)
{
    RBT_NODE *left = n->left;
    replace_node(tree, n, left);
    n->left = left->right;
    if (left->right != NULL)
        left->right->parent = n;
    left->right = n;
    n->parent = left;
}

static void
rbt_delete_fixup(RBT_TREE *tree, RBT_NODE *n)
{
    while (1) {
        if (n->parent == NULL) {
            /*
             * If N has become the root node, deletion resulted in
             * removing one black node (prior root) from every path, so
             * all properties still hold. 
             */
            return;
        } else {
            /*
             * If N has a red sibling, change the colors of the parent and 
             * sibling and rotate about the parent. Thus, the sibling
             * becomes grandparent and we can proceed to the next case. 
             */
            if (node_color(sibling(n)) == RED) {
                n->parent->color = RED;
                sibling(n)->color = BLACK;
                if (n == n->parent->left)
                    rotate_left(tree, n->parent);
                else
                    rotate_right(tree, n->parent);
            }

            /*
             * If the parent, sibling and nephews are all black, paint the 
             * sibling red.  This means one black node was removed from
             * all paths passing through the parent, so we recurse to the
             * beginning of the loop with parent as the argument to
             * restore the properties.  This is the only branch that
             * loops. 
             */
            if (node_color(n->parent) == BLACK
                && node_color(sibling(n)) == BLACK
                && node_color(sibling(n)->left) == BLACK
                && node_color(sibling(n)->right) == BLACK) {
                sibling(n)->color = RED;
                n = n->parent;
                continue;
            } else {
                /*
                 * If the sibling and nephews are black but the parent is
                 * red, swap the colors of the sibling and parent.  The
                 * properties are then restored. 
                 */
                if (node_color(n->parent) == RED
                    && node_color(sibling(n)) == BLACK
                    && node_color(sibling(n)->left) == BLACK
		    && node_color(sibling(n)->right) == BLACK) {
                    sibling(n)->color = RED;
                    n->parent->color = BLACK;
                } else {
                    /*
                     * N is the left child of its parent, its sibling is
                     * black, and the sibling's right child is black. Swap 
                     * the colors of the sibling and its left sibling and
                     * rotate right over the sibling. 
                     */
                    if (n == n->parent->left
                        && node_color(sibling(n)) == BLACK
                        && node_color(sibling(n)->left) == RED
                        && node_color(sibling(n)->right) == BLACK) {
                        sibling(n)->color = RED;
                        sibling(n)->left->color = BLACK;
                        rotate_right(tree, sibling(n));
                    } else if (n == n->parent->right
                               && node_color(sibling(n)) == BLACK
                               && node_color(sibling(n)->right) == RED
                               && node_color(sibling(n)->left) == BLACK) {
                        /*
                         * The mirror case is handled similarly. 
                         */
                        sibling(n)->color = RED;
                        sibling(n)->right->color = BLACK;
                        rotate_left(tree, sibling(n));
                    }
                    /*
                     * N is the left child of its parent, its sibling is
                     * black and the sibling's right child is red.  Swap
                     * the colors of the parent and sibling, paint the
                     * sibling's right child black and rotate left at the
                     * parent.  Similarly for the mirror case.  This
                     * achieves the following:
                     * 
                     *  . A black node is added to all paths passing
                     *    through N;
		     *  . A black node is removed from all paths 
                     *    through the sibling's red child.
		     *  . The latter is painted black which restores the
		     *    missing black node in all paths through the
		     *    sibling's red child.
                     * 
                     * Another sibling's child becomes a child of N's
                     * parent during the rotation and is therefore not
                     * affected. 
                     */
                    sibling(n)->color = node_color(n->parent);
                    n->parent->color = BLACK;
                    if (n == n->parent->left) {
                        sibling(n)->right->color = BLACK;
                        rotate_left(tree, n->parent);
                    } else {
                        sibling(n)->left->color = BLACK;
                        rotate_right(tree, n->parent);
                    }
                }
            }
        }
        break;
    }
}

/*
 * Remove N from the TREE. 
 */
void
rbt_delete_node(RBT_TREE *tree, RBT_NODE *n)
{
    RBT_NODE *child;

    /*
     * If N has both left and right children, reduce the problem to the
     * node with only one child.  To do so, find the in-order predecessor
     * of N, copy its value (elem) to N and then delete the predecessor. 
     */
    if (n->left != NULL && n->right != NULL) {
        RBT_NODE *p;
        for (p = n->left; p->right; p = p->right)
	    ;
        n->host = p->host;
        n = p;
    }

    /*
     * N has only one child. Select it. 
     */
    child = n->left ? n->left : n->right;
    if (node_color(n) == BLACK) {
        n->color = node_color(child);
        rbt_delete_fixup(tree, n);
    }
    replace_node(tree, n, child);
    if (n->parent == NULL && child != NULL)  /* root should be black */
        child->color = BLACK;

    free(n);
}

static void
rbt_insert_fixup(RBT_TREE *tree, RBT_NODE *n)
{
    while (1) {
        if (n->parent == NULL) {
            /*
             * Node was inserted at the root of the tree. The root node
             * must be black (property 2).  Changing its color to black
             * would add one black node to every path, which means the
             * property 5 would remain satisfied.  So we simply paint the
             * node black. 
             */
            n->color = BLACK;
        } else if (node_color(n->parent) == BLACK) {
            /*
             * The node has black parent. All properties are satisfied.
             * There's no need to change anything. 
             */
            return;
        } else if (node_color(uncle(n)) == RED) {
            /*
             * The uncle node is red. Repaint the parent and uncle black
             * and the grandparent red.  This would satisfy 4.  However,
             * if the grandparent is root, this would violate the property 
             * 2.  So we repaint the grandparent by re-entering the fixup
             * loop with grandparent as the node. This is the only branch
             * that loops. 
             */
            n->parent->color = BLACK;
            uncle(n)->color = BLACK;
            n = grandparent(n);
            n->color = RED;
            continue;
        } else {
            /*
             * The new node is the right child of its parent and the
             * parent is the left child of the grandparent.  Rotate left
             * about the parent. Mirror case: The new node is the left
             * child of its parent and the parent is the right child of
             * the grandparent.  Rotate right about the parent.  This
             * fixes the properties for the rbt_insert_5. 
             */
            if (n == n->parent->right && n->parent == grandparent(n)->left) {
                rotate_left(tree, n->parent);
                n = n->left;
            } else if (n == n->parent->left
                       && n->parent == grandparent(n)->right) {
                rotate_right(tree, n->parent);
                n = n->right;
            }

            /*
             * The new node is the left child of its parent and the parent 
             * is the left child of the grandparent. Rotate right about
             * the grandparent. Mirror case: The new node is the right
             * child of its parent and the parent is the right child of
             * the grandparent. Rotate left. 
             */
            n->parent->color = BLACK;
            grandparent(n)->color = RED;
            if (n == n->parent->left && n->parent == grandparent(n)->left) {
                rotate_right(tree, grandparent(n));
            } else {
                rotate_left(tree, grandparent(n));
            }
        }
        break;
    }
}

RBT_LOOKUP_RESULT
rbt_lookup_or_insert_node(RBT_TREE *tree, HOSTPING *key, int insert,
			  RBT_NODE **retval)
{
    RBT_LOOKUP_RESULT res;
    RBT_NODE *node, *parent = NULL;
    RBT_NODE **nodeptr;

    nodeptr = &tree->root;
    while ((node = *nodeptr) != NULL) {
        int rc = tree->cmp(key, node->host);
        if (rc == 0)
            break;
        parent = node;
        if (rc < 0)
            nodeptr = &node->left;
        else
            nodeptr = &node->right;
    }

    if (node) {
        res = RBT_LOOKUP_SUCCESS;
        *retval = node;
    } else {
	res = RBT_LOOKUP_NOENT;
	if (insert) {
	    node = malloc(sizeof(*node));
	    if (!node)
		return RBT_LOOKUP_FAILURE;
	    memset(node, 0, sizeof(*node));
	    *nodeptr = node;
	    node->parent = parent;
	    rbt_insert_fixup(tree, node);
	    *retval = node;
	}
    }
    return res;
}

HOSTPING *
rbt_lookup(RBT_TREE *tree, HOSTPING *key)
{
    RBT_NODE *node;
    switch (rbt_lookup_or_insert_node(tree, key, 0, &node)) {
    case RBT_LOOKUP_SUCCESS:
	return node->host;

    case RBT_LOOKUP_NOENT:
	return NULL;

    default:
	/* Should not happen, since no allocation takes place. */
	abort();
    }
}

RBT_LOOKUP_RESULT
rbt_insert(RBT_TREE *tree, HOSTPING *host)
{
    RBT_NODE *node;
    RBT_LOOKUP_RESULT res;
    res = rbt_lookup_or_insert_node(tree, host, 1, &node);
    switch (res) {
    case RBT_LOOKUP_SUCCESS:
    case RBT_LOOKUP_FAILURE:
	break;

    case RBT_LOOKUP_NOENT:
	node->host = host;
	break;

    default:
	/* Should not happen */
	abort();
    }
    return res;
}

void
rbt_delete(RBT_TREE *tree, HOSTPING *host)
{
    RBT_NODE *node;
    if (rbt_lookup_or_insert_node(tree, host, 0, &node) == RBT_LOOKUP_SUCCESS)
	rbt_delete_node(tree, node);
}    

/* Local Variables: */
/* mode: c */
/* c-basic-offset: 4 */
/* End: */