rbtree-ruby 0.3.6 → 0.4.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
data/lib/rbtree.rb CHANGED
@@ -54,6 +54,29 @@ require_relative "rbtree/version"
54
54
  class RBTree
55
55
  include Enumerable
56
56
 
57
+ # Key classes where <tt>(a <=> b) == 0</tt> and <tt>a.eql?(b)</tt> agree for any two
58
+ # instances of that same class, letting {#find_node} skip its fallback descent.
59
+ COHERENT_KEY_CLASSES = [
60
+ Integer, Float, Rational, String, Symbol, Time,
61
+ ].to_h { |klass| [klass, true] }
62
+ private_constant :COHERENT_KEY_CLASSES
63
+
64
+ # Declares that <tt>(a <=> b) == 0</tt> and +a.eql?(b)+ agree for instances of the
65
+ # given key class, speeding up lookups that miss the internal hash index.
66
+ #
67
+ # Optimization hint only: undeclared classes are still handled correctly, at
68
+ # O(log n) per missed lookup. Declaring a class that does not satisfy this makes
69
+ # such lookups report a present key as absent.
70
+ #
71
+ # @param klass [Class] the key class to declare coherent
72
+ # @return [Class] the declared class
73
+ # @example
74
+ # RBTree.coherent_key_class(Version)
75
+ def self.coherent_key_class(klass)
76
+ COHERENT_KEY_CLASSES[klass] = true
77
+ klass
78
+ end
79
+
57
80
  # Returns the number of key-value pairs stored in the tree.
58
81
  # @return [Integer] the number of entries in the tree
59
82
  attr_reader :key_count
@@ -100,9 +123,12 @@ class RBTree
100
123
  @root = @nil_node
101
124
  @min_node = @nil_node
102
125
  @max_node = @nil_node
103
- @hash_index = {} # Hash index for O(1) key lookup
126
+ @hash_index = {} # Hash index for O(1) key lookup, one entry per node
104
127
  @node_allocator = node_allocator
105
128
  @key_count = 0
129
+ @mod_count = 0 # bumped by changes a traversal can observe
130
+ @key_class = nil # single class of all keys, or false once mixed
131
+ @coherent_keys = false
106
132
 
107
133
  @overwrite = overwrite
108
134
 
@@ -114,17 +140,25 @@ class RBTree
114
140
  # Creates a deep copy of the tree.
115
141
  # Called automatically by `dup` and `clone`.
116
142
  #
143
+ # The copy shares the original's node allocator instance.
144
+ #
117
145
  # @param orig [RBTree] the original tree to copy
118
146
  # @return [void]
119
147
  def initialize_copy(orig)
120
- initialize(overwrite: orig.instance_variable_get(:@overwrite))
148
+ initialize(
149
+ overwrite: orig.instance_variable_get(:@overwrite),
150
+ node_allocator: orig.instance_variable_get(:@node_allocator))
121
151
  orig.each { |k, v| insert(k, v) }
122
152
  end
123
153
 
124
- # Returns a Hash containing all key-value pairs from the tree.
154
+ # Returns a Hash containing all key-value pairs from the tree, in ascending key order.
125
155
  #
126
156
  # @return [Hash] a new Hash with the tree's contents
127
- def to_h = @hash_index.transform_values(&:value)
157
+ def to_h
158
+ h = {}
159
+ each { |k, v| h[k] = v }
160
+ h
161
+ end
128
162
 
129
163
  # Checks if the tree is empty.
130
164
  #
@@ -146,11 +180,19 @@ class RBTree
146
180
 
147
181
  # Returns the minimum key-value pair without removing it.
148
182
  #
149
- # @return [Array(Object, Object), nil] a two-element array [key, value], or nil if tree is empty
183
+ # With no argument and no block this is an O(1) cached lookup; given a count or a
184
+ # comparison block, `Enumerable#min` semantics apply instead.
185
+ #
186
+ # @param args [Integer] optional number of smallest pairs to return
187
+ # @return [Array, nil] the pair, or the n smallest pairs, or nil if tree is empty
150
188
  # @example
151
189
  # tree = RBTree.new({3 => 'three', 1 => 'one', 2 => 'two'})
152
- # tree.min # => [1, "one"]
153
- def min = min_node&.pair
190
+ # tree.min # => [1, "one"]
191
+ # tree.min(2) # => [[1, "one"], [2, "two"]]
192
+ def min(*args, &block)
193
+ return super if !args.empty? || block
194
+ min_node&.pair
195
+ end
154
196
 
155
197
  # Returns the maximum key without removing it.
156
198
  #
@@ -162,37 +204,59 @@ class RBTree
162
204
 
163
205
  # Returns the maximum key-value pair without removing it.
164
206
  #
165
- # @return [Array(Object, Object), nil] a two-element array [key, value], or nil if tree is empty
207
+ # With no argument and no block this is an O(1) cached lookup; given a count or a
208
+ # comparison block, `Enumerable#max` semantics apply instead.
209
+ #
210
+ # @param args [Integer] optional number of largest pairs to return
211
+ # @return [Array, nil] the pair, or the n largest pairs, or nil if tree is empty
166
212
  # @example
167
213
  # tree = RBTree.new({3 => 'three', 1 => 'one', 2 => 'two'})
168
- # tree.max # => [3, "three"]
169
- def max = max_node&.pair
214
+ # tree.max # => [3, "three"]
215
+ # tree.max(2) # => [[3, "three"], [2, "two"]]
216
+ def max(*args, &block)
217
+ return super if !args.empty? || block
218
+ max_node&.pair
219
+ end
170
220
 
171
- # Returns the first key-value pair without removing it.
221
+ # Returns the first key-value pair, or the first +n+ pairs, without removing them.
172
222
  #
173
- # @return [Array(Object, Object), nil] a two-element array [key, value], or nil if tree is empty
223
+ # @param n [Integer, nil] number of leading pairs to return
224
+ # @return [Array, nil] the pair, or the n smallest pairs in ascending order
174
225
  # @example
