compiler: handle multi-iterators
[nit.git] / lib / trees / bintree.nit
index 4c647b3..610ec2e 100644 (file)
@@ -64,7 +64,7 @@ class BinTreeMap[K: Comparable, E]
        #     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
+       redef fun has_key(key) do
                if is_empty then return false
                var res = search_down(root.as(not null), key)
                if res != null then
@@ -85,7 +85,7 @@ class BinTreeMap[K: Comparable, E]
        #     assert tree[1] == "n1"
        #     assert tree.has_key(1)
        #     assert tree[2] == "n2"
-       redef fun [](key: K): E do
+       redef fun [](key) do
                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)
@@ -93,7 +93,9 @@ class BinTreeMap[K: Comparable, E]
                return res.value
        end
 
-       protected fun search_down(from: N, key: K): nullable N do
+       # Search `key` in `from` and its children nodes.
+       protected fun search_down(from: N, key: nullable Object): nullable N do
+               if not key isa Comparable then return null
                var cmp = key <=> from.key
                if cmp == 0 then return from
                if from.left != null and cmp < 0 then
@@ -115,6 +117,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))
@@ -131,6 +134,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))
@@ -148,6 +152,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
@@ -244,11 +249,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
@@ -270,11 +277,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
@@ -307,6 +316,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)
@@ -319,6 +329,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)
@@ -329,14 +340,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)
@@ -367,15 +379,11 @@ end
 class BinTreeNode[K: Comparable, E]
        super TreeNode[K, E]
 
-       private var prev: nullable BinTreeNode[K, E]
-       private var next: nullable BinTreeNode[K, E]
+       private var prev: nullable BinTreeNode[K, E] = null
+       private var next: nullable BinTreeNode[K, E] = null
 
        redef type N: BinTreeNode[K, E]
 
-       init(key: K, item: E) do
-               super(key, item)
-       end
-
        private var left_node: nullable N = null
 
        # `left` tree node child (null if node has no left child)
@@ -444,11 +452,10 @@ end
 private class BinTreeMapIterator[K: Comparable, E]
        super MapIterator[K, E]
 
-       var current: nullable BinTreeNode[K, E]
+       var tree: BinTreeMap[K, E]
+       var current: nullable BinTreeNode[K, E] = null
 
-       init(tree: BinTreeMap[K, E]) do
-               current = tree.first_node
-       end
+       init do current = tree.first_node
 
        redef fun is_ok do return not current == null
        redef fun next do current = current.next