/*
 * Usage:  avl3  list of integers ...
 *
 * Each integer will be checked to see if it is currently in
 * the AVL tree.  If not, it will be inserted. If so, it will
 * be deleted.  The tree starts out empty and the final tree is
 * printed (on its side) in an ASCII-art style
 */

#include <stdlib.h>
#include <stdio.h>
#include <string.h>
typedef int value_t;

#define LEFT 0
#define RIGHT 1
#define NEITHER -1
typedef int direction;
typedef struct node_s {
	value_t		value;
	struct node_s	*next[2];
	int		longer:2;
} *node;
#define Balanced(n) (n->longer < 0)

node avl_find(node tree, value_t target)
{
	while (tree && target != tree->value) {
		direction next_step = (target > tree->value);
		tree = tree->next[next_step];
	}
	return tree;
}

static node avl_rotate_2(node *path_top, direction dir)
{
	node B, C, D, E;

	B = *path_top;
	D = B->next[dir];
	C = D->next[1-dir];
	E = D->next[dir];

	*path_top = D;
	D->next[1-dir] = B;
	B->next[dir] = C;
	B->longer = NEITHER;
	D->longer = NEITHER;
	return E;
}

static node avl_rotate_3(node *path_top, direction dir, direction third)
{
	node B, F, D, C, E;
	B = *path_top;
	F = B->next[dir];
	D = F->next[1-dir];
	/* note: C and E can be NULL */
	C = D->next[1-dir];
	E = D->next[dir];
	*path_top = D;
	D->next[1-dir] = B;
	D->next[dir] = F;
	B->next[dir] = C;
	F->next[1-dir] = E;
	D->longer = NEITHER;

	/* assume both trees are balanced */
	B->longer = F->longer = NEITHER;

	if (third == NEITHER) {
		return NULL;
	} else if (third == dir) {
		/* E holds the insertion so B is unbalanced */
		B->longer = 1-dir;
		return E;
	} else {
		/* C holds the insertion so F is unbalanced */
		F->longer = dir;
		return C;
	}
}

/***************************************************
 * INSERTION                                       *
 ***************************************************/
static inline void avl_rebalance_path(node path, value_t target)
{
	/* Each node in path is currently balanced.
	 * Until we find target, mark each node as longer
	 * in the direction of target because we know we have
	 * inserted target there
	 */
	while (path && target != path->value) {
		direction next_step = (target > path->value);
		path->longer = next_step;
		path = path->next[next_step];
	}
}

static inline void avl_rebalance_insert(node *path_top, value_t target)
{
	node path = *path_top;
	direction first, second, third;
	if (Balanced(path)) {
		;
	} else if (path->longer != (first = (target > path->value))) {
		/* took the shorter path */
		path->longer = NEITHER;
		path = path->next[first];
	} else if (first == (second = (target > path->next[first]->value))) {
		/* just a two-point rotate */
		path = avl_rotate_2(path_top, first);
	} else {
		/* Fine details of the 3 point rotate depend on the
		 * third step.  However there may not be a third step,
		 * if the third point of the rotation is the newly
		 * inserted point.  In that case we record the third
		 * step as NEITHER
		 */
		path = path->next[first]->next[second];
		if (target == path->value)
			third = NEITHER;
		else
			third = (target > path->value);
		path = avl_rotate_3(path_top, first, third);
	}
	avl_rebalance_path(path, target);
}

int avl_insert(node *treep, value_t target)
{
	/* insert the target into the tree, returning 1 on success or 0 if it
	 * already existed
	 */
	node tree = *treep;
	node *path_top = treep;
	while (tree && target != tree->value) {
		direction next_step = (target > tree->value);
		if (!Balanced(tree))
			path_top = treep;
		treep = &tree->next[next_step];
		tree = *treep;
	}
	if (tree)
		return 0;
	tree = malloc(sizeof(*tree));
	tree->next[0] = tree->next[1] = NULL;
	tree->longer = NEITHER;
	tree->value = target;
	*treep = tree;
	avl_rebalance_insert(path_top, target);
	return 1;
}

/******************************************************
 * DELETION                                           *
 *****************************************************/

static inline void avl_swap_del(node *targetp, node *treep, direction dir)
{
	node targetn = *targetp;
	node tree = *treep;

	*targetp = tree;
	*treep = tree->next[1-dir];
	tree->next[LEFT] = targetn->next[LEFT];
	tree->next[RIGHT]= targetn->next[RIGHT];
	tree->longer = targetn->longer;

	free(targetn);
}