175
226
  # tree = RBTree.new({3 => 'three', 1 => 'one', 2 => 'two'})
176
- # tree.first # => [1, "one"]
177
- def first = min
227
+ # tree.first # => [1, "one"]
228
+ # tree.first(2) # => [[1, "one"], [2, "two"]]
229
+ def first(n = nil) = n.nil? ? min : take(n)
178
230
 
179
- # Returns the last key-value pair without removing it.
231
+ # Returns the last key-value pair, or the last +n+ pairs, without removing them.
180
232
  #
181
- # @return [Array(Object, Object), nil] a two-element array [key, value], or nil if tree is empty
233
+ # @param n [Integer, nil] number of trailing pairs to return
234
+ # @return [Array, nil] the pair, or the n largest pairs in ascending order
182
235
  # @example
183
236
  # tree = RBTree.new({3 => 'three', 1 => 'one', 2 => 'two'})
184
- # tree.last # => [3, "three"]
185
- def last = max
237
+ # tree.last # => [3, "three"]
238
+ # tree.last(2) # => [[2, "two"], [3, "three"]]
239
+ def last(n = nil)
240
+ return max if n.nil?
241
+ result = []
242
+ reverse_each { |pair| break if result.size >= n; result << pair }
243
+ result.reverse!
244
+ result
245
+ end
186
246
 
187
247
  # Checks if the tree contains the given key.
188
248
  #
249
+ # A key counts as present when <tt>(key <=> stored_key) == 0</tt>, even if the two
250
+ # objects are not +eql?+ (e.g. +1.0+ finds a stored +1+).
251
+ #
189
252
  # @param key [Object] the key to search for
190
253
  # @return [Boolean] true if the key exists in the tree, false otherwise
191
254
  # @example
192
255
  # tree = RBTree.new({1 => 'one', 2 => 'two'})
193
- # tree.key?(1) # => true
194
- # tree.key?(3) # => false
195
- def has_key?(key) = @hash_index.key?(key)
256
+ # tree.key?(1) # => true
257
+ # tree.key?(1.0) # => true
258
+ # tree.key?(3) # => false
259
+ def has_key?(key) = @hash_index.key?(key) || !find_node_by_order(key).nil?
196
260
  alias :key? :has_key?
197
261
 
198
262
  # Retrieves the value associated with the given key.
@@ -202,7 +266,7 @@ class RBTree
202
266
  # @example
203
267
  # tree = RBTree.new({1 => 'one', 2 => 'two'})
204
268
  # tree.get(1) # => "one"
205
- def value(key) = @hash_index[key]&.value
269
+ def value(key) = (@hash_index[key] || find_node_by_order(key))&.value
206
270
  alias :get :value
207
271
 
208
272
  # Retrieves a value associated with the given key, or a range of entries if a Range is provided.
@@ -363,6 +427,9 @@ class RBTree
363
427
  raise ArgumentError, "Source must be iterable"
364
428
  end
365
429
 
430
+ # Self-insertion (e.g. tree.merge!(tree)) must not mutate what it iterates.
431
+ source = source.to_a if source.equal?(self)
432
+
366
433
  source.each do |*pair|
367
434
  key, value = nil, nil
368
435
  if pair.size == 1 && pair[0].is_a?(Array)
@@ -406,6 +473,8 @@ class RBTree
406
473
  end
407
474
  if block
408
475
  other_enum = other.is_a?(Hash) || other.is_a?(RBTree) ? other : other.each
476
+ # Snapshot on self-merge; see RBTree#insert.
477
+ other_enum = other_enum.to_a if other_enum.equal?(self)
409
478
  other_enum.each do |k, v|
410
479
  if has_key?(k)
411
480
  insert_entry(k, block.call(k, value(k), v), overwrite: true)
@@ -421,6 +490,10 @@ class RBTree
421
490
 
422
491
  # Deletes the key-value pair with the specified key.
423
492
  #
493
+ # Entries whose value is +nil+ or +false+ are deleted correctly, so the return
494
+ # value alone cannot distinguish them from a missing key; use {#has_key?} first
495
+ # if that matters.
496
+ #
424
497
  # @param key [Object] the key to delete
425
498
  # @return [Object, nil] the value associated with the deleted key, or nil if not found
426
499
  # @example
@@ -428,8 +501,9 @@ class RBTree
428
501
  # tree.delete(1) # => "one"
429
502
  # tree.delete(3) # => nil
430
503
  def delete_key(key)
431
- return nil unless (value = (z = @hash_index[key])&.value)
432
- delete_indexed_node(key)
504
+ return nil unless (z = find_node(key))
505
+ value = z.value
506
+ delete_found_node(z)
433
507
  value
434
508
  end
435
509
  alias :delete :delete_key
@@ -464,11 +538,18 @@ class RBTree
464
538
 
465
539
  # Removes all key-value pairs from the tree.
466
540
  #
541
+ # Runs in O(1): the nodes are left to the garbage collector, and the allocator is
542
+ # told of the bulk discard so that pool statistics stay accurate.
543
+ #
467
544
  # @return [RBTree] self
468
545
  def clear
546
+ @node_allocator.discard(@key_count) if @key_count > 0
469
547
  @root = @min_node = @max_node = @nil_node
470
548
  @hash_index.clear
471
549
  @key_count = 0
550
+ @key_class = nil
551
+ @coherent_keys = false
552
+ @mod_count += 1
472
553
  self
473
554
  end
474
555
 
@@ -494,7 +575,8 @@ class RBTree
494
575
  # tree.delete(k) if k.even?
495
576
  # end
496
577
  def keys(reverse: false, safe: false, &block)
497
- return enum_for(__method__, reverse: reverse, safe: safe) { @key_count } unless block_given?
578
+ # `size`, not `key_count`: MultiRBTree yields a key once per value it holds.
579
+ return enum_for(__method__, reverse: reverse, safe: safe) { size } unless block_given?
498
580
  each(reverse: reverse, safe: safe) { |key, _| yield key }
