#N	Balanced binary tree package
#
#  "red-black" 2-3-4 trees represented with binary trees
#  Adapted from Sedgewick, _Algorithms_, 2nd edition, pp 219-228.
#

box rbtree { type T } is
	export insert, search, delete_point, all, size, point, `==`;
	export `<>`;

	type link := pointer {node};

	box node is
		T data;
		link left, right;
		bool red;
	end box;

	link head, z, rb_point;

	int high := 0;			# stats

	proc auto_init is
		z.new;
		z^.left := z;
		z^.right := z;
		z^.red := false;

		head.new;
		head^.right := z;
		rb_point := z;

		assert head.notnull;
		assert head^.right.notnull;
	end proc;

	proc auto_done is
#		free_a_node(head);
#		z.free;
	end proc;

	proc free_a_node (ref link a) is
		if a.notnull and a <> z then
			free_a_node(a^.left);
			free_a_node(a^.right);

			a.free;
		end if;
	end proc;

	func bool search (T v) is
		int count := 0;			# stats

		rb_point := head^.right;
		z^.data := v;

		while v <> rb_point^.data do
			count := count + 1;	# stats

			if v < rb_point^.data then
				rb_point := rb_point^.left;
			else
				rb_point := rb_point^.right;
			end if;
		end while;

		if count > high then
			high := count;
		end if;

		assert head.notnull;
		assert head^.right.notnull;
		return rb_point <> z;
	end func;

	proc delete_point is
		link x, p, c;
		bool goleft := false;

		assert (rb_point <> z);

		x := head^.right;
		p := head;

		while x <> rb_point do
			p := x;
			goleft := (rb_point^.data < x^.data);
			if goleft then
				x := x^.left;
			else
				x := x^.right;
			end if;
		end while;

		if rb_point^.right == z then
			x := x^.left;
		elsif rb_point^.right^.left == z then
			x := x^.right;
			x^.left := rb_point^.left;
		else
			c := x^.right;
			while c^.left^.left <> z do
				c := c^.left;
			end while;
			x := c^.left;
			c^.left := x^.right;
			x^.left := rb_point^.left;
			x^.right := rb_point^.right;
		end if;

		if goleft then
			p^.left := x;
		else
			p^.right := x;
		end if;

		rb_point := z;
	end proc;

	func link rb_rotate(T v; link y) is
		link c, gc;
		bool goleft;

		goleft := (y <> head) and (v < y^.data);
		if goleft then
			c := y^.left;
		else
			c := y^.right;
		end if;

		if v < c^.data then
			gc := c^.left;
			c^.left := gc^.right;
			gc^.right := c;
		else
			gc := c^.right;
			c^.right := gc^.left;
			gc^.left := c;
		end if;

		if goleft then
			y^.left := gc;
		else
			y^.right := gc;
		end if;

		assert head.notnull;
		assert head^.right.notnull;
		return gc;
	end func;

	func link split (T v; link gg, g, p, x) is

		x^.red := true;
		x^.left^.red := false;
		x^.right^.red := false;

		if p^.red then
			g^.red := true;
			if (v < g^.data) <> (v < p^.data) then
				p := rb_rotate(v, g);
			end if;
			x := rb_rotate(v, gg);
			x^.red := false;
		end if;

		head^.right^.red := false;
		assert head.notnull;
		assert head^.right.notnull;
		return x;
	end func;

	proc insert (T v) is
		link x, gg, g, p;
		bool goleft := false;

		x := head^.right;
		p := head;
		g := head;
		gg := head;

		if x^.left^.red and x^.right^.red then
			x := split(v, gg, g, p, x);
		end if;

		while x <> z do
			gg := g;
			g := p;
			p := x;
			goleft := (v < x^.data);
			if goleft then
				x := x^.left;
			else
				x := x^.right;
			end if;

			if x^.left^.red and x^.right^.red then
				x := split(v, gg, g, p, x);
			end if;
		end while;

		x.new;
		x^.data := v;
		x^.left := z;
		x^.right := z;

		if goleft then
			p^.left := x;
		else
			p^.right := x;
		end if;

		rb_point := x;
		x := split(v, gg, g, p, x);
		assert head.notnull;
		assert head^.right.notnull;
	end proc;

	gen ref T all is
		gen ref T sub (const link p) is
		link left, right;

			assert p.notnull;
			if p <> z then
				left := p^.left;
			 	right := p^.right;

				assert right.notnull;
				assert left.notnull;

				yield sub(left);

				rb_point := p;
				yield p^.data;

				assert right.notnull;
				yield sub(right);
			end if;
		end gen;

		assert z.notnull;
		assert head.notnull;
		assert head^.right.notnull;
		yield sub(head^.right);
	end gen;

	const func int size is
		func int recurse (link p) is
			if p == z then
				return 0;
			else
				return 1 + recurse(p^.left) + recurse(p^.right);
			end if;
		end func;

		return recurse(head^.right);
	end func;

	func ref T point is

		assert(rb_point <> z);
		return rb_point^.data;
	end func;

	proc print (ref file f) is
		f.putline(size.to_string + '\n');
		for T i := all do
			i.print(f);
			f.putline("\n");
		end for;
	end proc;

	proc scan (ref file f) is
	int length;
	T x;

		auto_done;
		length.scan(f);
		while length > 0 do
			x.scan(f);
			insert(x);
			length := length - 1;
		end while;
	end proc;

	func bool `==` (ref rbtree{T} x) is

		if size <> x.size then
			return false;
		end if;

		for T i := all, T j := x.all do
			if i <> j then
				return false;
			end if;
		end for;

		return true;
	end func;

	func bool `<>` (ref rbtree{T} x) is
		return not `==` (x);
	end func;
end box;