/*
 * Time-stamp: <05/11/21 06:55:54 riffraff>
 *
 * $Id: rbtree.c,v 1.4 2005/11/21 15:34:58 riffraff Exp $
 */

#include "global.h"
#include "module.h"
#include "interpret.h"
#include "svalue.h"
#include "stralloc.h"
#include "array.h"
#include "pike_macros.h"
#include "program.h"
#include "object.h"
#include "pike_types.h"
#include "pike_error.h"
#include "builtin_functions.h"
#include "opcodes.h"
#include "module_support.h"
#include "operators.h"
#include "mapping.h"

#include "tree.h"

#undef THIS
#define THIS ((struct atable *)(Pike_fp->current_storage))
#define THIS_PROGRAM (Pike_fp->context.prog)
#define RETURN_THIS() do{ ref_push_object(Pike_fp->current_object); }while(0)

#define RB(X) (X)->data.color

extern struct anode sentinel;

void rbtree_create(INT32 args) {
  struct svalue *cb=NULL;

  get_all_args("create",args,".%*",&cb);
  low_tree_create(cb);
  THIS->type=TREE_RB;
  RB(THIS->sent)=A_BLACK;
  pop_n_elems(args);
}

struct anode *rbtree_create_node() {
  struct anode *n;
  n=create_node();
  RB(n)=A_BLACK;
  return n;
}

void real_low_rbtree_insert(struct anode *node) {
  struct anode *parent;

  RB(node)=A_RED;
  parent=node->parent;

  while (RB(parent)==A_RED) {
    struct anode *g=parent->parent;
    
    if (parent==g->left) {
      struct anode *u=g->right;
      
      if (RB(u)==A_RED) {
	RB(parent)=A_BLACK;
	RB(u)=A_BLACK;
	RB(g)=A_RED;
	node=g;
	parent=g->parent;
      } else {
	if (node==parent->right) {
	  tree_rotate_left(node,parent);
	  parent=node;
	  if (g!=parent->parent) {
	    fputs("real_low_rbtree_insert:  g!=parent->parent\n",stderr);
	  }
	}
	RB(parent)=A_BLACK;
	RB(g)=A_RED;
	tree_rotate_right(parent,g);
	break;
      }
    } else {
      struct anode *u=g->left;
      if (RB(u)==A_RED) {
	RB(parent)=A_BLACK;
	RB(u)=A_BLACK;
	RB(g)=A_RED;
	node=g;
	parent=g->parent;
      } else {
	if (node==parent->left) {
	  tree_rotate_right(node,parent);
	  parent=node;
	  if (g!=parent->parent) {
	    fputs("real_low_rbtree_insert: g!=parent->parent\n",stderr);
	  }
	}
	RB(parent)=A_BLACK;
	RB(g)=A_RED;
	tree_rotate_left(parent,g);
	break;
      }
    }
  }
  RB(THIS->root->left)=A_BLACK;
}

struct anode *low_rbtree_insert(struct svalue *key, struct svalue *val) {
  struct anode *n;

  n=low_tree_insert(key,val);
  if (n==NULL)
    return NULL;
  real_low_rbtree_insert(n);
  return n;
}

void rbtree_insert(INT32 args) {
  struct svalue *key,*val;
  struct anode *n;

  get_all_args("insert",args,"%*.%*",&key,&val);
  if (key==NULL) {
    pop_n_elems(args);
    RETURN_THIS();
    return;
  }
  
  n=low_rbtree_insert(key,val);
  pop_n_elems(args);
  if (n==NULL && THIS->abend)
    Pike_error("Key already exists.\n");
  RETURN_THIS();
}

