/* Binary search tree without balancing
 * 15-122 Principles of Imperative Computation, Fall 2010
 * Frank Pfenning
 */

/* compile with -lstring */

/* Client types and code */

/* elements */
struct elem {
  string word; /* key */
  int count;   /* information */
};

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

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

/* Interface type definitions */
/* provided above by client and used below in implementation */
typedef string key;
typedef struct elem* elem;  /* NULL must be an elem */
key elem_key(elem e);
int compare(key k1, key k2);

/* Interface */

typedef struct bst* bst;
bst bst_new();
void bst_insert(bst B, elem x); /* might change to return old entry with key(x)? */
elem bst_search(bst B, key k);  /* return NULL if not in tree */

/* Implementation */

bool is_sorted(key[] A, int lower, int upper)
//@requires 0 <= lower && lower <= upper && upper <= \length(A);
{ int i;
  for (i = lower; i < upper-1; i++)
    //@loop_invariant lower == upper || (lower <= i && i <= upper-1);
    if (!(compare(A[i],A[i+1]) <= 0)) return false;
  return true;
}

typedef struct tree* tree;
struct tree {
  elem data;
  tree left;
  tree right;
};

struct bst {
  tree root;
};

/* this is very weak, compared to the true ordering invariant */
bool is_ordtree(tree T) {
  key k;
  if (T == NULL) return true;  /* an empty tree is a BST */
  k = elem_key(T->data);
  return (T->left == NULL || (compare(elem_key(T->left->data), k) < 0
			      && is_ordtree(T->left)))
      && (T->right == NULL || (compare(k, elem_key(T->right->data)) < 0
			       && is_ordtree(T->right)));
}

int tree_size(tree T)
{
  if (T == NULL) return 0;
  return 1 + tree_size(T->left) + tree_size(T->right);
}

int bst_size(bst B)
// do not require this, so we can use it in invariant-checking code
// @requires is_bst(B);
{
  return tree_size(B->root);
} 

/* in-order traversal, adding all keys to keyqueue */
void tree_traverse(tree T, queue keyqueue)
{
  if (T == NULL) return;
  //@assert T->data != NULL;
  tree_traverse(T->left, keyqueue); /* first left subtree */
  enq(keyqueue, elem_key(T->data)); /* then key of node */
  tree_traverse(T->right, keyqueue); /* then right subtree */
}

key[] bst_traverse(bst B)
// don't check these invariants, so we can use it in invariant-checking code
// @requires is_bst(B);
// @ensures is_sorted(\result, 0, \length(\result));
{ key[] keyarray;
  queue keyqueue = q_new();
  tree_traverse(B->root, keyqueue);
  keyarray = q_toarray(keyqueue);    // new function for queues
  return keyarray;
}

/* weak, because is_ordtree is weak */
/*
bool is_bst(bst B) {
  return B != NULL && is_ordtree(B->root);
}
*/

