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                 for r in roots do
  58                         stream.write display(r)
  59                         stream.write "\n"
  60                         sub_write_to(stream, r, "")
  61                 end
  62         end
  63
  64         private fun sub_write_to(o: OStream, e: E, prefix: String)
  65         do
  66                 if not sub.has_key(e) then return
  67                 var subs = sub[e]
  68                 var last = subs.last
  69                 for e2 in subs do
  70                         if e2 != last then
  71                                 o.write "{prefix}|--{display(e2)}\n"
  72                                 sub_write_to(o, e2, prefix+"|  ")
  73                         else
  74                                 o.write "{prefix}`--{display(e2)}\n"
  75                                 sub_write_to(o, e2, prefix+"   ")
  76                         end
  77                 end
  78         end
  79
  80         # Sort roots and other elements using a comparator method
  81         # This method basically sorts roots then each group of children
  82         fun sort_with(comparator: Comparator)
  83         do
  84                 comparator.sort(roots)
  85                 for a in sub.values do
  86                         comparator.sort(a)
  87                 end
  88         end
  89
  90         # How to display a specific element of the tree
  91         # By defaut, uses `to_s`
  92         #
  93         # Subclasses should redefine this method to provide a specific output
  94         fun display(e: E): String do return e.to_s
  95
  96         # Get an array of the contained elements
  97         # Order is preserved
  98         #
  99         #     var tree = new OrderedTree[Int]
 100         #     tree.add(null, 1)
 101         #     tree.add(1, 11)
 102         #     tree.add(11, 111)
 103         #     tree.add(11, 112)
 104         #     tree.add(1, 12)
 105         #     tree.add(12, 121)
 106         #     tree.add(12, 122)
 107         #     tree.add(null, 2)
 108         #     tree.add(2, 21)
 109         #     tree.add(2, 22)
 110         #     assert tree.to_a == [1, 11, 111, 112, 12, 121, 122, 2, 21, 22]
 111         redef fun to_a: Array[E] do
 112                 var res = new Array[E]
 113                 for r in roots do sub_to_a(r, res)
 114                 return res
 115         end
 116
 117         private fun sub_to_a(e: E, res: Array[E]) do
 118                 res.add e
 119                 if sub.has_key(e) then for e2 in sub[e] do sub_to_a(e2, res)
 120         end
 121
 122         #     var tree = new OrderedTree[Int]
 123         #     assert tree.is_empty
 124         #     tree.add(null, 1)
 125         #     assert not tree.is_empty
 126         redef fun is_empty: Bool do return roots.is_empty
 127
 128         #     var tree = new OrderedTree[Int]
 129         #     tree.add(null, 1)
 130         #     tree.add(1, 11)
 131         #     assert tree.first == 1
 132         redef fun first do return roots.first
 133
 134         #     var tree = new OrderedTree[Int]
 135         #     tree.add(null, 1)
 136         #     tree.add(1, 11)
 137         #     tree.add(11, 111)
 138         #     tree.add(11, 112)
 139         #     tree.add(1, 12)
 140         #     tree.add(12, 121)
 141         #     tree.add(12, 122)
 142         #     tree.add(null, 2)
 143         #     tree.add(2, 21)
 144         #     tree.add(2, 22)
 145         #     var order = [1, 11, 111, 112, 12, 121, 122, 2, 21, 22]
 146         #     var res = new Array[Int]
 147         #     for i in tree do res.add(i)
 148         #     assert res == order
 149         redef fun iterator do return new OrderedTreeIterator[E](self)
 150 end
 151
 152 # An Iterator over an OrderedTree
 153 private class OrderedTreeIterator[E: Object]
 154         super Iterator[E]
 155
 156         var tree: OrderedTree[E]
 157
 158         var iterators = new Array[Iterator[E]]
 159
 160         init do
 161                 if not tree.is_empty then
 162                         iterators.add tree.roots.iterator
 163                 end
 164         end
 165
 166         redef fun is_ok do return not iterators.is_empty
 167
 168         redef fun item do
 169                 assert is_ok
 170                 return iterators.last.item
 171         end
 172
 173         redef fun next do
 174                 assert is_ok
 175                 if tree.sub.has_key(item) then
 176                         iterators.add tree.sub[item].iterator
 177                 else
 178                         iterators.last.next
 179                         while is_ok and not iterators.last.is_ok do
 180                                 iterators.pop
 181                                 if is_ok and iterators.last.is_ok then
 182                                         iterators.last.next
 183                                 end
 184                         end
 185                 end
 186         end
 187
 188         redef fun iterator do return new OrderedTreeIterator[E](tree)
 189 end