499
581
  self
500
582
  end
@@ -645,7 +727,7 @@ class RBTree
645
727
  # @return [RBTree, Enumerator] a new tree with selected pairs, or Enumerator if no block
646
728
  def select(&block)
647
729
  return enum_for(__method__) { size } unless block_given?
648
- result = self.class.new
730
+ result = new_derived_tree
649
731
  each { |k, v| result.insert(k, v) if block.call(k, v) }
650
732
  result
651
733
  end
@@ -656,7 +738,7 @@ class RBTree
656
738
  # @return [RBTree, Enumerator] a new tree with non-rejected pairs, or Enumerator if no block
657
739
  def reject(&block)
658
740
  return enum_for(__method__) { size } unless block_given?
659
- result = self.class.new
741
+ result = new_derived_tree
660
742
  each { |k, v| result.insert(k, v) unless block.call(k, v) }
661
743
  result
662
744
  end
@@ -700,7 +782,7 @@ class RBTree
700
782
  #
701
783
  # @return [RBTree, MultiRBTree] a new tree with keys and values inverted
702
784
  def invert
703
- result = self.class.new
785
+ result = new_derived_tree
704
786
  each { |k, v| result.insert(v, k) }
705
787
  result
706
788
  end
@@ -716,19 +798,41 @@ class RBTree
716
798
  "#<#{self.class}:0x#{object_id.to_s(16)} size=#{size} {#{content}#{suffix}}>"
717
799
  end
718
800
 
719
- # Validates the red-black tree properties.
801
+ # Validates the red-black tree properties and the auxiliary structures.
720
802
  #
721
803
  # Checks that:
722
- # 1. Root is black
804
+ # 1. Root is black and has no parent
723
805
  # 2. All paths from root to leaves have the same number of black nodes
724
806
  # 3. No red node has a red child
725
- # 4. Keys are properly ordered
807
+ # 4. Keys are ordered against bounds inherited from ancestors
808
+ # 5. Every child's parent pointer points back at its parent
809
+ # 6. The Hash index holds exactly one entry per node, mapped to that node
810
+ # 7. `@min_node` / `@max_node` are the actual extremes, and `key_count` is right
726
811
  #
727
812
  # @return [Boolean] true if all properties are satisfied, false otherwise
728
813
  def valid?
729
814
  return false if @root.color == Node::RED
815
+ return false if @root != @nil_node && @root.parent != @nil_node
730
816
  return false if check_black_height(@root) == -1
731
- return false unless check_order(@root)
817
+ return false unless check_order(@root, nil, nil)
818
+ return false unless check_links(@root)
819
+
820
+ count = 0
821
+ first_node = nil
822
+ last_node = nil
823
+ indexed = true
824
+ each_node_asc do |n|
825
+ count += 1
826
+ first_node ||= n
827
+ last_node = n
828
+ indexed &&= @hash_index[n.key].equal?(n)
829
+ end
830
+
831
+ return false unless indexed
832
+ return false unless count == @key_count
833
+ return false unless @hash_index.size == @key_count
834
+ return false unless (first_node || @nil_node).equal?(@min_node)
835
+ return false unless (last_node || @nil_node).equal?(@max_node)
732
836
  true
733
837
  end
734
838
 
@@ -738,6 +842,108 @@ class RBTree
738
842
  def min_node = (@min_node == @nil_node) ? nil : @min_node
739
843
  def max_node = (@max_node == @nil_node) ? nil : @max_node
740
844
 
845
+ # Creates an empty tree of the same class, keeping this tree's allocator.
846
+ #
847
+ # @return [RBTree] a new empty tree
848
+ def new_derived_tree = self.class.new(node_allocator: @node_allocator)
849
+
850
+ # Locates the node holding the given key, by O(1) hash index or by ordering.
851
+ #
852
+ # The index is keyed by +eql?+ while the tree orders by <tt><=></tt>, so a key
853
+ # such as +1.0+ misses the index but still matches a stored +1+.
854
+ #
855
+ # @param key [Object] the key to locate
856
+ # @return [Node, nil] the node holding an equal key, or nil if absent
857
+ def find_node(key) = @hash_index[key] || find_node_by_order(key)
858
+
859
+ # Resolves a key the hash index missed, using the tree's <tt><=></tt> ordering.
860
+ #
861
+ # Split from {#find_node} so hot callers can inline the hash probe and pay no
862
+ # method call on a hit. The descent is skipped when every stored key and the
863
+ # lookup key share a class in {COHERENT_KEY_CLASSES}, since a match would then
864
+ # have been found in the index already. Incomparable keys report absence rather
865
+ # than raising, so lookups stay total.
866
+ #
867
+ # @param key [Object] the key to locate
868
+ # @return [Node, nil] the node holding an equal key, or nil if absent
869
+ def find_node_by_order(key)
870
+ return nil if @coherent_keys && key.instance_of?(@key_class)
871
+
872
+ x = @root
873
+ while x != @nil_node
874
+ cmp = key <=> x.key
875
+ return nil unless cmp
876
+ return x if cmp == 0
877
+ x = cmp < 0 ? x.left : x.right
878
+ end
879
+ nil
880
+ end
881
+
882
+ # Records the class of a newly inserted key for {#find_node_by_order}.
883
+ #
884
+ # Deletion never relaxes the result: a stale pessimistic value costs a descent,
885
+ # never correctness.
886
+ #
887
+ # @param key [Object] the key of the node being created
888
+ # @return [void]
889
+ def note_key_class(key)
890
+ klass = key.class
891
+ return if @key_class.equal?(klass)
892
+
893
+ if @key_class.nil?
894
+ @key_class = klass
895
+ @coherent_keys = COHERENT_KEY_CLASSES.include?(klass)
896
+ else
897
+ @key_class = false
898
+ @coherent_keys = false
899
+ end
900
+ end
901
+
902
+ # Compares two keys, rejecting incomparable ones with a clear error.
903
+ #
904
+ # @param a [Object] the left-hand key
905
+ # @param b [Object] the right-hand key
906
+ # @return [Integer] the result of <tt>a <=> b</tt>
907
+ # @raise [ArgumentError] if the keys cannot be compared
908
+ def compare!(a, b)
909
+ cmp = a <=> b
910
+ return cmp if cmp
911
+ raise ArgumentError, "comparison of #{a.class} with #{b.inspect} failed"
912
+ end
913
+
914
+ # Reports mutation detected during a non-safe traversal.
915
+ #
916
+ # Such a traversal walks live node links, and released nodes are recycled by the
917
+ # allocator, so mutating during one can yield keys out of order rather than
918
+ # merely skip them. Callers compare `@mod_count` inline and only call this on
919
+ # mismatch, keeping the per-element cost to an integer comparison.
920
+ #
921
+ # @raise [RuntimeError] always
922
+ # @return [void]
923
+ def concurrent_modification!
924
+ raise RuntimeError,
925
+ "can't modify #{self.class} during iteration (pass `safe: true` to iterate safely)"
926
+ end
927
+
928
+ # Walks every node in ascending key order, following live tree links.
929
+ #
930
+ # @yield [node] each node
931
+ # @return [void]
932
+ def each_node_asc
933
+ stack = []
934
+ current = @root
935
+ while current != @nil_node || !stack.empty?
936
+ if current != @nil_node
937
+ stack << current
938
+ current = current.left
939
+ else
940
+ current = stack.pop
941
+ yield current
942
+ current = current.right
943
+ end
944
+ end
945
+ end
946
+
741
947
  # Inserts a single key-value pair.
