lib/ordered_tree.nit

   1 # This file is part of NIT ( http://www.nitlanguage.org ).
   2 #
   3 # Copyright 2012 Jean Privat <jean@pryen.org>
   4 #
   5 # Licensed under the Apache License, Version 2.0 (the "License");
   6 # you may not use this file except in compliance with the License.
   7 # You may obtain a copy of the License at
   8 #
   9 #     http://www.apache.org/licenses/LICENSE-2.0
  10 #
  11 # Unless required by applicable law or agreed to in writing, software
  12 # distributed under the License is distributed on an "AS IS" BASIS,
  13 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14 # See the License for the specific language governing permissions and
  15 # limitations under the License.
  16
  17 # Manipulation and presentation of ordered trees.
  18 module ordered_tree
  19
  20 # Generic structure to manage and display an ordered tree
  21 #
  22 # Ordered tree are tree where the elements of a same parent are in a specific order
  23 #
  24 # Elements of the trees are added with the `add` method that takes a parent and
  25 # a sub-element.
  26 # If the parent is `null`, then the element is considered a root.
  27 #
  28 # ~~~~
  29 # var t = new OrderedTree[String]
  30 # t.add(null, "root")
  31 # t.add("root", "child1")
  32 # t.add("root", "child2")
  33 # t.add("child1", "grand-child")
  34 # assert t.length == 4
  35 # ~~~~
  36 #
  37 # By default, the elements with a same parent
  38 # are visited in the order they are added.
  39 #
  40 # ~~~
  41 # assert t.to_a == ["root", "child1", "grand-child", "child2"]
  42 # assert t.write_to_string == """
  43 # root
  44 # |--child1
  45 # |  `--grand-child
  46 # `--child2
  47 # """
  48 # ~~~
  49 #
  50 # The `sort_with` method can be used reorder elements
  51 #
  52 # ~~~
  53 # t.add("root", "aaa")
  54 # assert t.to_a == ["root", "child1", "grand-child", "child2", "aaa"]
  55 # t.sort_with(alpha_comparator)
  56 # assert t.to_a == ["root", "aaa", "child1", "grand-child", "child2"]
  57 # ~~~
  58 #
  59 # This class can be used as it to work with generic trees but can also be specialized to provide more specific
  60 # behavior or display. It is why the internal attributes are mutable.
  61 class OrderedTree[E: Object]
  62         super Writable
  63         super Collection[E]
  64
  65         # The roots of the tree (in sequence)
  66         var roots = new Array[E]
  67
  68         # The branches of the trees.
  69         # For each element, the ordered array of its direct sub-elements.
  70         var sub = new HashMap[E, Array[E]]
  71
  72         # Add a new element `e` in the tree.
  73         # `p` is the parent of `e`.
  74         # if `p` is null, then `e` is a root element.
  75         fun add(p: nullable E, e: E)
  76         do
  77                 if p == null then
  78                         roots.add(e)
  79                 else if sub.has_key(p) then
  80                         sub[p].add(e)
  81                 else
  82                         sub[p] = [e]
  83                 end
  84         end
  85
  86         # print the full tree on `o`
  87         # Write a ASCII-style tree and use the `display` method to label elements
  88         redef fun write_to(stream: Writer)
  89         do
  90                 for r in roots do
  91                         stream.write display(r)
  92                         stream.write "\n"
  93                         sub_write_to(stream, r, "")
  94                 end
  95         end
  96
  97         private fun sub_write_to(o: Writer, e: E, prefix: String)
  98         do
  99                 if not sub.has_key(e) then return
 100                 var subs = sub[e]
 101                 var last = subs.last
 102                 for e2 in subs do
 103                         if e2 != last then
 104                                 o.write "{prefix}|--{display(e2)}\n"
 105                                 sub_write_to(o, e2, prefix+"|  ")
 106                         else
 107                                 o.write "{prefix}`--{display(e2)}\n"
 108                                 sub_write_to(o, e2, prefix+"   ")
 109                         end
 110                 end
 111         end
 112
 113         # Sort roots and other elements using a comparator method
 114         # This method basically sorts roots then each group of children
 115         fun sort_with(comparator: Comparator)
 116         do
 117                 comparator.sort(roots)
 118                 for a in sub.values do
 119                         comparator.sort(a)
 120                 end
 121         end
 122
 123         # How to display a specific element of the tree
 124         # By defaut, uses `to_s`
 125         #
 126         # Subclasses should redefine this method to provide a specific output
 127         fun display(e: E): String do return e.to_s
 128
 129         # Get an array of the contained elements
 130         # Order is preserved
 131         #
 132         #     var tree = new OrderedTree[Int]
 133         #     tree.add(null, 1)
 134         #     tree.add(1, 11)
 135         #     tree.add(11, 111)
 136         #     tree.add(11, 112)
 137         #     tree.add(1, 12)
 138         #     tree.add(12, 121)
 139         #     tree.add(12, 122)
 140         #     tree.add(null, 2)
 141         #     tree.add(2, 21)
 142         #     tree.add(2, 22)
 143         #     assert tree.to_a == [1, 11, 111, 112, 12, 121, 122, 2, 21, 22]
 144         redef fun to_a: Array[E] do
 145                 var res = new Array[E]
 146                 for r in roots do sub_to_a(r, res)
 147                 return res
 148         end
 149
 150         private fun sub_to_a(e: E, res: Array[E]) do
 151                 res.add e
 152                 if sub.has_key(e) then for e2 in sub[e] do sub_to_a(e2, res)
 153         end
 154
 155         #     var tree = new OrderedTree[Int]
 156         #     assert tree.is_empty
 157         #     tree.add(null, 1)
 158         #     assert not tree.is_empty
 159         redef fun is_empty: Bool do return roots.is_empty
 160
 161         #     var tree = new OrderedTree[Int]
 162         #     tree.add(null, 1)
 163         #     tree.add(1, 11)
 164         #     assert tree.first == 1
 165         redef fun first do return roots.first
 166
 167         #     var tree = new OrderedTree[Int]
 168         #     tree.add(null, 1)
 169         #     tree.add(1, 11)
 170         #     tree.add(11, 111)
 171         #     tree.add(11, 112)
 172         #     tree.add(1, 12)
 173         #     tree.add(12, 121)
 174         #     tree.add(12, 122)
 175         #     tree.add(null, 2)
 176         #     tree.add(2, 21)
 177         #     tree.add(2, 22)
 178         #     var order = [1, 11, 111, 112, 12, 121, 122, 2, 21, 22]
 179         #     var res = new Array[Int]
 180         #     for i in tree do res.add(i)
 181         #     assert res == order
 182         redef fun iterator do return new OrderedTreeIterator[E](self)
 183 end
 184
 185 # An Iterator over an OrderedTree
 186 private class OrderedTreeIterator[E: Object]
 187         super Iterator[E]
 188
 189         var tree: OrderedTree[E]
 190
 191         var iterators = new Array[Iterator[E]]
 192
 193         init do
 194                 if not tree.is_empty then
 195                         iterators.add tree.roots.iterator
 196                 end
 197         end
 198
 199         redef fun is_ok do return not iterators.is_empty
 200
 201         redef fun item do
 202                 assert is_ok
 203                 return iterators.last.item
 204         end
 205
 206         redef fun next do
 207                 assert is_ok
 208                 if tree.sub.has_key(item) then
 209                         iterators.add tree.sub[item].iterator
 210                 else
 211                         iterators.last.next
 212                         while is_ok and not iterators.last.is_ok do
 213                                 iterators.pop
 214                                 if is_ok and iterators.last.is_ok then
 215                                         iterators.last.next
 216                                 end
 217                         end
 218                 end
 219         end
 220
 221         redef fun iterator do return new OrderedTreeIterator[E](tree)
 222 end