X-Git-Url: http://nitlanguage.org

diff --git a/lib/trees/bintree.nit b/lib/trees/bintree.nit
index 8165a24..6b6b91c 100644
--- a/lib/trees/bintree.nit
+++ b/lib/trees/bintree.nit
@@ -45,6 +45,8 @@ class BinTreeMap[K: Comparable, E]
 	redef type N: BinTreeNode[K, E]
 
 	private var len = 0
+	private var first_node: nullable BinTreeNode[K, E] = null
+	private var last_node: nullable BinTreeNode[K, E] = null
 
 	# O(n) in worst case, average is O(h) with h: tree height
 	#
@@ -53,22 +55,49 @@ class BinTreeMap[K: Comparable, E]
 	#     tree[1] = "n1"
 	#     assert not tree.is_empty
 	redef fun is_empty do return root == null
+
+	# O(n) in worst case, average is O(h) with h: tree height
+	#
+	#     var tree = new BinTreeMap[Int, String]
+	#     assert not tree.has_key(1)
+	#     for i in [4, 2, 1, 5, 3] do tree[i] = "n{i}"
+	#     assert not tree.has_key(0)
+	#     assert tree.has_key(2)
+	#     assert not tree.has_key(6)
+	redef fun has_key(key: K): Bool do
+		if is_empty then return false
+		var res = search_down(root.as(not null), key)
+		if res != null then
+			cache_node = res
+			return true
+		end
+		return false
+	end
+
+	private var cache_node: nullable N = null
+
 	# Get the node value associated to `key`
 	# O(n) in worst case, average is O(h) with h: tree height
 	#
 	#     var tree = new BinTreeMap[Int, String]
 	#     for i in [4, 2, 1, 5, 3] do tree[i] = "n{i}"
+	#     assert tree.has_key(1)
 	#     assert tree[1] == "n1"
+	#     assert tree.has_key(1)
+	#     assert tree[2] == "n2"
 	redef fun [](key: K): E do
-		assert not_empty: root != null
+		assert not_empty: not is_empty
+		if cache_node != null and cache_node.key == key then return cache_node.value
 		var res = search_down(root.as(not null), key)
 		assert has_key: res != null
 		return res.value
 	end
 
+	# Search `key` in `from` and its children nodes.
 	protected fun search_down(from: N, key: K): nullable N do
-		if key == from.key then return from
-		if from.left != null and key < from.key then
+		var cmp = key <=> from.key
+		if cmp == 0 then return from
+		if from.left != null and cmp < 0 then
 			return search_down(from.left.as(not null), key)
 		else if from.right != null then
 			return search_down(from.right.as(not null), key)
@@ -87,6 +116,7 @@ class BinTreeMap[K: Comparable, E]
 		return min_from(root.as(not null)).value
 	end
 
+	# Get the left-most child from `node`.
 	protected fun min_from(node: N): N do
 		if node.left == null then return node
 		return min_from(node.left.as(not null))
@@ -103,6 +133,7 @@ class BinTreeMap[K: Comparable, E]
 		return max_from(root.as(not null)).value
 	end
 
+	# Get the right-most child from `node`.
 	protected fun max_from(node: N): N do
 		if node.right == null then return node
 		return max_from(node.right.as(not null))
@@ -120,6 +151,7 @@ class BinTreeMap[K: Comparable, E]
 		insert_node(new BinTreeNode[K, E](key, item))
 	end
 
+	# Add `node` in the tree.
 	protected fun insert_node(node: N) do
 		len += 1
 		if root == null then
@@ -127,18 +159,27 @@ class BinTreeMap[K: Comparable, E]
 		else
 			shift_down(root.as(not null), node)
 		end
+		if first_node == null then
+			first_node = node
+		end
+		if last_node != null then
+			last_node.next = node
+			node.prev = last_node
+		end
+		last_node = node
 	end
 
 	# Push down the `node` in tree from a specified `from` index
 	protected fun shift_down(from, node: N) do
-		if node.key < from.key then
+		var cmp = node.key <=> from.key
+		if cmp < 0 then
 			if from.left == null then
 				from.left = node
 				node.parent = from
 			else
 				shift_down(from.left.as(not null), node)
 			end
-		else if node.key > from.key then
+		else if cmp > 0 then
 			if from.right == null then
 				from.right = node
 				node.parent = from
@@ -178,6 +219,16 @@ class BinTreeMap[K: Comparable, E]
 			min.left = node.left
 			min.left.parent = min
 		end
+		if first_node == node then
+			first_node = null
+		end
+		if last_node == node then
+			last_node = node.prev
+			last_node.next = null
+		else
+			node.prev.next = node.next
+			node.next.prev = node.prev
+		end
 		return node.value
 	end
 
@@ -197,11 +248,13 @@ class BinTreeMap[K: Comparable, E]
 
 	# Perform left rotation on `node`
 	#
+	# ~~~raw
 	#     N             Y
 	#    / \     >     / \
 	#   a   Y         N   c
 	#      / \   <   / \
 	#     b   c     a   b
+	# ~~~
 	#
 	protected fun rotate_left(node: N) do
 		var y = node.right
@@ -223,11 +276,13 @@ class BinTreeMap[K: Comparable, E]
 
 	# Perform right rotation on `node`
 	#
