File content as of revision 1:cfe8984a32b4:

/*********************************************************************
PicoTCP. Copyright (c) 2012 TASS Belgium NV. Some rights reserved.
See LICENSE and COPYING for usage.

Author: Andrei Carp <andrei.carp@tass.be>
*********************************************************************/

#include "pico_tree.h"
#include "pico_config.h"

#define RED     0
#define BLACK 1

// By default the null leafs are black
struct pico_tree_node LEAF = {
  NULL, // key
  &LEAF, &LEAF, &LEAF, // parent, left,right
  BLACK // color
};

#define IS_LEAF(x) (x == &LEAF)
#define IS_NOT_LEAF(x) (x != &LEAF)
#define INIT_LEAF (&LEAF)

#define AM_I_LEFT_CHILD(x) (x == x->parent->leftChild)
#define AM_I_RIGHT_CHILD(x) (x == x->parent->rightChild)

#define PARENT(x) (x->parent)
#define GRANPA(x) (x->parent->parent)

/*
 * Local Functions
 */
static struct pico_tree_node * create_node(void *key);
static void rotateToLeft(struct pico_tree* tree, struct pico_tree_node* node);
static void rotateToRight(struct pico_tree* root, struct pico_tree_node* node);
static void fix_insert_collisions(struct pico_tree* tree, struct pico_tree_node* node);
static void fix_delete_collisions(struct pico_tree* tree, struct pico_tree_node * node);
static void switchNodes(struct pico_tree* tree, struct pico_tree_node* nodeA, struct pico_tree_node* nodeB);

/*
 * Exported functions
 */

struct pico_tree_node * pico_tree_firstNode(struct pico_tree_node * node)
{
    while(IS_NOT_LEAF(node->leftChild))
        node = node->leftChild;

    return node;
}

struct pico_tree_node * pico_tree_lastNode(struct pico_tree_node * node)
{
    while(IS_NOT_LEAF(node->rightChild))
        node = node->rightChild;

    return node;
}

struct pico_tree_node * pico_tree_next(struct pico_tree_node * node)
{
    if(IS_NOT_LEAF(node->rightChild))
    {
        node = node->rightChild;
        while(IS_NOT_LEAF(node->leftChild))
            node = node->leftChild;
    }
    else
    {
        if (IS_NOT_LEAF(node->parent) &&  AM_I_LEFT_CHILD(node))
                  node = node->parent;
        else {
            while (IS_NOT_LEAF(node->parent) &&    AM_I_RIGHT_CHILD(node))
                node = node->parent;

            node = node->parent;
        }
    }

    return node;
}

struct pico_tree_node * pico_tree_prev(struct pico_tree_node * node)
{
  if (IS_NOT_LEAF(node->leftChild)) {
      node = node->leftChild;
      while (IS_NOT_LEAF(node->rightChild))
          node = node->rightChild;
  } else {
      if (IS_NOT_LEAF(node->parent) && AM_I_RIGHT_CHILD(node))
          node = node->parent;
      else {
          while (IS_NOT_LEAF(node) &&    AM_I_LEFT_CHILD(node))
              node = node->parent;

          node = node->parent;
      }
  }

  return node;
}

void * pico_tree_insert(struct pico_tree* tree, void * key){

    struct pico_tree_node * last_node = INIT_LEAF;
  struct pico_tree_node * temp = tree->root;
  struct pico_tree_node * insert;
  void * LocalKey;
  int result = 0;

  LocalKey = (IS_NOT_LEAF(tree->root) ? pico_tree_findKey(tree,key) : NULL);
    // if node already in, bail out
  if(LocalKey)
      return LocalKey;
  else
      insert = create_node(key);

  // search for the place to insert the new node
  while(IS_NOT_LEAF(temp))
  {
        last_node = temp;
        result = tree->compare(insert->keyValue,temp->keyValue);

        temp = (result < 0 ? temp->leftChild : temp->rightChild);
  }

  // make the needed connections
  insert->parent = last_node;

  if(IS_LEAF(last_node))
    tree->root = insert;
  else{
      result = tree->compare(insert->keyValue,last_node->keyValue);
    if(result < 0)
      last_node->leftChild = insert;
    else
      last_node->rightChild = insert;
  }

  // fix colour issues
  fix_insert_collisions(tree, insert);

  return NULL;
}

struct pico_tree_node * pico_tree_findNode(struct pico_tree * tree, void * key)
{
        struct pico_tree_node *found;
        found = tree->root;

