1 files changed, 225 insertions, 0 deletions
diff --git a/src/nbc/tree.sml b/src/nbc/tree.sml
new file mode 100644
index 0000000..5d64f51
--- /dev/null
+++ b/src/nbc/tree.sml
@@ -0,0 +1,225 @@
+functor Tree (Key: sig
+	type key
+	val compare: key * key -> order
+end) :> sig
+	type key = Key.key
+	type 'datum tree
+	val size: 'datum tree -> int
+	val empty: 'datum tree
+	val single: key * 'datum -> 'datum tree
+	val add: 'datum tree * key * 'datum -> 'datum tree
+	val find: 'datum tree * key -> 'datum option
+	val exists: 'datum tree * key -> bool
+	val up: 'datum tree -> (key * 'datum) Stream.stream
+	val down: 'datum tree -> (key * 'datum) Stream.stream
+	val upFrom: 'datum tree * key -> (key * 'datum) Stream.stream
+	val downFrom: 'datum tree * key -> (key * 'datum) Stream.stream
+	val fromList: (key * 'datum) list -> 'datum tree
+	val fromStream: (key * 'datum) Stream.stream -> 'datum tree
+	(*
+	val remove: 'datum tree * key -> 'datum tree
+	*)
+end = struct
+	type key = Key.key
+	structure Height = Int8
+	datatype 'datum tree =
+		Leaf
+		| Branch of {
+			height: Height.int
+			, less: 'datum tree
+			, key: key
+			, datum: 'datum
+			, greater: 'datum tree
+		}
+	fun size tree = case tree of
+		Leaf => 0
+		| Branch {less, greater, ...} =>
+			1 + size less + size greater
+	val empty = Leaf
+	fun single (key, datum) = Branch {
+		height = 1
+		, less = Leaf
+		, key = key
+		, datum = datum
+		, greater = Leaf
+	}
+	fun height tree = case tree of
+		Leaf => 0
+		| Branch {height, ...} => height
+	fun calculateHeight {key, datum, less, greater} = Branch {
+		key = key
+		, datum = datum
+		, less = less
+		, greater = greater
+		, height = 1 + Height.max (height less, height greater)
+	}
+	fun rotateLess branch = case branch of
+	(*
+		   b          d
+		 a   d  =>  b   e
+		    c e    a c
+	*)
+		{
+			less = a
+			, key = bKey
+			, datum = bDatum
+			, greater = Branch {
+				less = c
+				, key = dKey
+				, datum = dDatum
+				, greater = e
+				, height = _
+			}
+		} => calculateHeight {
+			less = calculateHeight {
+				less = a
+				, key = bKey
+				, datum = bDatum
+				, greater = c
+			}, key = dKey
+			, datum = dDatum
+			, greater = e
+		} | _ => raise Fail "rotateLess"
+	fun rotateGreater branch = case branch of
+	(*
+		   d        b
+		 b   e => a   d
+		a c          c e
+	*)
+		{
+			less = Branch {
+				less = a
+				, key = bKey
+				, datum = bDatum
+				, greater = c
+				, height = _
+			}, key = dKey
+			, datum = dDatum
+			, greater = e
+		} => calculateHeight {
+			less = a
+			, key = bKey
+			, datum = bDatum
+			, greater = calculateHeight {
+				less = c
+				, key = dKey
+				, datum = dDatum
+				, greater = e
+			}
+		} | _ => raise Fail "rotateGreater"
+	fun balance (branch as {key, datum, less, greater}) =
+		let
+			val heightLess = height less
+			val heightGreater = height greater
+		in
+			if heightLess < heightGreater - 2 then
+				rotateLess branch
+			else if heightLess > heightGreater + 2 then
+				rotateGreater branch
+			else calculateHeight branch
+		end
+	fun add (tree, newKey, newDatum) = case tree of
+		Leaf => single (newKey, newDatum)
+		| Branch {height, less, key, datum, greater} =>
+			case Key.compare (newKey, key) of
+				EQUAL => Branch {
+					height = height
+					, less = less
+					, key = newKey
+					, datum = newDatum
+					, greater = greater
+				} | LESS => balance {
+					less = add (less, newKey, newDatum)
+					, key = key
+					, datum = datum
+					, greater = greater
+				} | GREATER => balance {
+					less = less
+					, key = key
+					, datum = datum
+					, greater = add (greater, newKey, newDatum)
+				}
+	fun find (tree, desiredKey) = case tree of
+		Leaf => NONE
+		| Branch {less, key, datum, greater, ...} =>
+			case Key.compare (desiredKey, key) of
+				EQUAL => SOME datum
+				| LESS => find (less, desiredKey)
+				| GREATER => find (greater, desiredKey)
+	fun exists (tree, desiredKey) = case find (tree, desiredKey) of
+		NONE => false
+		| SOME _ => true
+	fun up tree = Stream.create (fn () =>
+		case tree of
+			Leaf => NONE
+			| Branch {less, key, datum, greater, ...} =>
+				Stream.getItem (
+					Stream.append (
+						up less
+						, Stream.cons (
+							(key, datum)
+							, up greater
+						)
+					)
+				)
+	)
+	fun down tree = Stream.create (fn () =>
+		case tree of
+			Leaf => NONE
+			| Branch {greater, key, datum, less, ...} =>
+				Stream.getItem (
+					Stream.append (
+						down greater
+						, Stream.cons (
+							(key, datum)
+							, down less
+						)
+					)
+				)
+	)
+	fun upFrom (tree, firstKey) = Stream.create (fn () =>
+		case tree of
+			Leaf => NONE
+			| Branch {less, key, datum, greater, ...} =>
+				case Key.compare (firstKey, key) of
+					LESS => Stream.getItem (
+						Stream.append (
+							upFrom (less, firstKey)
+							, Stream.cons (
+								(key, datum)
+								, up greater
+							)
+						)
+					) | EQUAL => SOME (
+						(key, datum)
+						, up greater
+					) | GREATER => Stream.getItem (
+						upFrom (greater, firstKey)
+					)
+	)
+	fun downFrom (tree, firstKey) = Stream.create (fn () =>
+		case tree of
+			Leaf => NONE
+			| Branch {greater, key, datum, less, ...} =>
+				case Key.compare (firstKey, key) of
+					LESS => Stream.getItem (
+						downFrom (less, firstKey)
+					) | EQUAL => SOME (
+						(key, datum)
+						, down less
+					) | GREATER => Stream.getItem (
+						Stream.append (
+							downFrom (greater, firstKey)
+							, Stream.cons (
+								(key, datum)
+								, down less
+							)
+						)
+					)
+	)
+	fun fromStream stream =
+		Stream.fold (fn ((key, datum), tree) =>
+			add (tree, key, datum)
+		) empty stream
+	fun fromList list = fromStream (Stream.fromList list)
+end