lib/more_collections.nit

   1 # This file is part of NIT ( http://www.nitlanguage.org ).
   2 #
   3 # Licensed under the Apache License, Version 2.0 (the "License");
   4 # you may not use this file except in compliance with the License.
   5 # You may obtain a copy of the License at
   6 #
   7 #     http://www.apache.org/licenses/LICENSE-2.0
   8 #
   9 # Unless required by applicable law or agreed to in writing, software
  10 # distributed under the License is distributed on an "AS IS" BASIS,
  11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12 # See the License for the specific language governing permissions and
  13 # limitations under the License.
  14
  15 # Highly specific, but useful, collections-related classes.
  16 module more_collections is serialize
  17
  18 import serialization
  19
  20 # Simple way to store an `HashMap[K, Array[V]]`
  21 #
  22 # Unlike standard HashMap, MultiHashMap provides a new
  23 # empty array on the first access on a unknown key.
  24 #
  25 #     var m = new MultiHashMap[String, Char]
  26 #     assert not m.has_key("four")
  27 #     m["four"].add('i')
  28 #     m["four"].add('i')
  29 #     m["four"].add('i')
  30 #     m["four"].add('i')
  31 #     assert m.has_key("four")
  32 #     assert m["four"] == ['i', 'i', 'i', 'i']
  33 #     assert m["zzz"] == new Array[Char]
  34 class MultiHashMap[K, V]
  35         super HashMap[K, Array[V]]
  36
  37         # Add `v` to the array associated with `k`.
  38         # If there is no array associated, then create it.
  39         fun add_one(k: K, v: V)
  40         do
  41                 var x = self.get_or_null(k)
  42                 if x != null then
  43                         x.add(v)
  44                 else
  45                         self[k] = [v]
  46                 end
  47         end
  48
  49         redef fun provide_default_value(key) do
  50                 var res = new Array[V]
  51                 self[key] = res
  52                 return res
  53         end
  54 end
  55
  56 # Simple way to store an `HashMap[K1, HashMap[K2, V]]`
  57 #
  58 # ~~~~
  59 # var hm2 = new HashMap2[Int, String, Float]
  60 # hm2[1, "one"] = 1.0
  61 # hm2[2, "two"] = 2.0
  62 # assert hm2[1, "one"] == 1.0
  63 # assert hm2[2, "not-two"] == null
  64 # ~~~~
  65 class HashMap2[K1, K2, V]
  66
  67         private var level1 = new HashMap[K1, HashMap[K2, V]]
  68
  69         # Return the value associated to the keys `k1` and `k2`.
  70         # Return `null` if no such a value.
  71         fun [](k1: K1, k2: K2): nullable V
  72         do
  73                 var level1 = self.level1
  74                 var level2 = level1.get_or_null(k1)
  75                 if level2 == null then return null
  76                 return level2.get_or_null(k2)
  77         end
  78
  79         # Set `v` the value associated to the keys `k1` and `k2`.
  80         fun []=(k1: K1, k2: K2, v: V)
  81         do
  82                 var level1 = self.level1
  83                 var level2 = level1.get_or_null(k1)
  84                 if level2 == null then
  85                         level2 = new HashMap[K2, V]
  86                         level1[k1] = level2
  87                 end
  88                 level2[k2] = v
  89         end
  90
  91         # Remove the item at `k1` and `k2`
  92         fun remove_at(k1: K1, k2: K2)
  93         do
  94                 var level1 = self.level1
  95                 var level2 = level1.get_or_null(k1)
  96                 if level2 == null then return
  97                 level2.keys.remove(k2)
  98         end
  99
 100         # Remove all items
 101         fun clear do level1.clear
 102 end
 103
 104 # Simple way to store an `HashMap[K1, HashMap[K2, HashMap[K3, V]]]`
 105 #
 106 # ~~~~
 107 # var hm3 = new HashMap3[Int, String, Int, Float]
 108 # hm3[1, "one", 11] = 1.0
 109 # hm3[2, "two", 22] = 2.0
 110 # assert hm3[1, "one", 11] == 1.0
 111 # assert hm3[2, "not-two", 22] == null
 112 # ~~~~
 113 class HashMap3[K1, K2, K3, V]
 114
 115         private var level1 = new HashMap[K1, HashMap2[K2, K3, V]]
 116
 117         # Return the value associated to the keys `k1`, `k2`, and `k3`.
 118         # Return `null` if no such a value.
 119         fun [](k1: K1, k2: K2, k3: K3): nullable V
 120         do
 121                 var level1 = self.level1
 122                 var level2 = level1.get_or_null(k1)
 123                 if level2 == null then return null
 124                 return level2[k2, k3]
 125         end
 126
 127         # Set `v` the value associated to the keys `k1`, `k2`, and `k3`.
 128         fun []=(k1: K1, k2: K2, k3: K3, v: V)
 129         do
 130                 var level1 = self.level1
 131                 var level2 = level1.get_or_null(k1)
 132                 if level2 == null then
 133                         level2 = new HashMap2[K2, K3, V]
 134                         level1[k1] = level2
 135                 end
 136                 level2[k2, k3] = v
 137         end
 138
 139         # Remove the item at `k1`, `k2` and `k3`
 140         fun remove_at(k1: K1, k2: K2, k3: K3)
 141         do
 142                 var level1 = self.level1
 143                 var level2 = level1.get_or_null(k1)
 144                 if level2 == null then return
 145                 level2.remove_at(k2, k3)
 146         end
 147
 148         # Remove all items
 149         fun clear do level1.clear
 150 end
 151
 152 # A map with a default value.
 153 #
 154 # ~~~~
 155 # var dm = new DefaultMap[String, Int](10)
 156 # assert dm["a"] == 10
 157 # ~~~~
 158 #
 159 # The default value is used when the key is not present.
 160 # And getting a default value does not register the key.
 161 #
 162 # ~~~~
 163 # assert dm["a"] == 10
 164 # assert dm.length == 0
 165 # assert dm.has_key("a") == false
 166 # ~~~~
 167 #
 168 # It also means that removed key retrieve the default value.
 169 #
 170 # ~~~~
 171 # dm["a"] = 2
 172 # assert dm["a"] == 2
 173 # dm.keys.remove("a")
 174 # assert dm["a"] == 10
 175 # ~~~~
 176 #
 177 # Warning: the default value is used as is, so using mutable object might
 178 # cause side-effects.
 179 #
 180 # ~~~~
 181 # var dma = new DefaultMap[String, Array[Int]](new Array[Int])
 182 #
 183 # dma["a"].add(65)
 184 # assert dma["a"] == [65]
 185 # assert dma.default == [65]
 186 # assert dma["c"] == [65]
 187 #
 188 # dma["b"] += [66]
 189 # assert dma["b"] == [65, 66]
 190 # assert dma.default == [65]
 191 # ~~~~
 192 class DefaultMap[K, V]
 193         super HashMap[K, V]
 194
 195         # The default value.
 196         var default: V
 197
 198         redef fun provide_default_value(key) do return default
 199 end
 200
 201 # An unrolled linked list
 202 #
 203 # A sequence implemented as a linked list of arrays.
 204 #
 205 # This data structure is similar to the `List` but it can benefit from
 206 # better cache performance, lower data overhead for the nodes metadata and
 207 # it spares the GC to allocate many small nodes.
 208 class UnrolledList[E]
 209         super Sequence[E]
 210
 211         # Desired capacity for each nodes
 212         #
 213         # By default at `32`, it can be increased for very large lists.
 214         #
 215         # It can be modified anytime, but newly created nodes may still be assigned
 216         # the same capacity of old nodes when created by `insert`.
 217         var nodes_length = 32 is writable
 218
 219         private var head_node: UnrolledNode[E] = new UnrolledNode[E](nodes_length)
 220
 221         private var tail_node: UnrolledNode[E] = head_node
 222
 223         redef var length = 0
 224
 225         redef fun clear
 226         do
 227                 head_node = new UnrolledNode[E](nodes_length)
 228                 tail_node = head_node
 229                 length = 0
 230         end
 231
 232         # Out parameter of `node_at`
 233         private var index_within_node = 0
 234
 235         private fun node_at(index: Int): UnrolledNode[E]
 236         do
 237                 assert index >= 0 and index < length
 238
 239                 var node = head_node
 240                 while index >= node.length do
 241                         index -= node.length
 242                         node = node.next.as(not null)
 243                 end
 244
 245                 index_within_node = index
 246                 return node
 247         end
 248
 249         private fun insert_node(node: UnrolledNode[E], prev, next: nullable UnrolledNode[E])
 250         do
 251                 if prev != null then
 252                         prev.next = node
 253                 else head_node = node
 254
 255                 if next != null then
 256                         next.prev = node
 257                 else tail_node = node
 258
 259                 node.next = next
 260                 node.prev = prev
 261         end
 262
 263         redef fun [](index)
 264         do
 265                 var node = node_at(index)
 266                 index = index_within_node + node.head_index
 267                 return node.items[index]
 268         end
 269
 270         redef fun []=(index, value)
 271         do
 272                 var node = node_at(index)
 273                 index = index_within_node + node.head_index
 274                 node.items[index] = value
 275         end
 276
 277         redef fun push(item)
 278         do
 279                 var node = tail_node
 280                 if not node.full then
 281                         if node.tail_index < node.capacity then
 282                                 # There's room at the tail
 283                                 node.tail_index += 1
 284                         else
 285                                 # Move everything over by `d`
 286                                 assert node.head_index > 0
 287                                 var d = node.head_index / 2 + 1
 288                                 node.move_head(node.length, d)
 289                                 for i in d.times do node.items[node.tail_index - i] = null
 290                                 node.head_index -= d
 291                                 node.tail_index += -d+1
 292                         end
 293                         node.items[node.tail_index-1] = item
 294                 else
 295                         # New node!
 296                         node = new UnrolledNode[E](nodes_length)
 297                         insert_node(node, tail_node, null)
 298                         node.tail_index = 1
 299                         node.items[0] = item
 300                 end
 301                 length += 1
 302         end
 303
 304         redef fun unshift(item)
 305         do
 306                 var node = head_node
 307                 if not node.full then
 308                         if node.head_index > 0 then
 309                                 # There's room at the head
 310                                 node.head_index -= 1
 311                         else
 312                                 # Move everything over by `d`
 313                                 assert node.tail_index < node.capacity
 314                                 var d = (node.capacity-node.tail_index) / 2 + 1
 315                                 node.move_tail(0, d)
 316                                 for i in d.times do node.items[node.head_index+i] = null
 317                                 node.head_index += d-1
 318                                 node.tail_index += d
 319                         end
 320                         node.items[node.head_index] = item
 321                 else
 322                         # New node!
 323                         node = new UnrolledNode[E](nodes_length)
 324                         insert_node(node, null, head_node)
 325                         node.head_index = node.capacity-1
 326                         node.tail_index = node.capacity
 327                         node.items[node.capacity-1] = item
 328                 end
 329                 length += 1
 330         end
 331
 332         redef fun pop
 333         do
 334                 assert not_empty
 335
 336                 var node = tail_node
 337                 while node.length == 0 do
 338                         # Delete empty node
 339                         var nullable_node = node.prev
 340                         assert is_not_empty: nullable_node != null
 341                         node = nullable_node
 342                         node.next = null
 343                         self.tail_node = node
 344                 end
 345
 346                 var item = node.items[node.tail_index-1]
 347                 node.tail_index -= 1
 348                 length -= 1
 349                 return item
 350         end
 351
 352         redef fun shift
 353         do
 354                 assert not_empty
 355
 356                 var node = head_node
 357                 while node.length == 0 do
 358                         # Delete empty node
 359                         var nullable_node = node.next
 360                         assert is_not_empty: nullable_node != null
 361                         node = nullable_node
 362                         node.prev = null
 363                         self.head_node = node
 364                 end
 365
 366                 var item = node.items[node.head_index]
 367                 node.head_index += 1
 368                 length -= 1
 369                 return item
 370         end
 371
 372         redef fun insert(item, index)
 373         do
 374                 if index == length then
 375                         push item
 376                         return
 377                 end
 378
 379                 var node = node_at(index)
 380                 index = index_within_node
 381                 if node.full then
 382                         # Move half to a new node
 383                         var new_node = new UnrolledNode[E](nodes_length.max(node.capacity))
 384
 385                         # Plug in the new node
 386                         var next_node = node.next
 387                         insert_node(new_node, node, next_node)
 388
 389                         # Move items at and after `index` to the new node
 390                         var to_displace = node.length-index
 391                         var offset = (new_node.capacity-to_displace)/2
 392                         for i in [0..to_displace[ do
 393                                 new_node.items[offset+i] = node.items[index+i]
 394                                 node.items[index+i] = null
 395                         end
 396                         new_node.head_index = offset
 397                         new_node.tail_index = offset + to_displace
 398                         node.tail_index -= to_displace
 399
 400                         # Store `item`
 401                         if index > node.capacity / 2 then
 402                                 new_node.items[offset-1] = item
 403                                 new_node.head_index -= 1
 404                         else
 405                                 node.items[node.head_index+index] = item
 406                                 node.tail_index += 1
 407                         end
 408                 else
 409                         if node.tail_index < node.capacity then
 410                                 # Move items towards the tail
 411                                 node.move_tail(index, 1)
 412                                 node.tail_index += 1
 413                                 node.items[node.head_index + index] = item
 414                         else
 415                                 # Move items towards the head
 416                                 node.move_head(index, 1)
 417                                 node.items[node.head_index + index-1] = item
 418                                 node.head_index -= 1
 419                         end
 420                 end
 421                 length += 1
 422         end
 423
 424         redef fun remove_at(index)
 425         do
 426                 var node = node_at(index)
 427                 index = index_within_node + node.head_index
 428
 429                 # Shift left all the elements after `index`
 430                 for i in [index+1 .. node.tail_index[ do
 431                         node.items[i-1] = node.items[i]
 432                 end
 433                 node.tail_index -= 1
 434                 node.items[node.tail_index] = null
 435
 436                 length -= 1
 437
 438                 var next_node = node.next
 439                 var prev_node = node.prev
 440                 if node.is_empty and (next_node != null or prev_node != null) then
 441                         # Empty and non-head or tail node, delete
 442                         if next_node != null then
 443                                 next_node.prev = node.prev
 444                         else tail_node = prev_node.as(not null)
 445
 446                         if prev_node != null then
 447                                 prev_node.next = node.next
 448                         else head_node = next_node.as(not null)
 449                 end
 450         end
 451
 452         redef fun iterator do return new UnrolledIterator[E](self)
 453 end
 454
 455 # Node composing an `UnrolledList`
 456 #
 457 # Stores the elements in the `items` array. The elements in the `items` array
 458 # begin at `head_index` and end right before `tail_index`. The data is contiguous,
 459 # but there can be empty cells at the beginning and the end of the array.
 460 private class UnrolledNode[E]
 461
 462         var prev: nullable UnrolledNode[E] = null
 463
 464         var next: nullable UnrolledNode[E] = null
 465
 466         # Desired length of `items`
 467         var capacity: Int
 468
 469         # `Array` of items in this node, filled with `null`
 470         var items = new Array[nullable E].filled_with(null, capacity) is lazy
 471
 472         # Index of the first element in `items`
 473         var head_index = 0
 474
 475         # Index after the last element in `items`
 476         var tail_index = 0
 477
 478         fun length: Int do return tail_index - head_index
 479
 480         fun full: Bool do return length == capacity
 481
 482         fun is_empty: Bool do return tail_index == head_index
 483
 484         # Move towards the head all elements before `index` of `displace` cells
 485         fun move_tail(index, displace: Int)
 486         do
 487                 for i in [tail_index-1..head_index+index].step(-1) do
 488                         items[i+displace] = items[i]
 489                 end
 490         end
 491
 492         # Move towards the tail all elements at and after `index` of `displace` cells
 493         fun move_head(index, displace: Int)
 494         do
 495                 for i in [head_index..head_index+index[ do
 496                         items[i-displace] = items[i]
 497                 end
 498         end
 499 end
 500
 501 private class UnrolledIterator[E]
 502         super IndexedIterator[E]
 503
 504         var list: UnrolledList[E]
 505
 506         var node: nullable UnrolledNode[E] = list.head_node is lazy
 507
 508         # Index of the current `item`
 509         redef var index = 0
 510
 511         # Index within the current `node`
 512         var index_in_node: Int = node.head_index is lazy
 513
 514         redef fun item do return node.items[index_in_node]
 515
 516         redef fun is_ok do return index < list.length
 517
 518         redef fun next
 519         do
 520                 index += 1
 521                 index_in_node += 1
 522
 523                 if index_in_node >= node.tail_index then
 524                         node = node.next
 525                         if node != null then index_in_node = node.head_index
 526                 end
 527         end
 528 end