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 import poset
  20
  21 # Simple way to store an `HashMap[K, Array[V]]`
  22 #
  23 # Unlike standard HashMap, MultiHashMap provides a new
  24 # empty array on the first access on a unknown key.
  25 #
  26 #     var m = new MultiHashMap[String, Char]
  27 #     assert not m.has_key("four")
  28 #     m["four"].add('i')
  29 #     m["four"].add('i')
  30 #     m["four"].add('i')
  31 #     m["four"].add('i')
  32 #     assert m.has_key("four")
  33 #     assert m["four"] == ['i', 'i', 'i', 'i']
  34 #     assert m["zzz"] == new Array[Char]
  35 class MultiHashMap[K, V]
  36         super HashMap[K, Array[V]]
  37
  38         # Add `v` to the array associated with `k`.
  39         #
  40         # If there is no array associated, then create it.
  41         #
  42         # For the inverse operation, see `remove_one`.
  43         #
  44         # ```
  45         # var m = new MultiHashMap[String, Char]
  46         # m.add_one("four", 'i')
  47         # m.add_one("four", 'i')
  48         # m.add_one("four", 'i')
  49         # m.add_one("four", 'i')
  50         # assert m.has_key("four")
  51         # assert m["four"] == ['i', 'i', 'i', 'i']
  52         # ```
  53         fun add_one(k: K, v: V)
  54         do
  55                 var x = self.get_or_null(k)
  56                 if x != null then
  57                         x.add(v)
  58                 else
  59                         self[k] = [v]
  60                 end
  61         end
  62
  63         redef fun provide_default_value(key) do
  64                 var res = new Array[V]
  65                 self[key] = res
  66                 return res
  67         end
  68
  69         # Remove an occurrence of `v` from the array associated with `k`.
  70         #
  71         # If the associated array does not contain `v`, do nothing. If the
  72         # associated array only contain one element and this element is `v`, remove
  73         # the key `k`.
  74         #
  75         # In a nutshell, does the inverse operation of `add_one`.
  76         #
  77         # ```
  78         # var m = new MultiHashMap[String, Char]
  79         # m["four"]        =  ['4', 'i', 'i', 'i', 'i']
  80         # m.remove_one("four", 'i')
  81         # assert m["four"] == ['4', 'i', 'i', 'i']
  82         #
  83         # m = new MultiHashMap[String, Char]
  84         # m.add_one("one", '1')
  85         # m.remove_one("one", '?')
  86         # assert m["one"] == ['1']
  87         # m.remove_one("one", '1')
  88         # assert not m.has_key("one")
  89         # assert m["one"] == new Array[Char]
  90         #
  91         # m = new MultiHashMap[String, Char]
  92         # m.add_one("one", '1')
  93         # m.remove_one("two", '2')
  94         # assert not m.has_key("two")
  95         # assert m["one"] == ['1']
  96         # assert m["two"] == new Array[Char]
  97         # ```
  98         fun remove_one(k: K, v: V)
  99         do
 100                 var x = get_or_null(k)
 101                 if x != null then
 102                         x.remove(v)
 103                         if x.is_empty then keys.remove(k)
 104                 end
 105         end
 106
 107         # Search the values in `pe.greaters` from the most smaller elements that have a value.
 108         #
 109         # Elements without values are ignored.
 110         #
 111         # Basically, values defined in nearest greater elements of `pe` are inherited.
 112         #
 113         # ~~~
 114         # var pos = new POSet[String]
 115         # pos.add_chain(["E", "D", "C", "B", "A"])
 116         # pos.add_chain(["D", "X", "B"])
 117         #
 118         # var map = new MultiHashMap[String, String]
 119         # map["A"].append(["a", "1"])
 120         # map["C"].append(["c", "2"])
 121         # map["X"].append(["x", "2"])
 122         # map["E"].add "e"
 123         #
 124         # assert map.lookup_joined_values(pos["B"]).has_exactly(["a", "1"])
 125         # assert map.lookup_joined_values(pos["C"]).has_exactly(["c", "2"])
 126         # assert map.lookup_joined_values(pos["D"]).has_exactly(["c", "x", "2"])
 127         # ~~~
 128         fun lookup_joined_values(pe: POSetElement[K]): Set[V]
 129         do
 130                 var res = new Set[V]
 131                 for k in pe.poset.select_smallest(filter_keys(pe.greaters)) do res.add_all self[k]
 132                 return res
 133
 134         end
 135 end
 136
 137 # Simple way to store an `HashMap[K1, HashMap[K2, V]]`
 138 #
 139 # ~~~~
 140 # var hm2 = new HashMap2[Int, String, Float]
 141 # hm2[1, "one"] = 1.0
 142 # hm2[2, "two"] = 2.0
 143 # assert hm2[1, "one"] == 1.0
 144 # assert hm2[2, "not-two"] == null
 145 # ~~~~
 146 class HashMap2[K1, K2, V]
 147
 148         private var level1 = new HashMap[K1, HashMap[K2, V]]
 149
 150         # Return the value associated to the keys `k1` and `k2`.
 151         # Return `null` if no such a value.
 152         fun [](k1: K1, k2: K2): nullable V
 153         do
 154                 var level1 = self.level1
 155                 var level2 = level1.get_or_null(k1)
 156                 if level2 == null then return null
 157                 return level2.get_or_null(k2)
 158         end
 159
 160         # Set `v` the value associated to the keys `k1` and `k2`.
 161         fun []=(k1: K1, k2: K2, v: V)
 162         do
 163                 var level1 = self.level1
 164                 var level2 = level1.get_or_null(k1)
 165                 if level2 == null then
 166                         level2 = new HashMap[K2, V]
 167                         level1[k1] = level2
 168                 end
 169                 level2[k2] = v
 170         end
 171
 172         # Remove the item at `k1` and `k2`
 173         fun remove_at(k1: K1, k2: K2)
 174         do
 175                 var level1 = self.level1
 176                 var level2 = level1.get_or_null(k1)
 177                 if level2 == null then return
 178                 level2.keys.remove(k2)
 179         end
 180
 181         # Is there a value at `k1, k2`?
 182         fun has(k1: K1, k2: K2): Bool
 183         do
 184                 if not level1.keys.has(k1) then return false
 185                 return level1[k1].keys.has(k2)
 186         end
 187
 188         # Remove all items
 189         fun clear do level1.clear
 190 end
 191
 192 # Simple way to store an `HashMap[K1, HashMap[K2, HashMap[K3, V]]]`
 193 #
 194 # ~~~~
 195 # var hm3 = new HashMap3[Int, String, Int, Float]
 196 # hm3[1, "one", 11] = 1.0
 197 # hm3[2, "two", 22] = 2.0
 198 # assert hm3[1, "one", 11] == 1.0
 199 # assert hm3[2, "not-two", 22] == null
 200 # ~~~~
 201 class HashMap3[K1, K2, K3, V]
 202
 203         private var level1 = new HashMap[K1, HashMap2[K2, K3, V]]
 204
 205         # Return the value associated to the keys `k1`, `k2`, and `k3`.
 206         # Return `null` if no such a value.
 207         fun [](k1: K1, k2: K2, k3: K3): nullable V
 208         do
 209                 var level1 = self.level1
 210                 var level2 = level1.get_or_null(k1)
 211                 if level2 == null then return null
 212                 return level2[k2, k3]
 213         end
 214
 215         # Set `v` the value associated to the keys `k1`, `k2`, and `k3`.
 216         fun []=(k1: K1, k2: K2, k3: K3, v: V)
 217         do
 218                 var level1 = self.level1
 219                 var level2 = level1.get_or_null(k1)
 220                 if level2 == null then
 221                         level2 = new HashMap2[K2, K3, V]
 222                         level1[k1] = level2
 223                 end
 224                 level2[k2, k3] = v
 225         end
 226
 227         # Remove the item at `k1`, `k2` and `k3`
 228         fun remove_at(k1: K1, k2: K2, k3: K3)
 229         do
 230                 var level1 = self.level1
 231                 var level2 = level1.get_or_null(k1)
 232                 if level2 == null then return
 233                 level2.remove_at(k2, k3)
 234         end
 235
 236         # Is there a value at `k1, k2, k3`?
 237         fun has(k1: K1, k2: K2, k3: K3): Bool
 238         do
 239                 if not level1.keys.has(k1) then return false
 240                 return level1[k1].has(k2, k3)
 241         end
 242
 243         # Remove all items
 244         fun clear do level1.clear
 245 end
 246
 247 # A map with a default value.
 248 #
 249 # ~~~~
 250 # var dm = new DefaultMap[String, Int](10)
 251 # assert dm["a"] == 10
 252 # ~~~~
 253 #
 254 # The default value is used when the key is not present.
 255 # And getting a default value does not register the key.
 256 #
 257 # ~~~~
 258 # assert dm["a"] == 10
 259 # assert dm.length == 0
 260 # assert dm.has_key("a") == false
 261 # ~~~~
 262 #
 263 # It also means that removed key retrieve the default value.
 264 #
 265 # ~~~~
 266 # dm["a"] = 2
 267 # assert dm["a"] == 2
 268 # dm.keys.remove("a")
 269 # assert dm["a"] == 10
 270 # ~~~~
 271 #
 272 # Warning: the default value is used as is, so using mutable object might
 273 # cause side-effects.
 274 #
 275 # ~~~~
 276 # var dma = new DefaultMap[String, Array[Int]](new Array[Int])
 277 #
 278 # dma["a"].add(65)
 279 # assert dma["a"] == [65]
 280 # assert dma.default == [65]
 281 # assert dma["c"] == [65]
 282 #
 283 # dma["b"] += [66]
 284 # assert dma["b"] == [65, 66]
 285 # assert dma.default == [65]
 286 # ~~~~
 287 class DefaultMap[K, V]
 288         super HashMap[K, V]
 289
 290         # The default value.
 291         var default: V
 292
 293         redef fun provide_default_value(key) do return default
 294 end
 295
 296 # An unrolled linked list
 297 #
 298 # A sequence implemented as a linked list of arrays.
 299 #
 300 # This data structure is similar to the `List` but it can benefit from
 301 # better cache performance, lower data overhead for the nodes metadata and
 302 # it spares the GC to allocate many small nodes.
 303 class UnrolledList[E]
 304         super Sequence[E]
 305
 306         # Desired capacity for each nodes
 307         #
 308         # By default at `32`, it can be increased for very large lists.
 309         #
 310         # It can be modified anytime, but newly created nodes may still be assigned
 311         # the same capacity of old nodes when created by `insert`.
 312         var nodes_length = 32 is writable
 313
 314         private var head_node: UnrolledNode[E] = new UnrolledNode[E](nodes_length)
 315
 316         private var tail_node: UnrolledNode[E] = head_node
 317
 318         redef var length = 0
 319
 320         redef fun clear
 321         do
 322                 head_node = new UnrolledNode[E](nodes_length)
 323                 tail_node = head_node
 324                 length = 0
 325         end
 326
 327         # Out parameter of `node_at`
 328         private var index_within_node = 0
 329
 330         private fun node_at(index: Int): UnrolledNode[E]
 331         do
 332                 assert index >= 0 and index < length
 333
 334                 var node = head_node
 335                 while index >= node.length do
 336                         index -= node.length
 337                         node = node.next.as(not null)
 338                 end
 339
 340                 index_within_node = index
 341                 return node
 342         end
 343
 344         private fun insert_node(node: UnrolledNode[E], prev, next: nullable UnrolledNode[E])
 345         do
 346                 if prev != null then
 347                         prev.next = node
 348                 else head_node = node
 349
 350                 if next != null then
 351                         next.prev = node
 352                 else tail_node = node
 353
 354                 node.next = next
 355                 node.prev = prev
 356         end
 357
 358         redef fun [](index)
 359         do
 360                 var node = node_at(index)
 361                 index = index_within_node + node.head_index
 362                 return node.items[index]
 363         end
 364
 365         redef fun []=(index, value)
 366         do
 367                 var node = node_at(index)
 368                 index = index_within_node + node.head_index
 369                 node.items[index] = value
 370         end
 371
 372         redef fun push(item)
 373         do
 374                 var node = tail_node
 375                 if not node.full then
 376                         if node.tail_index < node.capacity then
 377                                 # There's room at the tail
 378                                 node.tail_index += 1
 379                         else
 380                                 # Move everything over by `d`
 381                                 assert node.head_index > 0
 382                                 var d = node.head_index / 2 + 1
 383                                 node.move_head(node.length, d)
 384                                 for i in d.times do node.items[node.tail_index - i] = null
 385                                 node.head_index -= d
 386                                 node.tail_index += -d+1
 387                         end
 388                         node.items[node.tail_index-1] = item
 389                 else
 390                         # New node!
 391                         node = new UnrolledNode[E](nodes_length)
 392                         insert_node(node, tail_node, null)
 393                         node.tail_index = 1
 394                         node.items[0] = item
 395                 end
 396                 length += 1
 397         end
 398
 399         redef fun unshift(item)
 400         do
 401                 var node = head_node
 402                 if not node.full then
 403                         if node.head_index > 0 then
 404                                 # There's room at the head
 405                                 node.head_index -= 1
 406                         else
 407                                 # Move everything over by `d`
 408                                 assert node.tail_index < node.capacity
 409                                 var d = (node.capacity-node.tail_index) / 2 + 1
 410                                 node.move_tail(0, d)
 411                                 for i in d.times do node.items[node.head_index+i] = null
 412                                 node.head_index += d-1
 413                                 node.tail_index += d
 414                         end
 415                         node.items[node.head_index] = item
 416                 else
 417                         # New node!
 418                         node = new UnrolledNode[E](nodes_length)
 419                         insert_node(node, null, head_node)
 420                         node.head_index = node.capacity-1
 421                         node.tail_index = node.capacity
 422                         node.items[node.capacity-1] = item
 423                 end
 424                 length += 1
 425         end
 426
 427         redef fun pop
 428         do
 429                 assert not_empty
 430
 431                 var node = tail_node
 432                 while node.length == 0 do
 433                         # Delete empty node
 434                         var nullable_node = node.prev
 435                         assert is_not_empty: nullable_node != null
 436                         node = nullable_node
 437                         node.next = null
 438                         self.tail_node = node
 439                 end
 440
 441                 var item = node.items[node.tail_index-1]
 442                 node.tail_index -= 1
 443                 length -= 1
 444                 return item
 445         end
 446
 447         redef fun shift
 448         do
 449                 assert not_empty
 450
 451                 var node = head_node
 452                 while node.length == 0 do
 453                         # Delete empty node
 454                         var nullable_node = node.next
 455                         assert is_not_empty: nullable_node != null
 456                         node = nullable_node
 457                         node.prev = null
 458                         self.head_node = node
 459                 end
 460
 461                 var item = node.items[node.head_index]
 462                 node.head_index += 1
 463                 length -= 1
 464                 return item
 465         end
 466
 467         redef fun insert(item, index)
 468         do
 469                 if index == length then
 470                         push item
 471                         return
 472                 end
 473
 474                 var node = node_at(index)
 475                 index = index_within_node
 476                 if node.full then
 477                         # Move half to a new node
 478                         var new_node = new UnrolledNode[E](nodes_length.max(node.capacity))
 479
 480                         # Plug in the new node
 481                         var next_node = node.next
 482                         insert_node(new_node, node, next_node)
 483
 484                         # Move items at and after `index` to the new node
 485                         var to_displace = node.length-index
 486                         var offset = (new_node.capacity-to_displace)/2
 487                         for i in [0..to_displace[ do
 488                                 new_node.items[offset+i] = node.items[index+i]
 489                                 node.items[index+i] = null
 490                         end
 491                         new_node.head_index = offset
 492                         new_node.tail_index = offset + to_displace
 493                         node.tail_index -= to_displace
 494
 495                         # Store `item`
 496                         if index > node.capacity / 2 then
 497                                 new_node.items[offset-1] = item
 498                                 new_node.head_index -= 1
 499                         else
 500                                 node.items[node.head_index+index] = item
 501                                 node.tail_index += 1
 502                         end
 503                 else
 504                         if node.tail_index < node.capacity then
 505                                 # Move items towards the tail
 506                                 node.move_tail(index, 1)
 507                                 node.tail_index += 1
 508                                 node.items[node.head_index + index] = item
 509                         else
 510                                 # Move items towards the head
 511                                 node.move_head(index, 1)
 512                                 node.items[node.head_index + index-1] = item
 513                                 node.head_index -= 1
 514                         end
 515                 end
 516                 length += 1
 517         end
 518
 519         redef fun remove_at(index)
 520         do
 521                 var node = node_at(index)
 522                 index = index_within_node + node.head_index
 523
 524                 # Shift left all the elements after `index`
 525                 for i in [index+1 .. node.tail_index[ do
 526                         node.items[i-1] = node.items[i]
 527                 end
 528                 node.tail_index -= 1
 529                 node.items[node.tail_index] = null
 530
 531                 length -= 1
 532
 533                 var next_node = node.next
 534                 var prev_node = node.prev
 535                 if node.is_empty and (next_node != null or prev_node != null) then
 536                         # Empty and non-head or tail node, delete
 537                         if next_node != null then
 538                                 next_node.prev = node.prev
 539                         else tail_node = prev_node.as(not null)
 540
 541                         if prev_node != null then
 542                                 prev_node.next = node.next
 543                         else head_node = next_node.as(not null)
 544                 end
 545         end
 546
 547         redef fun iterator do return new UnrolledIterator[E](self)
 548 end
 549
 550 # Node composing an `UnrolledList`
 551 #
 552 # Stores the elements in the `items` array. The elements in the `items` array
 553 # begin at `head_index` and end right before `tail_index`. The data is contiguous,
 554 # but there can be empty cells at the beginning and the end of the array.
 555 private class UnrolledNode[E]
 556
 557         var prev: nullable UnrolledNode[E] = null
 558
 559         var next: nullable UnrolledNode[E] = null
 560
 561         # Desired length of `items`
 562         var capacity: Int
 563
 564         # `Array` of items in this node, filled with `null`
 565         var items = new Array[nullable E].filled_with(null, capacity) is lazy
 566
 567         # Index of the first element in `items`
 568         var head_index = 0
 569
 570         # Index after the last element in `items`
 571         var tail_index = 0
 572
 573         fun length: Int do return tail_index - head_index
 574
 575         fun full: Bool do return length == capacity
 576
 577         fun is_empty: Bool do return tail_index == head_index
 578
 579         # Move towards the head all elements before `index` of `displace` cells
 580         fun move_tail(index, displace: Int)
 581         do
 582                 for i in [tail_index-1..head_index+index].step(-1) do
 583                         items[i+displace] = items[i]
 584                 end
 585         end
 586
 587         # Move towards the tail all elements at and after `index` of `displace` cells
 588         fun move_head(index, displace: Int)
 589         do
 590                 for i in [head_index..head_index+index[ do
 591                         items[i-displace] = items[i]
 592                 end
 593         end
 594 end
 595
 596 private class UnrolledIterator[E]
 597         super IndexedIterator[E]
 598
 599         var list: UnrolledList[E]
 600
 601         var node: nullable UnrolledNode[E] = list.head_node is lazy
 602
 603         # Index of the current `item`
 604         redef var index = 0
 605
 606         # Index within the current `node`
 607         var index_in_node: Int = node.head_index is lazy
 608
 609         redef fun item do return node.items[index_in_node]
 610
 611         redef fun is_ok do return index < list.length
 612
 613         redef fun next
 614         do
 615                 index += 1
 616                 index_in_node += 1
 617
 618                 if index_in_node >= node.tail_index then
 619                         node = node.next
 620                         if node != null then index_in_node = node.head_index
 621                 end
 622         end
 623 end
 624
 625 # Keep track of the best elements according to a distance value.
 626 #
 627 # ~~~
 628 # var bests = new BestDistance[String](5)
 629 # bests.update(10, "Too big")
 630 # assert bests.best_items.is_empty
 631 # bests.update(5, "Just fine")
 632 # bests.update(5, "Another one")
 633 # assert bests.best_items.has_exactly(["Just fine", "Another one"])
 634 # bests.update(2, "A better one")
 635 # bests.update(4, "Not good enough")
 636 # assert bests.best_distance == 2
 637 # assert bests.best_items.has_exactly(["A better one"])
 638 # ~~~
 639 class BestDistance[E]
 640         # Current smallest distance
 641         var best_distance: Int is writable
 642
 643         # Known elements with the smallest distance
 644         var best_items = new Set[E] is writable
 645
 646         # Register a `candidate` with a `distance`
 647         #
 648         # * To high, it is ignored.
 649         # * Equal to the current best, it is added
 650         # * Better that them, is is the new best element
 651         #
 652         # Return `true` if the candidate is kept (alone or with other)
 653         # returns `false` if the candidate is ignored.
 654         fun update(distance: Int, candidate: E): Bool
 655         do
 656                 if distance > best_distance then return false
 657                 if distance < best_distance then
 658                         best_distance = distance
 659                         best_items.clear
 660                 end
 661                 best_items.add candidate
 662                 return true
 663         end
 664 end