static inline node *avl_rebalance_del(node *treep, value_t target,
				      node *targetp)
{
	/* each node from treep down towards target, but
	 * excluding the last, will have a subtree grow
	 * and need rebalancing
	 */
	node targetn = *targetp;

	while (1) {
		node tree = *treep;
		direction dir = (target > tree->value);

		if (tree->next[dir]==NULL)
			break;

		if (Balanced(tree)) {
			tree->longer = 1-dir;
		} else if (tree->longer == dir) {
			tree->longer = NEITHER;
		} else {
			int second = tree->next[1-dir]->longer;
			if (second == dir)
				avl_rotate_3(treep, 1-dir,
					     tree->next[1-dir]->next[dir]->longer);
			else if (second == NEITHER) {
				avl_rotate_2(treep, 1-dir);
				tree->longer = 1-dir;
				(*treep)->longer = dir;
			} else
				avl_rotate_2(treep, 1-dir);
			if (tree == targetn)
				targetp = &(*treep)->next[dir];
		}

		treep = &tree->next[dir];
	}
	return targetp;
}

int avl_delete(node *treep, value_t target)
{
	/* delete the target from the tree, returning 1 on success or 0 if
	 * it wasn't found
	 */
	node tree = *treep, targetn;
	node *path_top = treep;
	node *targetp = NULL;
	direction dir;

	while (tree) {
		dir = (target > tree->value);
		if (target == tree->value)
			targetp = treep;
		if (tree->next[dir] == NULL)
			break;
		if (Balanced(tree)
		    || (tree->longer == (1-dir) && Balanced(tree->next[1-dir]))
			) path_top = treep;
		treep = &tree->next[dir];
		tree = *treep;
	}
	if (!targetp)
		return 0;

	/* adjust balance, but don't lose 'targetp' */
	targetp = avl_rebalance_del(path_top, target, targetp);

	/*
	 * We have re-balanced everything, it remains only to swap the
	 * end of the path (*treep) with the deleted item (*targetp)
	 */
	avl_swap_del(targetp, treep, dir);

	return 1;
}

/****************************************************************
 * PRINTING and TESTING                                         *
 ***************************************************************/
void avl_print(node tree, char *prefix, char *x, direction side)
{
	int pl = strlen(prefix);
	if (tree) {
		strcat(prefix, "   ");
		prefix[pl-1] = x[0];
		avl_print(tree->next[LEFT], prefix, " ,|", LEFT);
		prefix[pl-1] = x[1];
		printf("%.*s--+%c %d\n", pl, prefix, ":/\\"[tree->longer+1],
		       tree->value);
		prefix[pl-1] = x[2];
		avl_print(tree->next[RIGHT],prefix, "|` ", RIGHT);
		prefix[pl] = '\0';
	} else
		printf("%.*s\n", pl-1, prefix);
}

int dir_check_depth(node tree)
{
	int err = 0;
	int rv;
	int left, right;

	if (!tree)
		return 0;
	left = dir_check_depth(tree->next[LEFT]);
	right = dir_check_depth(tree->next[RIGHT]);
	if (left == right) {
		if (!Balanced(tree))
			err = 1;
		rv = left+1;
	} else if (left == right-1) {
		if (tree->longer != RIGHT)
			err = 1;
		rv = right+1;
	} else if (left-1 == right) {
		if (tree->longer != LEFT)
			err = 1;
		rv = left+1;
	} else {
		err = 1;
		rv = 0;
	}
	if (err)
		printf("err at %d: left=%d right=%d longer=%d\n", tree->value,
		       left, right, tree->longer);
	return rv;
}

int main(int argc, char *argv[])
{
	node tree = NULL;
	char prefix[1000];
	int arg;

	for (arg=1; arg<argc; arg++) {
		value_t v = atoi(argv[arg]);
		if (!avl_find(tree, v)) {
			if (avl_insert(&tree, v))
				printf("Inserted %d\n", v);
			else
				printf("Strange failure to insert %d\n", v);
		} else {
			if (avl_delete(&tree, v))
				printf("Deleted: %d\n", v);
			else
				printf("Strange failure to delete %d\n", v);
		}
		dir_check_depth(tree);
	}

	strcpy(prefix, "-");
	avl_print(tree, prefix, " > ", NEITHER);
	exit(0);
}