+	# ~~~raw
 	#     N             Y
 	#    / \     >     / \
 	#   a   Y         N   c
 	#      / \   <   / \
 	#     b   c     a   b
+	# ~~~
 	#
 	protected fun rotate_right(node: N) do
 		var y = node.left
@@ -260,6 +315,7 @@ class BinTreeMap[K: Comparable, E]
 		return sorted
 	end
 
+	# Sort the tree from `node` and place results in `sorted`.
 	protected fun sort_down(node: N, sorted: Array[E]) do
 		if node.left != null then sort_down(node.left.as(not null), sorted)
 		sorted.add(node.value)
@@ -272,6 +328,7 @@ class BinTreeMap[K: Comparable, E]
 		return "[{print_tree(root)}]"
 	end
 
+	# Print the tree starting from `node`.
 	protected fun print_tree(node: N): String do
 		var s = new FlatBuffer
 		s.append(node.to_s)
@@ -282,14 +339,15 @@ class BinTreeMap[K: Comparable, E]
 
 	redef fun show_dot do
 		assert not_empty: root != null
-		var f = new OProcess("dot", "-Txlib")
+		var f = new ProcessWriter("dot", "-Txlib")
 		f.write "digraph \{\n"
 		dot_down(root.as(not null), f)
 		f.write "\}\n"
 		f.close
 	end
 
-	protected fun dot_down(node: N, f: OProcess) do
+	# Translate the tree in dot format starting from `node`.
+	protected fun dot_down(node: N, f: ProcessWriter) do
 		if node.left != null then dot_down(node.left.as(not null), f)
 		f.write node.to_dot
 		if node.right != null then dot_down(node.right.as(not null), f)
@@ -302,46 +360,55 @@ class BinTreeMap[K: Comparable, E]
 	#     for i in [4, 2, 1, 5, 3] do tree[i] = "n{i}"
 	#     assert tree.length == 5
 	redef fun length do return len
+
+	# Nodes are iterated in the same order in which they were added to the tree.
+	# O(n)
+	#
+	#     var tree = new BinTreeMap[Int, String]
+	#     for i in [4, 2, 1, 5, 3] do tree[i] = "n{i}"
+	#     var keys = new Array[Int]
+	#     for k, v in tree do
+	#         keys.add k
+	#     end
+	#     assert keys == [4, 2, 1, 5, 3]
+	redef fun iterator do return new BinTreeMapIterator[K, E](self)
 end
 
 # TreeNode used by BinTree
 class BinTreeNode[K: Comparable, E]
 	super TreeNode[K, E]
 
-	redef type SELF: BinTreeNode[K, E]
+	private var prev: nullable BinTreeNode[K, E] = null
+	private var next: nullable BinTreeNode[K, E] = null
 
-	init(key: K, item: E) do
-		super(key, item)
-	end
+	redef type N: BinTreeNode[K, E]
 
-	private var left_node: nullable SELF = null
+	private var left_node: nullable N = null
 
 	# `left` tree node child (null if node has no left child)
-	fun left: nullable SELF do return left_node
+	fun left: nullable N do return left_node
 
 	# set `left` child for this node (or null if left no child)
 	# ENSURE: node.key < key (only if node != null)
-	fun left=(node: nullable SELF) do
-		assert node != null implies node.key < key
+	fun left=(node: nullable N) do
+		#assert node != null implies node.key < key
 		left_node = node
 	end
 
-	private var right_node: nullable SELF = null
+	private var right_node: nullable N = null
 
 	# `right` tree node child (null if node has no right child)
-	fun right: nullable SELF do return right_node
+	fun right: nullable N do return right_node
 
 	# set `right` child for this node (or null if right no child)
 	# ENSURE: node.key < key (only if node != null)
-	fun right=(node: nullable SELF) do
-		if node != null then
-			assert node.key > key
-		end
+	fun right=(node: nullable N) do
+		#assert node != null implies node.key > key
 		right_node = node
 	end
 
 	# `parent` of the `parent` of this node (null if root)
-	fun grandparent: nullable SELF do
+	fun grandparent: nullable N do
 		if parent == null then
 			return null
 		else
@@ -351,7 +418,7 @@ class BinTreeNode[K: Comparable, E]
 
 	# Other child of the `grandparent`
 	# `left` or `right` depends on the position of the current node against its parent
-	fun uncle: nullable SELF do
+	fun uncle: nullable N do
 		var g = grandparent
 		if g == null then
 			return null
@@ -366,7 +433,7 @@ class BinTreeNode[K: Comparable, E]
 
 	# Other child of the parent
 	# `left` or `right` depends on the position of the current node against its parent
-	fun sibling: nullable SELF do
+	fun sibling: nullable N do
 		if parent == null then
 			return null
 		else if self == parent.left then
@@ -378,6 +445,19 @@ class BinTreeNode[K: Comparable, E]
 		end
 	end
 
-	redef fun to_s do return "\{{key}: {value}\}"
+	redef fun to_s do return "\{{key}: {value or else ""}\}"
 end
 
+private class BinTreeMapIterator[K: Comparable, E]
+	super MapIterator[K, E]
+
+	var tree: BinTreeMap[K, E]
+	var current: nullable BinTreeNode[K, E] = null
+
+	init do current = tree.first_node
+
+	redef fun is_ok do return not current == null
+	redef fun next do current = current.next
+	redef fun item do return current.value
+	redef fun key do do return current.key
+end