      while(IS_NOT_LEAF(found))
      {
          int result;
          result = tree->compare(found->keyValue, key);
          if(result == 0)
             return found;
          else if(result < 0)
           found = found->rightChild;
         else
           found = found->leftChild;
       }

      return NULL;
}

void * pico_tree_findKey(struct pico_tree * tree, void * key)
{
  struct pico_tree_node *found;
    found = tree->root;

  while(IS_NOT_LEAF(found))
  {
      int result;

      result = tree->compare(found->keyValue, key);
      if(result == 0)
      {
         return found->keyValue;
      }
      else if(result < 0)
       found = found->rightChild;
     else
       found = found->leftChild;

   }

  return NULL;
}

void * pico_tree_first(struct pico_tree * tree)
{
    return pico_tree_firstNode(tree->root)->keyValue;
}

void * pico_tree_last(struct pico_tree * tree)
{
    return pico_tree_lastNode(tree->root)->keyValue;
}

void * pico_tree_delete(struct pico_tree * tree, void * key){

    uint8_t nodeColor; // keeps the color of the node to be deleted
  void * lkey; // keeps a copy of the key which will be removed
    struct pico_tree_node * delete; // keeps a copy of the node to be extracted
    struct pico_tree_node * temp; // temporary
    struct pico_tree_node * ltemp; // used for switching nodes, as a copy

    delete = pico_tree_findNode(tree, key);
    ltemp = delete;

  // this key isn't in the tree, bail out
  if(!delete)
    return NULL;

  lkey = delete->keyValue;
  nodeColor = delete->color;

  if(IS_LEAF(delete->leftChild))
  {
    temp = ltemp->rightChild;
    switchNodes(tree, ltemp, ltemp->rightChild);
  }
  else
    if(IS_LEAF(delete->rightChild))
    {
        struct pico_tree_node * ltemp = delete;
      temp = ltemp->leftChild;
      switchNodes(tree, ltemp, ltemp->leftChild);
    }
    else{
        struct pico_tree_node * min;
        min = pico_tree_firstNode(delete->rightChild);
      nodeColor = min->color;

       temp = min->rightChild;
       if(min->parent == ltemp && IS_NOT_LEAF(temp))
      temp->parent = min;
       else{
           switchNodes(tree, min, min->rightChild);
           min->rightChild = ltemp->rightChild;
           if(IS_NOT_LEAF(min->rightChild)) min->rightChild->parent = min;
       }
       switchNodes(tree, ltemp, min);
       min->leftChild = ltemp->leftChild;
       if(IS_NOT_LEAF(min->leftChild)) min->leftChild->parent = min;
       min->color = ltemp->color;
    }

  // deleted node is black, this will mess up the black path property
  if(nodeColor == BLACK)
    fix_delete_collisions(tree, temp);

  pico_free(delete);

  return lkey;
}

int pico_tree_empty(struct pico_tree * tree)
{
    return (!tree->root || IS_LEAF(tree->root));
}

/*
 * Private functions
 */
static void rotateToLeft(struct pico_tree* tree, struct pico_tree_node* node)
{
  struct pico_tree_node* temp;

  temp = node->rightChild;

  if(temp == &LEAF) return;

  node->rightChild = temp->leftChild;

  if(IS_NOT_LEAF(temp->leftChild))
    temp->leftChild->parent = node;

  temp->parent = node->parent;

  if(IS_LEAF(node->parent))
    tree->root = temp;
  else
    if(node == node->parent->leftChild)
        node->parent->leftChild = temp;
    else
        node->parent->rightChild = temp;

  temp->leftChild = node;
  node->parent = temp;
}


static void rotateToRight(struct pico_tree * tree, struct pico_tree_node * node)
{
  struct pico_tree_node* temp;

  temp = node->leftChild;
  node->leftChild = temp->rightChild;

  if(temp == &LEAF) return;

  if(IS_NOT_LEAF(temp->rightChild))
    temp->rightChild->parent = node;

  temp->parent = node->parent;

  if(IS_LEAF(node->parent))
    tree->root = temp;
  else
    if(node == node->parent->rightChild)
      node->parent->rightChild = temp;
    else
      node->parent->leftChild = temp;

  temp->rightChild = node;
  node->parent = temp;
  return;
}

static struct pico_tree_node * create_node(void* key)
{
  struct pico_tree_node *temp;
  temp = (struct pico_tree_node *)pico_zalloc(sizeof(struct pico_tree_node));

