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 # management and display 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 # The class can be used as it to work with generic tree.
  25 # The class can also be specialized to provide more specific behavior.
  26 class OrderedTree[E: Object]
  27         super Streamable
  28         super Collection[E]
  29
  30         # Sequence
  31         var roots = new Array[E]
  32         var sub = new HashMap[E, Array[E]]
  33
  34         # Add a new element `e` in the tree
  35         # `p` is the parent of `e`.
  36         # if `p` is null, then `e` is a root element.
  37         #
  38         # By defauld, the elements with a same parent
  39         # are displayed in the order they are added.
  40         #
  41         # The `sort_with` method can be used reorder elements
  42         fun add(p: nullable E, e: E)
  43         do
  44                 if p == null then
  45                         roots.add(e)
  46                 else if sub.has_key(p) then
  47                         sub[p].add(e)
  48                 else
  49                         sub[p] = [e]
  50                 end
  51         end
  52
  53         # print the full tree on `o`
  54         # Write a ASCII-style tree and use the `display` method to label elements
  55         redef fun write_to(stream: OStream)
  56         do
  57                 var last = roots.last
  58                 for r in roots do
  59                         stream.write display(r)
  60                         stream.write "\n"
  61                         sub_write_to(stream, r, "")
  62                 end
  63         end
  64
  65         private fun sub_write_to(o: OStream, e: E, prefix: String)
  66         do
  67                 if not sub.has_key(e) then return
  68                 var subs = sub[e]
  69                 var last = subs.last
  70                 for e2 in subs do
  71                         if e2 != last then
  72                                 o.write "{prefix}|--{display(e2)}\n"
  73                                 sub_write_to(o, e2, prefix+"|  ")
  74                         else
  75                                 o.write "{prefix}`--{display(e2)}\n"
  76                                 sub_write_to(o, e2, prefix+"   ")
  77                         end
  78                 end
  79         end
  80
  81         # Sort roots and other elements using a comparator method
  82         # This method basically sorts roots then each group of children
  83         fun sort_with(comparator: Comparator)
  84         do
  85                 comparator.sort(roots)
  86                 for a in sub.values do
  87                         comparator.sort(a)
  88                 end
  89         end
  90
  91         # How to display a specific element of the tree
  92         # By defaut, uses `to_s`
  93         #
  94         # Subclasses should redefine this method to provide a specific output
  95         fun display(e: E): String do return e.to_s
  96
  97         # Get an array of the contained elements
  98         # Order is preserved
  99         #
 100         #     var tree = new OrderedTree[Int]
 101         #     tree.add(null, 1)
 102         #     tree.add(1, 11)
 103         #     tree.add(11, 111)
 104         #     tree.add(11, 112)
 105         #     tree.add(1, 12)
 106         #     tree.add(12, 121)
 107         #     tree.add(12, 122)
 108         #     tree.add(null, 2)
 109         #     tree.add(2, 21)
 110         #     tree.add(2, 22)
 111         #     assert tree.to_a == [1, 11, 111, 112, 12, 121, 122, 2, 21, 22]
 112         redef fun to_a: Array[E] do
 113                 var res = new Array[E]
 114                 for r in roots do sub_to_a(r, res)
 115                 return res
 116         end
 117
 118         private fun sub_to_a(e: E, res: Array[E]) do
 119                 res.add e
 120                 if sub.has_key(e) then for e2 in sub[e] do sub_to_a(e2, res)
 121         end
 122
 123         #     var tree = new OrderedTree[Int]
 124         #     assert tree.is_empty
 125         #     tree.add(null, 1)
 126         #     assert not tree.is_empty
 127         redef fun is_empty: Bool do return roots.is_empty
 128
 129         #     var tree = new OrderedTree[Int]
 130         #     tree.add(null, 1)
 131         #     tree.add(1, 11)
 132         #     assert tree.first == 1
 133         redef fun first do return roots.first
 134
 135         #     var tree = new OrderedTree[Int]
 136         #     tree.add(null, 1)
 137         #     tree.add(1, 11)
 138         #     tree.add(11, 111)
 139         #     tree.add(11, 112)
 140         #     tree.add(1, 12)
 141         #     tree.add(12, 121)
 142         #     tree.add(12, 122)
 143         #     tree.add(null, 2)
 144         #     tree.add(2, 21)
 145         #     tree.add(2, 22)
 146         #     var order = [1, 11, 111, 112, 12, 121, 122, 2, 21, 22]
 147         #     var res = new Array[Int]
 148         #     for i in tree do res.add(i)
 149         #     assert res == order
 150         redef fun iterator do return new OrderedTreeIterator[E](self)
 151 end
 152
 153 # An Iterator over an OrderedTree
 154 private class OrderedTreeIterator[E: Object]
 155         super Iterator[E]
 156
 157         var tree: OrderedTree[E]
 158
 159         var iterators = new Array[Iterator[E]]
 160
 161         init(tree: OrderedTree[E]) do
 162                 self.tree = tree
 163
 164                 if not tree.is_empty then
 165                         iterators.add tree.roots.iterator
 166                 end
 167         end
 168
 169         redef fun is_ok do return not iterators.is_empty
 170
 171         redef fun item do
 172                 assert is_ok
 173                 return iterators.last.item
 174         end
 175
 176         redef fun next do
 177                 assert is_ok
 178                 if tree.sub.has_key(item) then
 179                         iterators.add tree.sub[item].iterator
 180                 else
 181                         iterators.last.next
 182                         while is_ok and not iterators.last.is_ok do
 183                                 iterators.pop
 184                                 if is_ok and iterators.last.is_ok then
 185                                         iterators.last.next
 186                                 end
 187                         end
 188                 end
 189         end
 190
 191         redef fun iterator do return new OrderedTreeIterator[E](tree)
 192 end