/* Red/black trees
 * 15-122 Principles of Imperative Computation, Fall 2010
 * Frank Pfenning
 */

// include bst.h0

/* Implementation */
typedef struct tree* tree;
struct tree {
  elem data;			/* data element */
  bool red;			/* is node red? */
  tree left;			/* left subtree */
  tree right;			/* right subtree */
};

struct bst {
  tree root;			/* root of the tree */
};

typedef bst rbt;

void print_tree(tree T);

bool is_ordered(tree T, elem min, elem max) {
  if (T == NULL) return true; // empty tree is okay!
  if (T->data == NULL) return false; // no data -- bogus
  else {
    key k = elem_key(T->data);
    return
      // make sure the data is between min and max
         (min == NULL || compare(elem_key(min), k) < 0)
      && (max == NULL || compare(k, elem_key(max)) < 0)
      // check the subtrees
      && is_ordered(T->left, min, T->data)
      && is_ordered(T->right, T->data, max);
  }
}

bool is_ordtree(tree T) {
  return is_ordered(T, NULL, NULL);
}

/* sat_colorinv(T, red_parent)
 * iff T satisfies the color invariant
 */
bool sat_colorinv(tree T, bool red_parent)
{ if (T == NULL) return true;
  if (red_parent && T->red) return false;
  return sat_colorinv(T->left, T->red)
    && sat_colorinv(T->right, T->red);
}

/* black_height(T) >= 0 is the black height of T if it exists,
 * -1 otherwise
 */
int black_height(tree T) {
  if (T == NULL) return 0;
  {
    int hleft = black_height(T->left);
    int hright = black_height(T->right);
    if (hleft < 0 || hright < 0 || hleft != hright) return -1;
    //@assert hleft == hright;
    if (T->red) return hleft;
    else return hleft+1;
  }
}

bool is_baltree(tree T) {
  int h = black_height(T);
  return h >= 0 && sat_colorinv(T, false);
}

bool is_rbtree(tree T) {
  return is_ordtree(T) && is_baltree(T);
}

bool is_rbtree_except_root(tree T) {
  bool ok = is_ordtree(T);	/* order invariant is satisfied */
  ok = ok && black_height(T) >= 0; /* height invariant is satisfied */
  if (T->red
      && T->left != NULL && T->left->red /* left child red */
      && (T->right == NULL || !T->right->red)) /* right child not red */
    ok = ok && sat_colorinv(T->left, false) /* pretend parent is not red */
            && sat_colorinv(T->right, false);
  else if (T->red
	   && T->right != NULL && T->right->red /* right child red */
	   && (T->left == NULL || !T->left->red)) /* left child not red */
    ok = ok && sat_colorinv(T->left, false) /* pretend parent is not red */
            && sat_colorinv(T->right, false);
  else
    ok = ok && sat_colorinv(T, false); /* otherwise, color inv. satisfied */
  return ok;
}

bool is_rbt(rbt RBT) {
  return RBT != NULL && is_rbtree(RBT->root);
}

/* for testing purposes... */
bool is_bst(rbt RBT) {
  return is_rbt(RBT);
}

rbt bst_new()
//@ensures is_rbt(\result);
{
  rbt RBT = alloc(struct bst);
  RBT->root = NULL;
  return RBT;
}

elem tree_search(tree T, key k)
//@requires is_ordtree(T);
//@ensures \result == NULL || compare(elem_key(\result), k) == 0;
{
  if (T == NULL) return NULL;
  { key kt = elem_key(T->data);   // key in tree node
    if (compare(k, kt) == 0)
      return T->data;
    else if (compare(k, kt) < 0)
      return tree_search(T->left, k);
    else //@assert compare(k, kt) > 0;
      return tree_search(T->right, k);
  }
}

elem bst_search(rbt RBT, key k)
//@requires is_rbt(RBT);
//@ensures \result == NULL || compare(elem_key(\result), k) == 0;
{
  return tree_search(RBT->root, k);
}