/* better version: check that in-order traversal is sorted */
bool is_bst(bst B) {
  if (B == NULL) return false;
  {
    key[] A = bst_traverse(B);
    return is_sorted(A, 0, bst_size(B));
  }
}


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;
  { key kt = elem_key(T->data);
    if (compare(k, kt) == 0) return T->data;
    else if (compare(k, kt) < 0)
      return tree_search(T->left, k);
    else
      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 tree_insert(tree T, elem x) 
//@requires is_ordtree(T);
//@ensures is_ordtree(T);
//@ensures tree_search(\result, elem_key(x)) != NULL;
{
  if (T == NULL) {  /* create new node */
    T = alloc(struct tree); T->data = x;
    T->left = NULL; T->right = NULL;
  } else {
    key kt = elem_key(T->data);
    key k = elem_key(x);
    if (compare(k, kt) == 0) {
      T->data = x;
    } else if (compare(k, kt) < 0) {
      T->left = tree_insert(T->left, x);
    } else {
      //@assert compare(k, kt) > 0;
      T->right = tree_insert(T->right, x);
    }
  }
  return T;
}

void bst_insert(bst B, elem x)
//@requires is_bst(B);
//@ensures is_bst(B);
{
  /* wrapper function to start the process at root */
  B->root = tree_insert(B->root, x);
  return;
}

int main()
{
  bst B;
  elem e;
  int size;
  key[] keyarray;
  int i;

  println("");

  B = bst_new();
  e = alloc(struct elem);
  println("Inserting word: queue (count = 16)");
  e->word = "queue"; e->count = 16;
  bst_insert(B,e);
  assert(bst_search(B,e->word) != NULL, "word not inserted correctly");

  println("Inserting word: stack (count = 12)");
  e = alloc(struct elem);
  e->word = "stack"; e->count = 12;
  bst_insert(B,e);
  assert(bst_search(B,e->word) != NULL, "word not inserted correctly");

  println("Inserting word: list (count = 5)");
  e = alloc(struct elem);
  e->word = "list"; e->count = 5;
  bst_insert(B,e);
  assert(bst_search(B,e->word) != NULL, "word not inserted correctly");

  println("Inserting word: graph (count = 12)");
  e = alloc(struct elem);
  e->word = "graph"; e->count = 12;
  bst_insert(B,e);
  assert(bst_search(B,e->word) != NULL, "word not inserted correctly");

  println("Inserting word: heap (count = 9)");
  e = alloc(struct elem);
  e->word = "heap"; e->count = 9;
  bst_insert(B,e);
  assert(bst_search(B,e->word) != NULL, "word not inserted correctly");

  println("Inserting word: tree (count = 4)");
  e = alloc(struct elem);
  e->word = "tree"; e->count = 4;
  bst_insert(B,e);
  assert(bst_search(B,e->word) != NULL, "word not inserted correctly");

  println("Inserting word: array (count = 7)");
  e = alloc(struct elem);
  e->word = "array"; e->count = 7;
  bst_insert(B,e);
  assert(bst_search(B,e->word) != NULL, "word not inserted correctly");

  println("Inserting word: hashtable (count = 101)");
  e = alloc(struct elem);
  e->word = "hashtable"; e->count = 101;
  bst_insert(B,e);
  assert(bst_search(B,e->word) != NULL, "word not inserted correctly");

  println("Inserting word: deadbeef (count = 66)");
  e = alloc(struct elem);
  e->word = "deadbeef"; e->count = 66;
  bst_insert(B,e);
  assert(bst_search(B,e->word) != NULL, "word not inserted correctly");

  println("Reinserting word: tree (count = 12)");
  e = alloc(struct elem);
  e->word = "tree"; e->count = 12;   // insert key with new value
  bst_insert(B,e);
  assert(bst_search(B,e->word) != NULL, "word not inserted correctly");

  println("");

  size = bst_size(B);
  print("Size of BST: ");
  printint(size);
  println("");
  assert(size == 9, "incorrect size");

  println("");

  print("Retrieving count for word: queue  ");
  e = bst_search(B, "queue");
  if (e != NULL)
    printint(e->count); 
  else
    print("unknown");
  println("");
  assert(e->count == 16, "incorrect count reported");

  print("Retrieving count for word: deadbeef  ");
  e = bst_search(B, "deadbeef");
  if (e != NULL)
    printint(e->count); 
  else
    print("unknown");
  println("");
  assert(e->count == 66, "incorrect count reported");

  print("Retrieving count for word: set  ");
  e = bst_search(B, "set");
  if (e != NULL)
    printint(e->count); 
  else
    print("unknown");
  println("");
  assert(e == NULL, "incorrect count reported");

  print("Retrieving count for word: array  ");
  e = bst_search(B, "array");
  if (e != NULL)
    printint(e->count); 
  else
    print("unknown");
  println("");
  assert(e->count == 7, "incorrect count reported");

  print("Retrieving count for word: tree  ");
  e = bst_search(B, "tree");
  if (e != NULL)
    printint(e->count); 
  else
    print("unknown");
  println("");
  assert(e->count == 12, "incorrect count reported");

  println("");

  println("Inorder traversal of BST:");
  keyarray = bst_traverse(B);
  for (i = 0; i < 9; i++) {
    print(keyarray[i]); print("  ");
  }
  println("");
  assert(is_sorted(keyarray, 0, 9), "inorder traversal is incorrect");

  println("");

  return 0;
}