  if(!temp)
      return NULL;

  temp->keyValue = key;
  temp->parent = &LEAF;
  temp->leftChild = &LEAF;
  temp->rightChild = &LEAF;
  // by default every new node is red
  temp->color = RED;

  return temp;
}

/*
 * This function fixes the possible collisions in the tree.
 * Eg. if a node is red his children must be black !
 */
static void fix_insert_collisions(struct pico_tree* tree, struct pico_tree_node* node)
{
  struct pico_tree_node* temp;

  while(node->parent->color == RED)
  {
      if(AM_I_RIGHT_CHILD(node->parent))
     {
            temp = GRANPA(node)->leftChild;
            if(temp->color == RED){
                node->parent->color = BLACK;
                temp->color = BLACK;
                GRANPA(node)->color = RED;
                node = GRANPA(node);
            }
            else if(temp->color == BLACK){
                if(node == node->parent->leftChild){
                node = node->parent;
                rotateToRight(tree, node);
            }
            node->parent->color = BLACK;
            GRANPA(node)->color = RED;
            rotateToLeft(tree, GRANPA(node));
            }
        }
      else if(AM_I_LEFT_CHILD(node->parent))
    {
      temp = GRANPA(node)->rightChild;
      if(temp->color == RED){
                node->parent->color = BLACK;
                temp->color = BLACK;
                GRANPA(node)->color = RED;
                node = GRANPA(node);
      }
      else if(temp->color == BLACK){
                if(AM_I_RIGHT_CHILD(node)){
                    node = node->parent;
                    rotateToLeft(tree, node);
                }
                node->parent->color = BLACK;
                GRANPA(node)->color = RED;
                rotateToRight(tree, GRANPA(node));
          }
    }
  }

  // make sure that the root of the tree stays black
  tree->root->color = BLACK;
}

static void switchNodes(struct pico_tree* tree, struct pico_tree_node* nodeA, struct pico_tree_node* nodeB)
{

    if(IS_LEAF(nodeA->parent))
    tree->root = nodeB;
  else
  if(IS_NOT_LEAF(nodeA))
  {    
    if(AM_I_LEFT_CHILD(nodeA))
      nodeA->parent->leftChild = nodeB;
    else
      nodeA->parent->rightChild = nodeB;
   }

  if(IS_NOT_LEAF(nodeB))  nodeB->parent = nodeA->parent;

}

/*
 * This function fixes the possible collisions in the tree.
 * Eg. if a node is red his children must be black !
 * In this case the function fixes the constant black path property.
 */
static void fix_delete_collisions(struct pico_tree* tree, struct pico_tree_node * node)
{
  struct pico_tree_node* temp;

  while( node != tree->root && node->color == BLACK && IS_NOT_LEAF(node))
  {
      if(AM_I_LEFT_CHILD(node)){

      temp = node->parent->rightChild;
      if(temp->color == RED)
      {
                temp->color = BLACK;
                node->parent->color = RED;
                rotateToLeft(tree, node->parent);
                temp = node->parent->rightChild;
      }
      if(temp->leftChild->color == BLACK && temp->rightChild->color == BLACK)
      {
                temp->color = RED;
                node = node->parent;
      }
      else
      {
                if(temp->rightChild->color == BLACK)
                {
                    temp->leftChild->color = BLACK;
                    temp->color = RED;
                    rotateToRight(tree, temp);
                    temp = temp->parent->rightChild;
                }
                temp->color = node->parent->color;
                node->parent->color = BLACK;
                temp->rightChild->color = BLACK;
                rotateToLeft(tree, node->parent);
                node = tree->root;
      }
    }
    else{
      temp = node->parent->leftChild;
      if(temp->color == RED)
      {
                temp->color = BLACK;
                node->parent->color = RED;
                rotateToRight(tree, node->parent);
                temp = node->parent->leftChild;
      }
      if(temp->rightChild->color == BLACK && temp->leftChild->color == BLACK)
      {
                temp->color = RED;
                node = node->parent;
      }
      else{
                if(temp->leftChild->color == BLACK)
                {
                    temp->rightChild->color = BLACK;
                    temp->color = RED;
                    rotateToLeft(tree, temp);
                    temp = temp->parent->leftChild;
                }
            temp->color = node->parent->color;
            node->parent->color = BLACK;
            temp->leftChild->color = BLACK;
            rotateToRight(tree, node->parent);
            node = tree->root;
      }
    }
  }

  node->color = BLACK;
}