tree balance_left(tree T)
//@requires T != NULL;
//@requires T->red ? is_rbtree(T->left) : is_rbtree_except_root(T->left);
//@ensures \old(T->red) ? is_rbtree_except_root(\result) : is_rbtree(\result);
{ if (T->red) return T;
  //@assert !T->red;
  if (T->left->red
     && T->left->left != NULL && T->left->left->red)
    // rotate right
    { tree root = T->left; T->left = root->right; root->right = T;
      root->left->red = false;
      //@assert root->red == true;
      //@assert root->right->red == false;
      return root;
    } else if (T->left->red
               && T->left->right != NULL && T->left->right->red)
    // double rotate right
    { tree root = T->left->right; T->left->right = root->left;
      root->left = T->left; T->left = root->right; root->right = T;
      root->left->red = false;
      //@assert root->red == true;
      //@assert root->right->red == false;
      return root;
    } else return T;
}

tree balance_right(tree T)
//@requires T != NULL;
//@requires T->red ? is_rbtree(T->right) : is_rbtree_except_root(T->right);
//@ensures \old(T->red) ? is_rbtree_except_root(\result) : is_rbtree(\result);
{ if (T->red) return T;
  //@assert !T->red;
  if (T->right->red
     && T->right->right != NULL && T->right->right->red)
    // rotate left
    { tree root = T->right; T->right = root->left; root->left = T;
      root->right->red = false;
      //@assert root->red == true;
      //@assert root->left->red == false; 
      return root;
    } else if (T->right->red
	       && T->right->left != NULL && T->right->left->red)
    // double rotate left
    { tree root = T->right->left; T->right->left = root->right;
      root->right = T->right; T->right = root->left; root->left = T;
      root->right->red = false;
      //@assert root->red == true;
      //@assert root->left->red == false; 
      return root;
    } else return T;
}

tree tree_insert(tree T, elem e) 
//@requires is_rbtree(T);
/*@ensures \old(T != NULL && T->red)
           ? is_rbtree_except_root(\result)
           : is_rbtree(\result); @*/
//@ensures tree_search(\result, elem_key(e)) == e;
{
  if (T == NULL) {  /* create new node */
    T = alloc(struct tree);
    T->data = e; T->red = true;	/* new nodes are red */
    T->left = NULL; T->right = NULL;
  } else {
    key kt = elem_key(T->data);
    key k = elem_key(e);
    if (compare(k, kt) == 0) {
      T->data = e;
    } else if (compare(k, kt) < 0) {
      //@assert T->red ? (T->left == NULL || !T->left->red) : true;
      T->left = tree_insert(T->left, e);
      //@assert T->red ? is_rbtree(T->left) : is_rbtree_except_root(T->left);
      T = balance_left(T);
    } else {
      //@assert compare(k, kt) > 0;
      //@assert T->red ? (T->right == NULL || !T->right->red) : true;
      T->right = tree_insert(T->right, e);
      //@assert T->red ? is_rbtree(T->right) : is_rbtree_except_root(T->right);
      T = balance_right(T);
    }
  }
  return T;
}

void bst_insert(rbt RBT, elem e)
//@requires is_rbt(RBT);
//@ensures is_rbt(RBT);
{
  // wrapper function to start the process at root
  RBT->root = tree_insert(RBT->root, e);
  //@assert is_rbtree_except_root(RBT->root);
  RBT->root->red = false; /* color root black; might already be black */
  return;
}

void print_tree(tree T) {
  if (T == NULL) print("()");
  else {
    print("(");
    print(elem_key(T->data));
    print(" ");
    if (T->red) print("R "); else print("B ");
    print_tree(T->left);
    print(" ");
    print_tree(T->right);
    print(")");
  }
}

void print_bst(rbt RBT) {
  print_tree(RBT->root);
}