742
948
  #
743
949
  # @param key [Object] the key to insert
@@ -761,6 +967,9 @@ class RBTree
761
967
 
762
968
  # Generic entry insertion logic shared between RBTree and MultiRBTree.
763
969
  #
970
+ # A key that is <tt><=></tt>-equal to a stored key updates that node, which keeps
971
+ # the key object it was created with, so the index stays at one entry per node.
972
+ #
764
973
  # @param key [Object] the key to insert
765
974
  # @yield [node, is_new] yields the existing node (if any) and whether it's a new insertion
766
975
  # @yieldparam node [Node, nil] the existing node or nil
@@ -768,6 +977,7 @@ class RBTree
768
977
  # @yieldreturn [Object]
769
978
  # - if is_new: the initial value for the new node
770
979
  # - if !is_new: the value to return from insert_entry
980
+ # @raise [ArgumentError] if the key cannot be compared with the stored keys
771
981
  def insert_entry_generic(key)
772
982
  if (node = @hash_index[key])
773
983
  return yield(node, false)
@@ -778,6 +988,9 @@ class RBTree
778
988
  while x != @nil_node
779
989
  y = x
780
990
  cmp = key <=> x.key
991
+ unless cmp
992
+ raise ArgumentError, "comparison of #{key.class} with #{x.key.inspect} failed"
993
+ end
781
994
  if cmp == 0
782
995
  return yield(x, false)
783
996
  elsif cmp < 0
@@ -811,6 +1024,8 @@ class RBTree
811
1024
  end
812
1025
 
813
1026
  @hash_index[key] = z
1027
+ note_key_class(key)
1028
+ @mod_count += 1
814
1029
  true
815
1030
  end
816
1031
 
@@ -851,13 +1066,13 @@ class RBTree
851
1066
 
852
1067
  if safe
853
1068
  pair = !min ? find_min :
854
- include_min && @hash_index[min]&.pair || find_successor(min)
855
- while pair && (!max || pair[0] < max)
1069
+ include_min && find_node(min)&.pair || find_successor(min)
1070
+ while pair && (!max || (pair[0] <=> max) < 0)
856
1071
  current_key = pair[0]
857
1072
  yield pair
858
1073
  pair = find_successor(current_key)
859
1074
  end
860
- yield pair if pair && max && include_max && pair[0] == max
1075
+ yield pair if pair && max && include_max && (pair[0] <=> max) == 0
861
1076
  else
862
1077
  start_node, stack = resolve_startup_asc(min, include_min)
863
1078
  traverse_from_asc(start_node, stack, max, include_max, &block)
@@ -883,13 +1098,13 @@ class RBTree
883
1098
 
884
1099
  if safe
885
1100
  pair = !max ? find_max :
886
- include_max && @hash_index[max]&.pair || find_predecessor(max)
887
- while pair && (!min || pair[0] > min)
1101
+ include_max && find_node(max)&.pair || find_predecessor(max)
1102
+ while pair && (!min || (pair[0] <=> min) > 0)
888
1103
  current_key = pair[0]
889
1104
  yield pair
890
1105
  pair = find_predecessor(current_key)
891
1106
  end
892
- yield pair if pair && min && include_min && pair[0] == min
1107
+ yield pair if pair && min && include_min && (pair[0] <=> min) == 0
893
1108
  else
894
1109
  start_node, stack = resolve_startup_desc(max, include_max)
895
1110
  traverse_from_desc(start_node, stack, min, include_min, &block)
@@ -941,7 +1156,8 @@ class RBTree
941
1156
  return [@min_node, reconstruct_stack_asc(@min_node)]
942
1157
  end
943
1158
 
944
- # 2. Use Hash index if key exists
1159
+ # 2. Use Hash index if key exists. A pure O(1) probe, not {#find_node}: on a
1160
+ # miss, step 3's descent already places a <=>-equal node correctly.
945
1161
  if min && (node = @hash_index[min])
946
1162
  start_node = include_min ? node : successor_node_of(node)
947
1163
  return [start_node, reconstruct_stack_asc(start_node)]
