/* XXX possibly buggy!  an example for testing.  can you find all the bugs? */

/* AVL Trees
 * Ephemeral (imperative) version
 * 15-122 Principles of Imperative Computation, Spring 2011
 * Frank Pfenning
 */

#use <string>

/* Client-side implementation */

/* elements */
struct wcount {
  string word;			/* key */
  int count;			/* data = wordcount */
};
typedef struct wcount* wcount;

/* key comparison */
int compare(string s1, string s2) {
  return string_compare(s1,s2);
}

/* extracting keys from elements */
string elem_key(wcount wc)
//@requires wc != NULL;
{
  return wc->word;
}

/* Client-side interface declarations */
typedef string key;
typedef struct wcount* elem;	/* NULL must be an elem */
key elem_key(elem e);
int compare(key k1, key k2);

/* Library-side interface declarations */
typedef struct bst* bst;
bst bst_new();
void bst_insert(bst B, elem e)  /* replace if elem with same key as x in B */
//@requires e != NULL;
  ;
elem bst_search(bst B, key k);  /* return NULL if not in tree */

/* Implementation */
struct tree {
  int height;
  struct tree* left;
  elem data;
  struct tree* right;
};
typedef struct tree* tree;

struct bst {
  tree root;
};

/* is_ordered(T, lower, upper) checks if all elements in T
 * are strictly in the interval (elem_key(lower),elem_key(kupper)).
 * lower = NULL represents -infinity; upper = NULL represents +infinity
 */
bool is_ordered(tree T, elem lower, elem upper) {
  if (T == NULL) return true;
  if (T->data == NULL) return false;
  key k = elem_key(T->data);
  if (!(lower == NULL || compare(elem_key(lower),k) < 0))
    return false;
  if (!(upper == NULL || compare(k,elem_key(upper)) < 0))
    return false;
  return is_ordered(T->left, lower, T->data)
    && is_ordered(T->right, T->data, upper);
}

bool is_ordtree(tree T) {
  /* initially, we have no bounds - pass in NULL */
  return is_ordered(T, NULL, NULL);
}

int height(tree T) {
  return T == NULL ? 0 : T->height;
}

bool is_balanced(tree T) {
  if (T == NULL) return true;
  int h = T->height;
  int hl = height(T->left);
  int hr = height(T->right);
  if (!(h == (hl > hr ? hl+1 : hr+1))) return false;
  if (hl > hr+1 || hr > hl+1) return false;
  return is_balanced(T->left) && is_balanced(T->right);
}

bool is_avl(tree T) {
  return is_ordtree(T) && is_balanced(T);
}

tree leaf(elem e)
//@requires e != NULL;
//@ensures is_avl(\result);
{
  tree T = alloc(struct tree);
  T->left = NULL;
  T->data = e;
  T->right = NULL;
  T->height = 1;
  return T;
}

void fix_height(tree T)
//@requires T != NULL;
//@requires is_balanced(T->left) && is_balanced(T->right);
// //@ensures is_balanced(T);
{
  int hl = height(T->left);
  int hr = height(T->right);
  T->height = (hl > hr ? hl+1 : hr+1);
  return;
}

/* rotate_right(T) may not be balanced if first step
 * of a double rotation, but heights will be accurate.
 */
tree rotate_right(tree T)
//@requires T != NULL;
//@requires is_ordtree(T);
//@ensures is_ordtree(\result);
{
  tree root = T->left;
  T->left = root->right;
  root->right = T;
  fix_height(root->right);	/* must be first */
  fix_height(root);
  return root;
}

/* rotate_right(T) may not be balanced if first step
 * of a double rotation, but heights will be accurate.
 */
tree rotate_left(tree T)
//@requires T != NULL;
//@requires is_ordtree(T);
//@ensures is_ordtree(\result);
{
  tree root = T->right;
  T->right = root->left;
  root->left = T;
  fix_height(root->left);	/* must be first */
  fix_height(root);
  return root;
}

tree rebalance_left(tree T)
//@requires T != NULL;
//@requires is_avl(T->left) && is_avl(T->right);
/* also requires that T->left is result of insert into T */
//@ensures is_avl(\result);
{
  tree l = T->left;
  tree r = T->right;
  int hl = height(l);
  int hr = height(r);
  if (hl > hr+1) {
    //@assert hl == hr+2;
    if (height(l->left) > height(l->right)) {
      return rotate_right(T);
    } else {
      //@assert height(l->left) < height(l->right);
      /* double rotate right */
      T->left = rotate_left(T->left);
      return rotate_right(T);
    }
  } else { //@assert !(hl > hr+1);
    fix_height(T);
    return T;
  }
}

tree rebalance_right(tree T)
//@requires T != NULL;
//@requires is_avl(T->left) && is_avl(T->right);
/* also requires that T->right is result of insert into T */
//@ensures is_avl(\result);
{
  tree l = T->left;
  tree r = T->right;
  int hl = height(l);
  int hr = height(r);
  if (hr > hl+1) {
    //@assert hr == hl+2;
    if (height(r->right) > height(r->left)) {
      return rotate_left(T);
    } else {
      //@assert height(r->right) < height(r->left);
      /* double rotate left */
      T->right = rotate_right(T->right);
      return rotate_left(T);
    }
  } else { //@assert !(hr > hl+1);
    fix_height(T);
    return T;
  }
}
      
bool is_bst(bst B) {
  if (B == NULL) return false;
  return is_avl(B->root);
}

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

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;
  int r = compare(k, elem_key(T->data));
  if (r == 0)
    return T->data;
  else if (r < 0)
    return tree_search(T->left, k);
  else //@assert r > 0;
    return tree_search(T->right, k);
}

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

/* tree_insert(T, e) returns the modified tree
 * this avoids some complications in case T = NULL
 */
tree tree_insert(tree T, elem e)
//@requires is_avl(T);
//@ensures is_avl(\result);
{
  assert(e != NULL);		/* cannot insert NULL element */
  if (T == NULL) {
    T = leaf(e);		/* create new leaf with data e */
  } else {
    int r = compare(elem_key(e), elem_key(T->data));
    if (r < 0) {
      T = tree_insert(T->left, e);
      T = rebalance_left(T);
    } else if (r == 0) {
      T->data = e;
    } else { //@assert r > 0;
      T = tree_insert(T->right, e);
      T = rebalance_right(T);
    }
  }
  return T;
}

void bst_insert(bst B, elem e)
//@requires is_bst(B);
//@ensures is_bst(B);
{
  assert(e != NULL);		/* cannot insert NULL element */
  B->root = tree_insert(B->root, e);
  return;
}