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         super Cloneable
  65
  66         # The roots of the tree (in sequence)
  67         var roots = new Array[E]
  68
  69         # The branches of the trees.
  70         # For each element, the ordered array of its direct sub-elements.
  71         var sub = new HashMap[E, Array[E]]
  72
  73         # The parent of each element.
  74         #
  75         # Roots have `null` as parent.
  76         private var parents = new HashMap[E, nullable E]
  77
  78         redef fun length do return parents.length
  79
  80         redef fun has(e) do return parents.has_key(e)
  81
  82         # The parent of the element `e`
  83         #
  84         # Roots have `null` as parent.
  85         #
  86         # ~~~
  87         # var tree = new OrderedTree[Int]
  88         # tree.add(1, 2)
  89         # assert tree.parent(2) == 1
  90         # assert tree.parent(1) == null
  91         # ~~~
  92         fun parent(e: E): nullable E do return parents[e]
  93
  94         # Add a new element `e` in the tree.
  95         #
  96         # `p` is the parent of `e`.
  97         # If `p` is null, then `e` is a root element.
  98         #
  99         # If `e` is already in the tree, it is detached from its old
 100         # parent and attached to the new parent `p`.
 101         fun add(p: nullable E, e: E)
 102         do
 103                 detach(e)
 104                 parents[e] = p
 105                 if p == null then
 106                         roots.add(e)
 107                 else
 108                         if not has(p) then add(null, p)
 109                         if sub.has_key(p) then
 110                                 sub[p].add(e)
 111                         else
 112                                 sub[p] = [e]
 113                         end
 114                 end
 115         end
 116
 117         # Append all nodes `es` as children of `p`.
 118         fun add_all(p: nullable E, es: Collection[E])
 119         do
 120                 for e in es do add(p, e)
 121         end
 122
 123         # Temporary remove `e`.
 124         #
 125         # Children of `e` are left untouched in the tree.
 126         # This make the tree inconstant until `e` is added back.
 127         private fun detach(e: E)
 128         do
 129                 var old_parent = parents.get_or_null(e)
 130                 if old_parent != null then
 131                         var subs = sub[old_parent]
 132                         subs.remove(e)
 133                         if subs.is_empty then
 134                                 # remove the sub when all children are detached
 135                                 # so that `==` and `hash` are sane
 136                                 # Otherwise an empty array will be considered
 137                                 # differently than no array.
 138                                 sub.keys.remove(old_parent)
 139                         end
 140                 else if roots.has(e) then
 141                         roots.remove(e)
 142                 end
 143         end
 144
 145         # print the full tree on `o`
 146         # Write a ASCII-style tree and use the `display` method to label elements
 147         redef fun write_to(stream: Writer)
 148         do
 149                 for r in roots do
 150                         write_line(stream, r, "")
 151                         sub_write_to(stream, r, "")
 152                 end
 153         end
 154
 155         private fun sub_write_to(o: Writer, e: E, prefix: String)
 156         do
 157                 if not sub.has_key(e) then return
 158                 var subs = sub[e]
 159                 if subs.is_empty then return
 160                 var last = subs.last
 161                 for e2 in subs do
 162                         if e2 != last then
 163                                 write_line(o, e2, prefix+"|--")
 164                                 sub_write_to(o, e2, prefix+"|  ")
 165                         else
 166                                 write_line(o, e2, prefix+"`--")
 167                                 sub_write_to(o, e2, prefix+"   ")
 168                         end
 169                 end
 170         end
 171
 172         # Write the full line for the element `e` in `o`.
 173         #
 174         # Basically it does:
 175         #
 176         # ~~~nitish
 177         # o.write "{prefix}{display(e)}\n"
 178         # ~~~
 179         #
 180         # Usually, you should redefine `display` to change the display of an element.
 181         protected fun write_line(o: Writer, e: E, prefix: String)
 182         do
 183                 o.write "{prefix}{display(e)}\n"
 184         end
 185
 186         # Sort roots and other elements using a comparator method
 187         # This method basically sorts roots then each group of children
 188         fun sort_with(comparator: Comparator)
 189         do
 190                 comparator.sort(roots)
 191                 for a in sub.values do
 192                         comparator.sort(a)
 193                 end
 194         end
 195
 196         # How to display a specific element of the tree
 197         # By defaut, uses `to_s`
 198         #
 199         # Subclasses should redefine this method to provide a specific output
 200         fun display(e: E): String do return e.to_s
 201
 202         # Get an array of the contained elements
 203         # Order is preserved
 204         #
 205         #     var tree = new OrderedTree[Int]
 206         #     tree.add_all(null, [1, 2])
 207         #     tree.add_all(1, [11, 12])
 208         #     tree.add_all(11, [111, 112])
 209         #     tree.add_all(12, [121, 122])
 210         #     tree.add_all(2, [21, 22])
 211         #     assert tree.to_a == [1, 11, 111, 112, 12, 121, 122, 2, 21, 22]
 212         redef fun to_a: Array[E] do
 213                 var res = new Array[E]
 214                 for r in roots do sub_to_a(r, res)
 215                 return res
 216         end
 217
 218         private fun sub_to_a(e: E, res: Array[E]) do
 219                 res.add e
 220                 if sub.has_key(e) then for e2 in sub[e] do sub_to_a(e2, res)
 221         end
 222
 223         #     var tree = new OrderedTree[Int]
 224         #     assert tree.is_empty
 225         #     tree.add(null, 1)
 226         #     assert not tree.is_empty
 227         redef fun is_empty: Bool do return roots.is_empty
 228
 229         #     var tree = new OrderedTree[Int]
 230         #     tree.add(null, 1)
 231         #     tree.add(1, 11)
 232         #     assert tree.first == 1
 233         redef fun first do return roots.first
 234
 235         #     var tree = new OrderedTree[Int]
 236         #     tree.add_all(null, [1, 2])
 237         #     tree.add_all(1, [11, 12])
 238         #     tree.add_all(11, [111, 112])
 239         #     tree.add_all(12, [121, 122])
 240         #     tree.add_all(2, [21, 22])
 241         #     var order = [1, 11, 111, 112, 12, 121, 122, 2, 21, 22]
 242         #     assert tree.iterator.to_a == order
 243         redef fun iterator do return new OrderedTreeIterator[E](self)
 244
 245         # Two trees are equal if they have the same nodes in the same order
 246         #
 247         # ~~~
 248         # var t1 = new OrderedTree[Int]
 249         # t1.add_all(null, [1, 2])
 250         # t1.add_all(1, [11, 12])
 251         #
 252         # var t2 = new OrderedTree[Int]
 253         # t2.add_all(null, [1, 2])
 254         #
 255         # assert t1 != t2
 256         #
 257         # t2.add_all(1, [11, 12])
 258         #
 259         # assert t1 == t2
 260         # ~~~
 261         redef fun ==(other)
 262         do
 263                 if not other isa OrderedTree[Object] then return false
 264                 return roots == other.roots and sub == other.sub
 265         end
 266
 267         redef fun hash
 268         do
 269                 return roots.hash + sub.hash
 270         end
 271
 272         # Shallow clone of the tree.
 273         #
 274         # ~~~
 275         # var t = new OrderedTree[Int]
 276         # t.add_all(null, [1, 2])
 277         # t.add_all(1, [11, 12])
 278         #
 279         # assert t.clone == t
 280         # ~~~
 281         redef fun clone
 282         do
 283                 var res = new OrderedTree[E]
 284                 res.add_all(null, roots)
 285                 for p, es in sub do
 286                         res.add_all(p, es)
 287                 end
 288                 return res
 289         end
 290 end
 291
 292 # An Iterator over an OrderedTree
 293 private class OrderedTreeIterator[E: Object]
 294         super Iterator[E]
 295
 296         var tree: OrderedTree[E]
 297
 298         var iterators = new Array[Iterator[E]]
 299
 300         init do
 301                 if not tree.is_empty then
 302                         iterators.add tree.roots.iterator
 303                 end
 304         end
 305
 306         redef fun is_ok do return not iterators.is_empty
 307
 308         redef fun item do
 309                 assert is_ok
 310                 return iterators.last.item
 311         end
 312
 313         redef fun next do
 314                 assert is_ok
 315                 if tree.sub.has_key(item) then
 316                         iterators.add tree.sub[item].iterator
 317                 else
 318                         iterators.last.next
 319                         while is_ok and not iterators.last.is_ok do
 320                                 iterators.pop
 321                                 if is_ok and iterators.last.is_ok then
 322                                         iterators.last.next
 323                                 end
 324                         end
 325                 end
 326         end
 327
 328         redef fun iterator do return new OrderedTreeIterator[E](tree)
 329 end