@@ -987,16 +1203,21 @@ class RBTree
987
1203
  # @yieldparam key [Object] the key
988
1204
  # @yieldparam val [Object] the value
989
1205
  def traverse_from_asc(current, stack, max, include_max, &block)
1206
+ mod = @mod_count
990
1207
  while current != @nil_node || !stack.empty?
991
1208
  if current != @nil_node
992
1209
  if max
993
1210
  cmp = current.key <=> max
994
1211
  if cmp >= 0
995
- yield current.pair if include_max && cmp == 0
1212
+ if include_max && cmp == 0
1213
+ yield current.pair
1214
+ concurrent_modification! unless @mod_count == mod
1215
+ end
996
1216
  return
997
1217
  end
998
1218
  end
999
1219
  yield current.pair
1220
+ concurrent_modification! unless @mod_count == mod
1000
1221
  current = current.right
1001
1222
  while current != @nil_node
1002
1223
  stack << current
@@ -1019,7 +1240,7 @@ class RBTree
1019
1240
  return [@max_node, reconstruct_stack_desc(@max_node)]
1020
1241
  end
1021
1242
 
1022
- # 2. Use Hash index if key exists
1243
+ # 2. Use Hash index if key exists; see the note in #resolve_startup_asc.
1023
1244
  if max && (node = @hash_index[max])
1024
1245
  start_node = include_max ? node : predecessor_node_of(node)
1025
1246
  return [start_node, reconstruct_stack_desc(start_node)]
@@ -1065,16 +1286,21 @@ class RBTree
1065
1286
  # @yieldparam key [Object] the key
1066
1287
  # @yieldparam val [Object] the value
1067
1288
  def traverse_from_desc(current, stack, min, include_min, &block)
1289
+ mod = @mod_count
1068
1290
  while current != @nil_node || !stack.empty?
1069
1291
  if current != @nil_node
1070
1292
  if min
1071
1293
  cmp = current.key <=> min
1072
1294
  if cmp <= 0
1073
- yield current.pair if include_min && cmp == 0
1295
+ if include_min && cmp == 0
1296
+ yield current.pair
1297
+ concurrent_modification! unless @mod_count == mod
1298
+ end
1074
1299
  return
1075
1300
  end
1076
1301
  end
1077
1302
  yield current.pair
1303
+ concurrent_modification! unless @mod_count == mod
1078
1304
  current = current.left
1079
1305
  while current != @nil_node
1080
1306
  stack << current
@@ -1136,7 +1362,17 @@ class RBTree
1136
1362
  #
1137
1363
  # @param key [Object] the key to delete
1138
1364
  # @return [Object, nil] the value of the deleted node, or nil if not found
1139
- def delete_indexed_node(key) = (z = @hash_index.delete(key)) && delete_node(z)
1365
+ def delete_indexed_node(key) = (z = find_node(key)) && delete_found_node(z)
1366
+
1367
+ # Removes an already-located node, unindexing it under its own key so that a node
1368
+ # found through a non-+eql?+ but <tt><=></tt>-equal key is still removed correctly.
1369
+ #
1370
+ # @param z [Node] the node to remove
1371
+ # @return [Object] the value of the removed node
1372
+ def delete_found_node(z)
1373
+ @hash_index.delete(z.key)
1374
+ delete_node(z)
1375
+ end
1140
1376
 
1141
1377
  # Removes a node from the tree and restores red-black properties.
1142
1378
  #
@@ -1195,6 +1431,7 @@ class RBTree
1195
1431
 
1196
1432
  value = z.value
1197
1433
  release_node(z)
1434
+ @mod_count += 1
1198
1435
  value
1199
1436
  end
1200
1437
 
@@ -1332,8 +1569,9 @@ class RBTree
1332
1569
  # @return [Node] the predecessor node, or @nil_node if none exists
1333
1570
  def find_predecessor_node(key)
1334
1571
  # If key is larger than max_key, return max_node
1335
- return @max_node if max_key && (key <=> max_key) > 0
1336
- # If key exists using O(1) hash lookup, return predecessor node
1572
+ return @max_node if (mk = max_key) && compare!(key, mk) > 0
1573
+ # O(1) hash fast path; on a miss the descent below steps away from a
1574
+ # <=>-equal node in the right direction anyway.
1337
1575
  if (node = @hash_index[key])
1338
1576
  return predecessor_node_of(node)
1339
1577
  end
@@ -1372,10 +1610,10 @@ class RBTree
1372
1610
  # @return [Node] the successor node, or @nil_node if none exists
1373
1611
  def find_successor_node(key)
1374
1612
  # If key is larger than or equal to max_key, return nil
1375
- return @nil_node if max_key && (key <=> max_key) >= 0
1613
+ return @nil_node if (mk = max_key) && compare!(key, mk) >= 0
1376
1614
  # If key is smaller than min_key, return min_node
1377
- return @min_node if min_key && (key <=> min_key) < 0
1378
- # If key exists using O(1) hash lookup, return successor node
1615
+ return @min_node if (nk = min_key) && compare!(key, nk) < 0
1616
+ # O(1) hash fast path; see the note in #find_predecessor_node.
1379
1617
  if (node = @hash_index[key])
1380
1618
  return successor_node_of(node)
1381
1619
  end
@@ -1534,15 +1772,31 @@ class RBTree
1534
1772
  left_h + (node.color == Node::BLACK ? 1 : 0)
1535
1773
  end
1536
1774
 
1537
- def check_order(node)
1775
+ # Recursively checks the binary-search-tree ordering.
1776
+ #
1777
+ # Bounds are inherited from ancestors, so a key misplaced relative to a higher
1778
+ # ancestor is rejected even when it is ordered correctly against its parent.
1779
+ #
1780
+ # @param node [Node] the current node
1781
+ # @param lo [Object, nil] exclusive lower bound inherited from ancestors
1782
+ # @param hi [Object, nil] exclusive upper bound inherited from ancestors
1783
+ # @return [Boolean] true if the subtree is correctly ordered
1784
+ def check_order(node, lo, hi)
1538
1785
  return true if node == @nil_node
