#include <search.h>
#include <stdlib.h>

#define MAXH (sizeof(void*)*8*3/2)

struct node {
	const void *key;
	void *a[2];
	int h;
};

static int height(struct node *n) { return n ? n->h : 0; }

static struct node *rot(struct node *x, int dir)
{
	struct node *y = x->a[!dir];
	struct node *z = y->a[dir];
	int hz = height(z);

	if (hz > height(y->a[!dir])) {
		//      x
		// dir / \ !dir         z
		//    A   y            / \
		//       / \   -->    x   y
		//      z   D        /|   |\
		//     / \          A B   C D
		//    B   C
		x->a[!dir] = z->a[dir];
		y->a[dir] = z->a[!dir];
		z->a[dir] = x;
		z->a[!dir] = y;
		x->h = hz;
		y->h = hz;
		z->h = hz+1;
	} else {
		//      x               y
		//     / \             / \
		//    A   y    -->    x   D
		//       / \         / \
		//      z   D       A   z
		x->a[!dir] = z;
		y->a[dir] = x;
		x->h = hz+1;
		y->h = hz+2;
		z = y;
	}
	return z;
}

void *tsearch(const void *key, void **rootp,
	int (*compar)(const void *, const void *))
{
	if (!rootp)
		return 0;

	void **a[MAXH];
	struct node *n = *rootp;
	struct node *r;
	int h, i=0;

	a[i++] = rootp;
	for (;;) {
		if (!n) break;
		int c = compar(key, n->key);
		if (!c) return n;
		a[i++] = &n->a[c>0];
		n = n->a[c>0];
	}
	r = malloc(sizeof *r);
	if (!r)
		return 0;
	r->key = key;
	r->a[0] = r->a[1] = 0;
	r->h = h = 1;
	*a[--i] = r;
	while (i) {
		n = *a[--i];
		if (n->h > h)
			break;
		n->h = ++h;
		for (int d=0; d<2; d++)
			if (h-2 > height(n->a[d])) {
				*a[i] = rot(n,d);
				return r;
			}
	}
	return r;
}

void *tdelete(const void *restrict key, void **restrict rootp,
	int(*compar)(const void *, const void *))
{
	if (!rootp)
		return 0;

	void **a[MAXH];
	struct node *n = *rootp;
	struct node *parent;
	struct node *child;
	int h, i=0;
	a[i++] = rootp;
	a[i++] = rootp;
	for (;;) {
		if (!n) return 0;
		int c = compar(key, n->key);
		if (!c) break;
		a[i++] = &n->a[c>0];
		n = n->a[c>0];
	}
	parent = *a[i-2];
	if (n->a[0]) {
		struct node *found = n;
		for (a[i++]=&n->a[0], n=n->a[0]; n->a[1]; a[i++]=&n->a[1], n=n->a[1]);
		found->key = n->key;
		child = n->a[0];
	} else {
		child = n->a[1];
	}
	free(n);
	*a[--i] = child;
	h = height(child);
	while (--i) {
		int d, hd;

		h++;
		n = *a[i];
		hd = height(n->a[d=0]);
		if (h > hd)
			hd = height(n->a[d=1]);
		if (h < hd) {
			*a[i] = n = rot(n,!d);
			h = hd;
			if (h < n->h)
				break;
		} else if (h == hd) {
			break;
		} else {
			n->h = h;
		}
	}
	return parent;
}

void *tfind(const void *key, void *const *rootp,
	int(*compar)(const void *, const void *))
{
	if (!rootp)
		return 0;

	struct node *n = *rootp;
	for (;;) {
		if (!n) break;
		int c = compar(key, n->key);
		if (c < 0)
			n = n->a[0];
		else if (c > 0)
			n = n->a[1];
		else
			break;
	}
	return n;
}

static void walk(const struct node *r, void (*action)(const void *, VISIT, int), int d)
{
	if (r == 0)
		return;
	if (r->h == 1)
		action(r, leaf, d);
	else {
		action(r, preorder, d);
		walk(r->a[0], action, d+1);
		action(r, postorder, d);
		walk(r->a[1], action, d+1);
		action(r, endorder, d);
	}
}

void twalk(const void *root, void (*action)(const void *, VISIT, int))
{
	walk(root, action, 0);
}