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