1539
- if node.left != @nil_node && (node.left.key <=> node.key) >= 0
1540
- return false
1541
- end
1542
- if node.right != @nil_node && (node.right.key <=> node.key) <= 0
1543
- return false
1544
- end
1545
- check_order(node.left) && check_order(node.right)
1786
+ return false if lo && (node.key <=> lo) <= 0
1787
+ return false if hi && (node.key <=> hi) >= 0
1788
+ check_order(node.left, lo, node.key) && check_order(node.right, node.key, hi)
1789
+ end
1790
+
1791
+ # Recursively checks that every child points back at its parent.
1792
+ #
1793
+ # @param node [Node] the current node
1794
+ # @return [Boolean] true if all parent links in the subtree are consistent
1795
+ def check_links(node)
1796
+ return true if node == @nil_node
1797
+ return false if node.left != @nil_node && !node.left.parent.equal?(node)
1798
+ return false if node.right != @nil_node && !node.right.parent.equal?(node)
1799
+ check_links(node.left) && check_links(node.right)
1546
1800
  end
1547
1801
  end
1548
1802
 
@@ -1611,6 +1865,11 @@ end
1611
1865
  # @author Masahito Suzuki
1612
1866
  # @since 0.1.2
1613
1867
  class MultiRBTree < RBTree
1868
+ # Sentinel for "no key given", so +nil+ and +false+ stay usable as keys.
1869
+ # @api private
1870
+ NO_KEY = Object.new
1871
+ private_constant :NO_KEY
1872
+
1614
1873
  def initialize(*args, **kwargs)
1615
1874
  @value_count = 0
1616
1875
  super
@@ -1626,38 +1885,83 @@ class MultiRBTree < RBTree
1626
1885
  @value_count = 0
1627
1886
  super
1628
1887
  end
1888
+
1889
+ # Returns a Hash mapping each key to an Array of its values.
1890
+ #
1891
+ # Keys are inserted in ascending order, and each value Array is a fresh copy,
1892
+ # so mutating the result cannot corrupt the tree.
1893
+ #
1894
+ # @return [Hash] a new Hash with the tree's contents
1895
+ # @example
1896
+ # tree = MultiRBTree.new
1897
+ # tree.insert(1, 'a'); tree.insert(1, 'b')
1898
+ # tree.to_h # => {1 => ["a", "b"]}
1899
+ def to_h
1900
+ h = {}
1901
+ each { |k, v| (h[k] ||= []) << v }
1902
+ h
1903
+ end
1904
+
1905
+ # Validates everything `RBTree#valid?` does, plus that every node holds a
1906
+ # non-empty Array of values and that `size` matches the total value count.
1907
+ #
1908
+ # @return [Boolean] true if all properties are satisfied, false otherwise
1909
+ def valid?
1910
+ return false unless super
1911
+ total = 0
1912
+ ok = true
1913
+ each_node_asc do |n|
1914
+ ok &&= n.value.is_a?(Array) && !n.value.empty?
1915
+ total += n.value.size if n.value.is_a?(Array)
1916
+ end
1917
+ ok && total == @value_count
1918
+ end
1629
1919
 
1630
1920
  # Returns the minimum key-value pair without removing it.
1631
1921
  #
1922
+ # As with `RBTree#min`, a count or comparison block switches to `Enumerable#min`.
1923
+ #
1632
1924
  # @param last [Boolean] whether to return the last value (default: false)
1633
1925
  # @return [Array(Object, Object), nil] a two-element array [key, value], or nil if tree is empty
1634
1926
  # @example
1635
1927
  # tree = MultiRBTree.new({3 => 'three', 1 => 'one', 2 => 'two'})
1636
1928
  # tree.min # => [1, "one"]
1637
- def min(last: false) = (n = min_node) && [n.key, n.value.send(last ? :last : :first)]
1929
+ def min(*args, last: false, &block)
1930
+ return super(*args, &block) if !args.empty? || block
1931
+ (n = min_node) && [n.key, n.value.send(last ? :last : :first)]
1932
+ end
1638
1933
 
1639
1934
  # Returns the maximum key-value pair without removing it.
1640
1935
  #
1936
+ # As with `RBTree#max`, a count or comparison block switches to `Enumerable#max`.
1937
+ #
1641
1938
  # @param last [Boolean] whether to return the last value (default: false)
1642
1939
  # @return [Array(Object, Object), nil] a two-element array [key, value], or nil if tree is empty
1643
1940
  # @example
1644
1941
  # tree = MultiRBTree.new({3 => 'three', 1 => 'one', 2 => 'two'})
1645
1942
  # tree.max # => [3, "three"]
1646
- def max(last: false) = (n = max_node) && [n.key, n.value.send(last ? :last : :first)]
1943
+ def max(*args, last: false, &block)
1944
+ return super(*args, &block) if !args.empty? || block
1945
+ (n = max_node) && [n.key, n.value.send(last ? :last : :first)]
1946
+ end
1647
1947
 
1648
- # Returns the last key-value pair without removing it.
1948
+ # Returns the last key-value pair, or the last +n+ pairs, without removing them.
1649
1949
  #
1650
- # @return [Array(Object, Object), nil] a two-element array [key, value], or nil if tree is empty
1950
+ # @param n [Integer, nil] number of trailing pairs to return
1951
+ # @return [Array, nil] the pair, or the n last pairs in ascending order
1651
1952
  # @example
1652
1953
  # tree = MultiRBTree.new({3 => 'three', 1 => 'one', 2 => 'two'})
1653
1954
  # tree.last # => [3, "three"]
1654
- def last = max(last: true)
1955
+ def last(n = nil) = n.nil? ? max(last: true) : super(n)
1655
1956
 