void real_low_rbtree_remove(struct anode *node) {
  struct anode *swap,*child;
  int savecolor;


  low_tree_delete(node,&swap,&child);

  savecolor=RB(node);
  RB(node)=RB(swap);
  RB(swap)=savecolor;

  if (RB(node)==A_BLACK) {
    struct anode *parent,*sis;

    RB(THIS->root->left)=A_RED;

    while (RB(child)==A_BLACK) {
      parent=child->parent;
      if (child==parent->left) {
	sis=parent->right;
	if (sis==&sentinel) {
	  fputs("real_low_rbtree_remove:  sis==&sentinel\n",stderr);
	}
	if (RB(sis)==A_RED) {
	  RB(sis)=A_BLACK;
	  RB(parent)=A_RED;
	  tree_rotate_left(sis,parent);
	  sis=parent->right;
	  if (sis==&sentinel) {
	    fputs("real_low_rbtree_remove:  sis==&sentinel\n",stderr);
	  }
	}
	if (RB(sis->left)==A_BLACK &&
	    RB(sis->right)==A_BLACK) {
	  RB(sis)=A_RED;
	  child=parent;
	} else {
	  if (RB(sis->right)==A_BLACK) {
	    if (RB(sis)!=A_RED) {
	      fputs("real_low_rbtree_remove:  sis->left->color!=A_RED\n",stderr);
	    }
	    RB(sis->left)=A_BLACK;
	    RB(sis)=A_RED;
	    tree_rotate_right(sis->left,sis);
	    sis=parent->right;
	    if (sis==&sentinel) {
	      fputs("real_low_rbtree_remove:  sis==&sentinel\n",stderr);
	    }
	  }
	  RB(sis)=RB(parent);
	  RB(sis->right)=A_BLACK;
	  RB(parent)=A_BLACK;
	  tree_rotate_left(sis,parent);
	  break;
	}
      } else {
	if (child!=parent->right) {
	  fputs("real_low_rbtree_remove:  child!=parent->right\n",stderr);
	}
	sis=parent->left;
	if (sis==&sentinel) {
	  fputs("real_low_rbtree_remove:  sis==&sentinel\n",stderr);
	}
	if (RB(sis)==A_RED) {
	  RB(sis)=A_BLACK;
	  RB(parent)=A_RED;
	  tree_rotate_right(sis,parent);
	  sis=parent->left;
	  if (sis==&sentinel) {
	    fputs("real_low_rbtree_remove:  sis==&sentinel\n",stderr);
	  }
	}
	if (RB(sis->right)==A_BLACK &&
	    RB(sis->left)==A_BLACK) {
	  RB(sis)=A_RED;
	  child=parent;
	} else {
	  if (RB(sis->left)==A_BLACK) {
	    if (RB(sis->right)!=A_RED) {
	      fputs("real_low_rbtree_remove:  sis->right->color!=A_RED\n",stderr);
	    }
	    RB(sis->right)=A_BLACK;
	    RB(sis)=A_RED;
	    tree_rotate_left(sis->right,sis);
	    sis=parent->left;
	    if (sis==&sentinel) {
	      fputs("real_low_rbtree_remove:  sis==&sentinel\n",stderr);
	    }
	  }
	  RB(sis)=RB(parent);
	  RB(sis->left)=A_BLACK;
	  RB(parent)=A_BLACK;
	  tree_rotate_right(sis,parent);
	  break;
	}
      }
    }
    RB(child)=A_BLACK;
    RB(THIS->root->left)=A_BLACK;
  }
}

struct anode *low_rbtree_remove(struct svalue *key) {
  struct anode *n;

  n=low_tree_find(key);
  if (n==NULL)
    return NULL;
  
  real_low_rbtree_remove(n);
  return n;
}

void rbtree_remove(INT32 args) {
  struct svalue *key;
  struct anode *n;

  get_all_args("remove",args,"%*",&key);
  if (key==NULL) {
    pop_n_elems(args);
    push_undefined();
    return;
  }
  n=low_rbtree_remove(key);
  pop_n_elems(args);
  if (n) {
    push_svalue(n->key);
    push_svalue(n->val);
    f_aggregate(2);
  } else
    push_undefined();
}

extern struct program *tree_program;

struct program *create_rbtree_program() {
  start_new_program();
  low_inherit(tree_program,0,0,0,0,0);

  ADD_FUNCTION("create",rbtree_create,tFunc(tOr(tVoid,tFunc(tMixed tMixed,tInt)),tVoid),ID_STATIC);
  ADD_FUNCTION("insert",rbtree_insert,tFunc(tMixed tOr(tMixed,tVoid),tOr(tObj,tInt)),0);
  ADD_FUNCTION("remove",rbtree_remove,tFunc(tMixed,tArray),0);

  return end_program();
}