1656
- # Returns the number of values for a given key or the total number of key-value pairs if no key is given.
1957
+ # Returns the number of values for a given key, or the total number of key-value
1958
+ # pairs when called without an argument.
1959
+ #
1960
+ # The default is a sentinel, so +nil+ and +false+ are looked up as ordinary keys.
1657
1961
  #
1658
- # @param key [Object, nil] the key to look up, or nil for total count
1962
+ # @param key [Object] the key to look up; omit for the total count
1659
1963
  # @return [Integer] the number of values for the key, or total count if no key is given
1660
- def value_count(key = nil) = !key ? size : (@hash_index[key]&.value&.size || 0)
1964
+ def value_count(key = NO_KEY) = key.equal?(NO_KEY) ? size : (find_node(key)&.value&.size || 0)
1661
1965
 
1662
1966
  # Retrieves a value associated with the given key.
1663
1967
  #
@@ -1670,7 +1974,7 @@ class MultiRBTree < RBTree
1670
1974
  # tree.insert(1, 'second')
1671
1975
  # tree.get(1) # => "first"
1672
1976
  # tree.get(1, last: true) # => "second"
1673
- def value(key, last: false) = @hash_index[key]&.value&.send(last ? :last : :first)
1977
+ def value(key, last: false) = (@hash_index[key] || find_node_by_order(key))&.value&.send(last ? :last : :first)
1674
1978
  alias :get :value
1675
1979
 
1676
1980
  # Retrieves the first value associated with the given key.
@@ -1690,15 +1994,17 @@ class MultiRBTree < RBTree
1690
1994
  # Retrieves all values associated with the given key.
1691
1995
  #
1692
1996
  # @param key [Object] the key to look up
1693
- # @return [Array, nil] an Array containing all values, or nil if not found
1997
+ # @param reverse [Boolean] if true, yield the values in reverse insertion order (default: false)
1998
+ # @return [Enumerator, nil] an Enumerator over the values, or nil if the key is absent (when a block is given)
1694
1999
  # @example
1695
2000
  # tree = MultiRBTree.new
1696
2001
  # tree.insert(1, 'first')
1697
2002
  # tree.insert(1, 'second')
1698
- # tree.values(1).to_a # => ["first", "second"]
2003
+ # tree.values(1).to_a # => ["first", "second"]
2004
+ # tree.values(1, reverse: true).to_a # => ["second", "first"]
1699
2005
  def values(key, reverse: false)
1700
- return enum_for(__method__, key) { value_count(key) } unless block_given?
1701
- @hash_index[key]&.value&.send(reverse ? :reverse_each : :each) { |v| yield v }
2006
+ return enum_for(__method__, key, reverse: reverse) { value_count(key) } unless block_given?
2007
+ find_node(key)&.value&.send(reverse ? :reverse_each : :each) { |v| yield v }
1702
2008
  end
1703
2009
  alias :get_all :values
1704
2010
 
@@ -1758,9 +2064,13 @@ class MultiRBTree < RBTree
1758
2064
  # tree.delete_value(1) # => "first"
1759
2065
  # tree.delete_value(1, last: true) # => "second" (if more values existed)
1760
2066
  def delete_value(key, last: false)
1761
- (z = @hash_index[key]) or return nil
2067
+ (z = find_node(key)) or return nil
1762
2068
  value = z.value.send(last ? :pop : :shift)
1763
- z.value.empty? && delete_indexed_node(key)
2069
+ if z.value.empty?
2070
+ delete_found_node(z)
2071
+ else
2072
+ @mod_count += 1
2073
+ end
1764
2074
  @value_count -= 1
1765
2075
  value
1766
2076
  end
@@ -1793,9 +2103,9 @@ class MultiRBTree < RBTree
1793
2103
  # vals = tree.delete(1) # removes both values
1794
2104
  # vals.size # => 2
1795
2105
  def delete_key(key)
1796
- return nil unless (z = @hash_index[key])
2106
+ return nil unless (z = find_node(key))
1797
2107
  @value_count -= (value = z.value).size
1798
- delete_indexed_node(z.key)
2108
+ delete_found_node(z)
1799
2109
  value
1800
2110
  end
1801
2111
  alias :delete :delete_key
@@ -1807,9 +2117,14 @@ class MultiRBTree < RBTree
1807
2117
  # tree = MultiRBTree.new({3 => 'three', 1 => 'one', 2 => 'two'})
1808
2118
  # tree.shift # => [1, "one"]
1809
2119
  def shift
1810
- (key, vals = min_node&.pair) or return nil
2120
+ (n = min_node) or return nil
2121
+ key, vals = n.pair
1811
2122
  val = vals.shift
1812
- vals.empty? && delete_indexed_node(key)
2123
+ if vals.empty?
2124
+ delete_found_node(n)
2125
+ else
2126
+ @mod_count += 1
2127
+ end
1813
2128
  @value_count -= 1
1814
2129
  [key, val]
1815
2130
  end
@@ -1821,9 +2136,14 @@ class MultiRBTree < RBTree
1821
2136
  # tree = MultiRBTree.new({3 => 'three', 1 => 'one', 2 => 'two'})
1822
2137
  # tree.pop # => [3, "three"]
1823
2138
  def pop
1824
- (key, vals = max_node&.pair) or return nil
2139
+ (n = max_node) or return nil
2140
+ key, vals = n.pair
1825
2141
  val = vals.pop
1826
- vals.empty? && delete_indexed_node(key)
2142
+ if vals.empty?
2143
+ delete_found_node(n)
2144
+ else
2145
+ @mod_count += 1
2146
+ end
1827
2147
  @value_count -= 1
1828
2148
  [key, val]
1829
2149
  end
@@ -1858,14 +2178,17 @@ class MultiRBTree < RBTree
1858
2178
  # Removes nodes whose value arrays become empty.
1859
2179
  # Updates @value_count accordingly.
1860
2180
  def filter_values!
1861
- keys_to_delete = []
1862
- @hash_index.each do |key, node|
2181
+ emptied = []
2182
+ @hash_index.each_value do |node|
1863
2183
  before = node.value.size
1864
- node.value.select! { |v| yield key, v }
1865
- @value_count -= before - node.value.size
1866
- keys_to_delete << key if node.value.empty?
2184
+ node.value.select! { |v| yield node.key, v }
2185
+ removed = before - node.value.size
2186
+ next if removed == 0
2187
+ @value_count -= removed
2188
+ @mod_count += 1
2189
+ emptied << node if node.value.empty?
1867
2190
  end
1868
- keys_to_delete.each { |k| delete_indexed_node(k) }
2191
+ emptied.each { |node| delete_found_node(node) }
1869
2192
  end
1870
2193
 
1871
2194
  # Inserts a value for the given key.
@@ -1888,6 +2211,7 @@ class MultiRBTree < RBTree
1888
2211
  [value]
1889
2212
  else
1890
2213
  node.value << value
2214
+ @mod_count += 1 # observable to a traversal walking this value list
1891
2215
  true
1892
2216
  end
1893
2217
  end
@@ -1895,22 +2219,49 @@ class MultiRBTree < RBTree
1895
2219
 
1896
2220
  # Traverses the tree in ascending order, yielding each key-value pair.
1897
2221
  #
1898
- # @param range [Range] the range of keys to traverse
2222
+ # Each node's value list is flattened to one [key, value] pair per value. Safe
2223
+ # mode snapshots the list, so removing values of the key being iterated cannot
2224
+ # skip its siblings, though such values may still be yielded.
2225
+ #
2226
+ # @param min [Object] the lower bound
2227
+ # @param max [Object] the upper bound
2228
+ # @param include_min [Boolean] whether to include the lower bound
2229
+ # @param include_max [Boolean] whether to include the upper bound
2230
+ # @param safe [Boolean] whether to use safe traversal
1899
2231
  # @yield [Array(Object, Object)] each key-value pair
1900
- # @yieldparam key [Object] the key
1901
- # @yieldparam val [Object] the value
1902
- def traverse_range_asc(...)
1903
- super { |k, vals| vals.each { |v| yield [k, v] } }
2232
+ def traverse_range_asc(min, max, include_min, include_max, safe: false, &block)
2233
+ if safe
2234
+ super(min, max, include_min, include_max, safe: true) do |k, vals|
2235
+ vals.dup.each { |v| yield [k, v] }
2236
+ end
2237
+ else
2238
+ mod = @mod_count
2239
+ super(min, max, include_min, include_max, safe: false) do |k, vals|
2240
+ vals.each { |v| yield [k, v]; concurrent_modification! unless @mod_count == mod }
2241
+ end
2242
+ end
1904
2243
  end
1905
2244
 
1906
2245
  # Traverses the tree in descending order, yielding each key-value pair.
1907
2246
  #
1908
- # @param range [Range] the range of keys to traverse
2247
+ # @param min [Object] the lower bound
2248
+ # @param max [Object] the upper bound
2249
+ # @param include_min [Boolean] whether to include the lower bound
2250
+ # @param include_max [Boolean] whether to include the upper bound
2251
+ # @param safe [Boolean] whether to use safe traversal
1909
2252
  # @yield [Array(Object, Object)] each key-value pair
1910
- # @yieldparam key [Object] the key
1911
- # @yieldparam val [Object] the value
1912
- def traverse_range_desc(...)
1913
- super { |k, vals| vals.reverse_each { |v| yield [k, v] } }
2253
+ # @see #traverse_range_asc
2254
+ def traverse_range_desc(min, max, include_min, include_max, safe: false, &block)
2255
+ if safe
2256
+ super(min, max, include_min, include_max, safe: true) do |k, vals|
2257
+ vals.dup.reverse_each { |v| yield [k, v] }
2258
+ end
2259
+ else
2260
+ mod = @mod_count
2261
+ super(min, max, include_min, include_max, safe: false) do |k, vals|
2262
+ vals.reverse_each { |v| yield [k, v]; concurrent_modification! unless @mod_count == mod }
2263
+ end
2264
+ end
1914
2265
  end
1915
2266
  end
1916
2267
 
@@ -1981,6 +2332,15 @@ class RBTree::NodeAllocator
1981
2332
  #
1982
2333
  # @param node [Node] the node to release
1983
2334
  def release(node) = nil
2335
+
2336
+ # Notifies the allocator that +count+ nodes were dropped in bulk rather than
2337
+ # released individually, as `RBTree#clear` does.
2338
+ #
2339
+ # The nodes are not reclaimed for reuse; this only keeps live-node statistics right.
2340
+ #
2341
+ # @param count [Integer] the number of nodes discarded
2342
+ # @return [void]
2343
+ def discard(count) = nil
1984
2344
  end
1985
2345
 
1986
2346
  # Internal node pool for RBTree.
@@ -2094,12 +2454,22 @@ class RBTree::AutoShrinkNodePool < RBTree::NodePool
2094
2454
  @min_active_in_interval = @active_nodes if @active_nodes < @min_active_in_interval
2095
2455
 
2096
2456
  @check_count += 1
2097
-
2457
+
2098
2458
  perform_maintenance if @check_count >= @check_interval
2099
2459
 
2100
2460
  super if @pool.size < @current_target_capacity
2101
2461
  end
2102
2462
 
2463
+ # Accounts for nodes dropped in bulk, so that the live-node count comes back down
2464
+ # after a `clear` instead of inflating the observed fluctuation range.
2465
+ #
2466
+ # @param count [Integer] the number of nodes discarded
2467
+ # @return [void]
2468
+ def discard(count)
2469
+ @active_nodes -= count
2470
+ @min_active_in_interval = @active_nodes if @active_nodes < @min_active_in_interval
2471
+ end
2472
+
2103
2473
  private
2104
2474
 
2105
2475
  def perform_maintenance