1 /*
2 * Copyright (C) 2007 The Guava Authors
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17 package com.google.common.collect;
18
19 import static com.google.common.base.Preconditions.checkArgument;
20 import static com.google.common.base.Preconditions.checkNotNull;
21 import static com.google.common.base.Preconditions.checkState;
22 import static com.google.common.collect.CollectPreconditions.checkNonnegative;
23 import static com.google.common.collect.NullnessCasts.uncheckedCastNullableTToT;
24 import static java.util.Objects.requireNonNull;
25
26 import com.google.common.annotations.GwtCompatible;
27 import com.google.common.annotations.GwtIncompatible;
28 import com.google.common.base.MoreObjects;
29 import com.google.common.primitives.Ints;
30 import com.google.errorprone.annotations.CanIgnoreReturnValue;
31 import java.io.IOException;
32 import java.io.ObjectInputStream;
33 import java.io.ObjectOutputStream;
34 import java.io.Serializable;
35 import java.util.Comparator;
36 import java.util.ConcurrentModificationException;
37 import java.util.Iterator;
38 import java.util.NoSuchElementException;
39 import java.util.function.ObjIntConsumer;
40 import javax.annotation.CheckForNull;
41 import org.checkerframework.checker.nullness.qual.Nullable;
42
43 /**
44 * A multiset which maintains the ordering of its elements, according to either their natural order
45 * or an explicit {@link Comparator}. In all cases, this implementation uses {@link
46 * Comparable#compareTo} or {@link Comparator#compare} instead of {@link Object#equals} to determine
47 * equivalence of instances.
48 *
49 * <p><b>Warning:</b> The comparison must be <i>consistent with equals</i> as explained by the
50 * {@link Comparable} class specification. Otherwise, the resulting multiset will violate the {@link
51 * java.util.Collection} contract, which is specified in terms of {@link Object#equals}.
52 *
53 * <p>See the Guava User Guide article on <a href=
54 * "https://github.com/google/guava/wiki/NewCollectionTypesExplained#multiset"> {@code
55 * Multiset}</a>.
56 *
57 * @author Louis Wasserman
58 * @author Jared Levy
59 * @since 2.0
60 */
61 @GwtCompatible(emulated = true)
62 @ElementTypesAreNonnullByDefault
63 public final class TreeMultiset<E extends @Nullable Object> extends AbstractSortedMultiset<E>
64 implements Serializable {
65
66 /**
67 * Creates a new, empty multiset, sorted according to the elements' natural order. All elements
68 * inserted into the multiset must implement the {@code Comparable} interface. Furthermore, all
69 * such elements must be <i>mutually comparable</i>: {@code e1.compareTo(e2)} must not throw a
70 * {@code ClassCastException} for any elements {@code e1} and {@code e2} in the multiset. If the
71 * user attempts to add an element to the multiset that violates this constraint (for example, the
72 * user attempts to add a string element to a set whose elements are integers), the {@code
73 * add(Object)} call will throw a {@code ClassCastException}.
74 *
75 * <p>The type specification is {@code <E extends Comparable>}, instead of the more specific
76 * {@code <E extends Comparable<? super E>>}, to support classes defined without generics.
77 */
78 public static <E extends Comparable> TreeMultiset<E> create() {
79 return new TreeMultiset<E>(Ordering.natural());
80 }
81
82 /**
83 * Creates a new, empty multiset, sorted according to the specified comparator. All elements
84 * inserted into the multiset must be <i>mutually comparable</i> by the specified comparator:
85 * {@code comparator.compare(e1, e2)} must not throw a {@code ClassCastException} for any elements
86 * {@code e1} and {@code e2} in the multiset. If the user attempts to add an element to the
87 * multiset that violates this constraint, the {@code add(Object)} call will throw a {@code
88 * ClassCastException}.
89 *
90 * @param comparator the comparator that will be used to sort this multiset. A null value
91 * indicates that the elements' <i>natural ordering</i> should be used.
92 */
93 @SuppressWarnings("unchecked")
94 public static <E extends @Nullable Object> TreeMultiset<E> create(
95 @CheckForNull Comparator<? super E> comparator) {
96 return (comparator == null)
97 ? new TreeMultiset<E>((Comparator) Ordering.natural())
98 : new TreeMultiset<E>(comparator);
99 }
100
101 /**
102 * Creates an empty multiset containing the given initial elements, sorted according to the
103 * elements' natural order.
104 *
105 * <p>This implementation is highly efficient when {@code elements} is itself a {@link Multiset}.
106 *
107 * <p>The type specification is {@code <E extends Comparable>}, instead of the more specific
108 * {@code <E extends Comparable<? super E>>}, to support classes defined without generics.
109 */
110 public static <E extends Comparable> TreeMultiset<E> create(Iterable<? extends E> elements) {
111 TreeMultiset<E> multiset = create();
112 Iterables.addAll(multiset, elements);
113 return multiset;
114 }
115
116 private final transient Reference<AvlNode<E>> rootReference;
117 private final transient GeneralRange<E> range;
118 private final transient AvlNode<E> header;
119
120 TreeMultiset(Reference<AvlNode<E>> rootReference, GeneralRange<E> range, AvlNode<E> endLink) {
121 super(range.comparator());
122 this.rootReference = rootReference;
123 this.range = range;
124 this.header = endLink;
125 }
126
127 TreeMultiset(Comparator<? super E> comparator) {
128 super(comparator);
129 this.range = GeneralRange.all(comparator);
130 this.header = new AvlNode<>();
131 successor(header, header);
132 this.rootReference = new Reference<>();
133 }
134
135 /** A function which can be summed across a subtree. */
136 private enum Aggregate {
137 SIZE {
138 @Override
139 int nodeAggregate(AvlNode<?> node) {
140 return node.elemCount;
141 }
142
143 @Override
144 long treeAggregate(@CheckForNull AvlNode<?> root) {
145 return (root == null) ? 0 : root.totalCount;
146 }
147 },
148 DISTINCT {
149 @Override
150 int nodeAggregate(AvlNode<?> node) {
151 return 1;
152 }
153
154 @Override
155 long treeAggregate(@CheckForNull AvlNode<?> root) {
156 return (root == null) ? 0 : root.distinctElements;
157 }
158 };
159
160 abstract int nodeAggregate(AvlNode<?> node);
161
162 abstract long treeAggregate(@CheckForNull AvlNode<?> root);
163 }
164
165 private long aggregateForEntries(Aggregate aggr) {
166 AvlNode<E> root = rootReference.get();
167 long total = aggr.treeAggregate(root);
168 if (range.hasLowerBound()) {
169 total -= aggregateBelowRange(aggr, root);
170 }
171 if (range.hasUpperBound()) {
172 total -= aggregateAboveRange(aggr, root);
173 }
174 return total;
175 }
176
177 private long aggregateBelowRange(Aggregate aggr, @CheckForNull AvlNode<E> node) {
178 if (node == null) {
179 return 0;
180 }
181 // The cast is safe because we call this method only if hasLowerBound().
182 int cmp =
183 comparator()
184 .compare(uncheckedCastNullableTToT(range.getLowerEndpoint()), node.getElement());
185 if (cmp < 0) {
186 return aggregateBelowRange(aggr, node.left);
187 } else if (cmp == 0) {
188 switch (range.getLowerBoundType()) {
189 case OPEN:
190 return aggr.nodeAggregate(node) + aggr.treeAggregate(node.left);
191 case CLOSED:
192 return aggr.treeAggregate(node.left);
193 default:
194 throw new AssertionError();
195 }
196 } else {
197 return aggr.treeAggregate(node.left)
198 + aggr.nodeAggregate(node)
199 + aggregateBelowRange(aggr, node.right);
200 }
201 }
202
203 private long aggregateAboveRange(Aggregate aggr, @CheckForNull AvlNode<E> node) {
204 if (node == null) {
205 return 0;
206 }
207 // The cast is safe because we call this method only if hasUpperBound().
208 int cmp =
209 comparator()
210 .compare(uncheckedCastNullableTToT(range.getUpperEndpoint()), node.getElement());
211 if (cmp > 0) {
212 return aggregateAboveRange(aggr, node.right);
213 } else if (cmp == 0) {
214 switch (range.getUpperBoundType()) {
215 case OPEN:
216 return aggr.nodeAggregate(node) + aggr.treeAggregate(node.right);
217 case CLOSED:
218 return aggr.treeAggregate(node.right);
219 default:
220 throw new AssertionError();
221 }
222 } else {
223 return aggr.treeAggregate(node.right)
224 + aggr.nodeAggregate(node)
225 + aggregateAboveRange(aggr, node.left);
226 }
227 }
228
229 @Override
230 public int size() {
231 return Ints.saturatedCast(aggregateForEntries(Aggregate.SIZE));
232 }
233
234 @Override
235 int distinctElements() {
236 return Ints.saturatedCast(aggregateForEntries(Aggregate.DISTINCT));
237 }
238
239 static int distinctElements(@CheckForNull AvlNode<?> node) {
240 return (node == null) ? 0 : node.distinctElements;
241 }
242
243 @Override
244 public int count(@CheckForNull Object element) {
245 try {
246 @SuppressWarnings("unchecked")
247 E e = (E) element;
248 AvlNode<E> root = rootReference.get();
249 if (!range.contains(e) || root == null) {
250 return 0;
251 }
252 return root.count(comparator(), e);
253 } catch (ClassCastException | NullPointerException e) {
254 return 0;
255 }
256 }
257
258 @CanIgnoreReturnValue
259 @Override
260 public int add(@ParametricNullness E element, int occurrences) {
261 checkNonnegative(occurrences, "occurrences");
262 if (occurrences == 0) {
263 return count(element);
264 }
265 checkArgument(range.contains(element));
266 AvlNode<E> root = rootReference.get();
267 if (root == null) {
268 comparator().compare(element, element);
269 AvlNode<E> newRoot = new AvlNode<E>(element, occurrences);
270 successor(header, newRoot, header);
271 rootReference.checkAndSet(root, newRoot);
272 return 0;
273 }
274 int[] result = new int[1]; // used as a mutable int reference to hold result
275 AvlNode<E> newRoot = root.add(comparator(), element, occurrences, result);
276 rootReference.checkAndSet(root, newRoot);
277 return result[0];
278 }
279
280 @CanIgnoreReturnValue
281 @Override
282 public int remove(@CheckForNull Object element, int occurrences) {
283 checkNonnegative(occurrences, "occurrences");
284 if (occurrences == 0) {
285 return count(element);
286 }
287 AvlNode<E> root = rootReference.get();
288 int[] result = new int[1]; // used as a mutable int reference to hold result
289 AvlNode<E> newRoot;
290 try {
291 @SuppressWarnings("unchecked")
292 E e = (E) element;
293 if (!range.contains(e) || root == null) {
294 return 0;
295 }
296 newRoot = root.remove(comparator(), e, occurrences, result);
297 } catch (ClassCastException | NullPointerException e) {
298 return 0;
299 }
300 rootReference.checkAndSet(root, newRoot);
301 return result[0];
302 }
303
304 @CanIgnoreReturnValue
305 @Override
306 public int setCount(@ParametricNullness E element, int count) {
307 checkNonnegative(count, "count");
308 if (!range.contains(element)) {
309 checkArgument(count == 0);
310 return 0;
311 }
312
313 AvlNode<E> root = rootReference.get();
314 if (root == null) {
315 if (count > 0) {
316 add(element, count);
317 }
318 return 0;
319 }
320 int[] result = new int[1]; // used as a mutable int reference to hold result
321 AvlNode<E> newRoot = root.setCount(comparator(), element, count, result);
322 rootReference.checkAndSet(root, newRoot);
323 return result[0];
324 }
325
326 @CanIgnoreReturnValue
327 @Override
328 public boolean setCount(@ParametricNullness E element, int oldCount, int newCount) {
329 checkNonnegative(newCount, "newCount");
330 checkNonnegative(oldCount, "oldCount");
331 checkArgument(range.contains(element));
332
333 AvlNode<E> root = rootReference.get();
334 if (root == null) {
335 if (oldCount == 0) {
336 if (newCount > 0) {
337 add(element, newCount);
338 }
339 return true;
340 } else {
341 return false;
342 }
343 }
344 int[] result = new int[1]; // used as a mutable int reference to hold result
345 AvlNode<E> newRoot = root.setCount(comparator(), element, oldCount, newCount, result);
346 rootReference.checkAndSet(root, newRoot);
347 return result[0] == oldCount;
348 }
349
350 @Override
351 public void clear() {
352 if (!range.hasLowerBound() && !range.hasUpperBound()) {
353 // We can do this in O(n) rather than removing one by one, which could force rebalancing.
354 for (AvlNode<E> current = header.succ(); current != header; ) {
355 AvlNode<E> next = current.succ();
356
357 current.elemCount = 0;
358 // Also clear these fields so that one deleted Entry doesn't retain all elements.
359 current.left = null;
360 current.right = null;
361 current.pred = null;
362 current.succ = null;
363
364 current = next;
365 }
366 successor(header, header);
367 rootReference.clear();
368 } else {
369 // TODO(cpovirk): Perhaps we can optimize in this case, too?
370 Iterators.clear(entryIterator());
371 }
372 }
373
374 private Entry<E> wrapEntry(final AvlNode<E> baseEntry) {
375 return new Multisets.AbstractEntry<E>() {
376 @Override
377 @ParametricNullness
378 public E getElement() {
379 return baseEntry.getElement();
380 }
381
382 @Override
383 public int getCount() {
384 int result = baseEntry.getCount();
385 if (result == 0) {
386 return count(getElement());
387 } else {
388 return result;
389 }
390 }
391 };
392 }
393
394 /** Returns the first node in the tree that is in range. */
395 @CheckForNull
396 private AvlNode<E> firstNode() {
397 AvlNode<E> root = rootReference.get();
398 if (root == null) {
399 return null;
400 }
401 AvlNode<E> node;
402 if (range.hasLowerBound()) {
403 // The cast is safe because of the hasLowerBound check.
404 E endpoint = uncheckedCastNullableTToT(range.getLowerEndpoint());
405 node = root.ceiling(comparator(), endpoint);
406 if (node == null) {
407 return null;
408 }
409 if (range.getLowerBoundType() == BoundType.OPEN
410 && comparator().compare(endpoint, node.getElement()) == 0) {
411 node = node.succ();
412 }
413 } else {
414 node = header.succ();
415 }
416 return (node == header || !range.contains(node.getElement())) ? null : node;
417 }
418
419 @CheckForNull
420 private AvlNode<E> lastNode() {
421 AvlNode<E> root = rootReference.get();
422 if (root == null) {
423 return null;
424 }
425 AvlNode<E> node;
426 if (range.hasUpperBound()) {
427 // The cast is safe because of the hasUpperBound check.
428 E endpoint = uncheckedCastNullableTToT(range.getUpperEndpoint());
429 node = root.floor(comparator(), endpoint);
430 if (node == null) {
431 return null;
432 }
433 if (range.getUpperBoundType() == BoundType.OPEN
434 && comparator().compare(endpoint, node.getElement()) == 0) {
435 node = node.pred();
436 }
437 } else {
438 node = header.pred();
439 }
440 return (node == header || !range.contains(node.getElement())) ? null : node;
441 }
442
443 @Override
444 Iterator<E> elementIterator() {
445 return Multisets.elementIterator(entryIterator());
446 }
447
448 @Override
449 Iterator<Entry<E>> entryIterator() {
450 return new Iterator<Entry<E>>() {
451 @CheckForNull AvlNode<E> current = firstNode();
452 @CheckForNull Entry<E> prevEntry;
453
454 @Override
455 public boolean hasNext() {
456 if (current == null) {
457 return false;
458 } else if (range.tooHigh(current.getElement())) {
459 current = null;
460 return false;
461 } else {
462 return true;
463 }
464 }
465
466 @Override
467 public Entry<E> next() {
468 if (!hasNext()) {
469 throw new NoSuchElementException();
470 }
471 // requireNonNull is safe because current is only nulled out after iteration is complete.
472 Entry<E> result = wrapEntry(requireNonNull(current));
473 prevEntry = result;
474 if (current.succ() == header) {
475 current = null;
476 } else {
477 current = current.succ();
478 }
479 return result;
480 }
481
482 @Override
483 public void remove() {
484 checkState(prevEntry != null, "no calls to next() since the last call to remove()");
485 setCount(prevEntry.getElement(), 0);
486 prevEntry = null;
487 }
488 };
489 }
490
491 @Override
492 Iterator<Entry<E>> descendingEntryIterator() {
493 return new Iterator<Entry<E>>() {
494 @CheckForNull AvlNode<E> current = lastNode();
495 @CheckForNull Entry<E> prevEntry = null;
496
497 @Override
498 public boolean hasNext() {
499 if (current == null) {
500 return false;
501 } else if (range.tooLow(current.getElement())) {
502 current = null;
503 return false;
504 } else {
505 return true;
506 }
507 }
508
509 @Override
510 public Entry<E> next() {
511 if (!hasNext()) {
512 throw new NoSuchElementException();
513 }
514 // requireNonNull is safe because current is only nulled out after iteration is complete.
515 requireNonNull(current);
516 Entry<E> result = wrapEntry(current);
517 prevEntry = result;
518 if (current.pred() == header) {
519 current = null;
520 } else {
521 current = current.pred();
522 }
523 return result;
524 }
525
526 @Override
527 public void remove() {
528 checkState(prevEntry != null, "no calls to next() since the last call to remove()");
529 setCount(prevEntry.getElement(), 0);
530 prevEntry = null;
531 }
532 };
533 }
534
535 @Override
536 public void forEachEntry(ObjIntConsumer<? super E> action) {
537 checkNotNull(action);
538 for (AvlNode<E> node = firstNode();
539 node != header && node != null && !range.tooHigh(node.getElement());
540 node = node.succ()) {
541 action.accept(node.getElement(), node.getCount());
542 }
543 }
544
545 @Override
546 public Iterator<E> iterator() {
547 return Multisets.iteratorImpl(this);
548 }
549
550 @Override
551 public SortedMultiset<E> headMultiset(@ParametricNullness E upperBound, BoundType boundType) {
552 return new TreeMultiset<E>(
553 rootReference,
554 range.intersect(GeneralRange.upTo(comparator(), upperBound, boundType)),
555 header);
556 }
557
558 @Override
559 public SortedMultiset<E> tailMultiset(@ParametricNullness E lowerBound, BoundType boundType) {
560 return new TreeMultiset<E>(
561 rootReference,
562 range.intersect(GeneralRange.downTo(comparator(), lowerBound, boundType)),
563 header);
564 }
565
566 private static final class Reference<T> {
567 @CheckForNull private T value;
568
569 @CheckForNull
570 public T get() {
571 return value;
572 }
573
574 public void checkAndSet(@CheckForNull T expected, @CheckForNull T newValue) {
575 if (value != expected) {
576 throw new ConcurrentModificationException();
577 }
578 value = newValue;
579 }
580
581 void clear() {
582 value = null;
583 }
584 }
585
586 private static final class AvlNode<E extends @Nullable Object> {
587 /*
588 * For "normal" nodes, the type of this field is `E`, not `@Nullable E` (though note that E is a
589 * type that can include null, as in a TreeMultiset<@Nullable String>).
590 *
591 * For the header node, though, this field contains `null`, regardless of the type of the
592 * multiset.
593 *
594 * Most code that operates on an AvlNode never operates on the header node. Such code can access
595 * the elem field without a null check by calling getElement().
596 */
597 @CheckForNull private final E elem;
598
599 // elemCount is 0 iff this node has been deleted.
600 private int elemCount;
601
602 private int distinctElements;
603 private long totalCount;
604 private int height;
605 @CheckForNull private AvlNode<E> left;
606 @CheckForNull private AvlNode<E> right;
607 /*
608 * pred and succ are nullable after construction, but we always call successor() to initialize
609 * them immediately thereafter.
610 *
611 * They may be subsequently nulled out by TreeMultiset.clear(). I think that the only place that
612 * we can reference a node whose fields have been cleared is inside the iterator (and presumably
613 * only under concurrent modification).
614 *
615 * To access these fields when you know that they are not null, call the pred() and succ()
616 * methods, which perform null checks before returning the fields.
617 */
618 @CheckForNull private AvlNode<E> pred;
619 @CheckForNull private AvlNode<E> succ;
620
621 AvlNode(@ParametricNullness E elem, int elemCount) {
622 checkArgument(elemCount > 0);
623 this.elem = elem;
624 this.elemCount = elemCount;
625 this.totalCount = elemCount;
626 this.distinctElements = 1;
627 this.height = 1;
628 this.left = null;
629 this.right = null;
630 }
631
632 /** Constructor for the header node. */
633 AvlNode() {
634 this.elem = null;
635 this.elemCount = 1;
636 }
637
638 // For discussion of pred() and succ(), see the comment on the pred and succ fields.
639
640 private AvlNode<E> pred() {
641 return requireNonNull(pred);
642 }
643
644 private AvlNode<E> succ() {
645 return requireNonNull(succ);
646 }
647
648 int count(Comparator<? super E> comparator, @ParametricNullness E e) {
649 int cmp = comparator.compare(e, getElement());
650 if (cmp < 0) {
651 return (left == null) ? 0 : left.count(comparator, e);
652 } else if (cmp > 0) {
653 return (right == null) ? 0 : right.count(comparator, e);
654 } else {
655 return elemCount;
656 }
657 }
658
659 private AvlNode<E> addRightChild(@ParametricNullness E e, int count) {
660 right = new AvlNode<E>(e, count);
661 successor(this, right, succ());
662 height = Math.max(2, height);
663 distinctElements++;
664 totalCount += count;
665 return this;
666 }
667
668 private AvlNode<E> addLeftChild(@ParametricNullness E e, int count) {
669 left = new AvlNode<E>(e, count);
670 successor(pred(), left, this);
671 height = Math.max(2, height);
672 distinctElements++;
673 totalCount += count;
674 return this;
675 }
676
677 AvlNode<E> add(
678 Comparator<? super E> comparator, @ParametricNullness E e, int count, int[] result) {
679 /*
680 * It speeds things up considerably to unconditionally add count to totalCount here,
681 * but that destroys failure atomicity in the case of count overflow. =(
682 */
683 int cmp = comparator.compare(e, getElement());
684 if (cmp < 0) {
685 AvlNode<E> initLeft = left;
686 if (initLeft == null) {
687 result[0] = 0;
688 return addLeftChild(e, count);
689 }
690 int initHeight = initLeft.height;
691
692 left = initLeft.add(comparator, e, count, result);
693 if (result[0] == 0) {
694 distinctElements++;
695 }
696 this.totalCount += count;
697 return (left.height == initHeight) ? this : rebalance();
698 } else if (cmp > 0) {
699 AvlNode<E> initRight = right;
700 if (initRight == null) {
701 result[0] = 0;
702 return addRightChild(e, count);
703 }
704 int initHeight = initRight.height;
705
706 right = initRight.add(comparator, e, count, result);
707 if (result[0] == 0) {
708 distinctElements++;
709 }
710 this.totalCount += count;
711 return (right.height == initHeight) ? this : rebalance();
712 }
713
714 // adding count to me! No rebalance possible.
715 result[0] = elemCount;
716 long resultCount = (long) elemCount + count;
717 checkArgument(resultCount <= Integer.MAX_VALUE);
718 this.elemCount += count;
719 this.totalCount += count;
720 return this;
721 }
722
723 @CheckForNull
724 AvlNode<E> remove(
725 Comparator<? super E> comparator, @ParametricNullness E e, int count, int[] result) {
726 int cmp = comparator.compare(e, getElement());
727 if (cmp < 0) {
728 AvlNode<E> initLeft = left;
729 if (initLeft == null) {
730 result[0] = 0;
731 return this;
732 }
733
734 left = initLeft.remove(comparator, e, count, result);
735
736 if (result[0] > 0) {
737 if (count >= result[0]) {
738 this.distinctElements--;
739 this.totalCount -= result[0];
740 } else {
741 this.totalCount -= count;
742 }
743 }
744 return (result[0] == 0) ? this : rebalance();
745 } else if (cmp > 0) {
746 AvlNode<E> initRight = right;
747 if (initRight == null) {
748 result[0] = 0;
749 return this;
750 }
751
752 right = initRight.remove(comparator, e, count, result);
753
754 if (result[0] > 0) {
755 if (count >= result[0]) {
756 this.distinctElements--;
757 this.totalCount -= result[0];
758 } else {
759 this.totalCount -= count;
760 }
761 }
762 return rebalance();
763 }
764
765 // removing count from me!
766 result[0] = elemCount;
767 if (count >= elemCount) {
768 return deleteMe();
769 } else {
770 this.elemCount -= count;
771 this.totalCount -= count;
772 return this;
773 }
774 }
775
776 @CheckForNull
777 AvlNode<E> setCount(
778 Comparator<? super E> comparator, @ParametricNullness E e, int count, int[] result) {
779 int cmp = comparator.compare(e, getElement());
780 if (cmp < 0) {
781 AvlNode<E> initLeft = left;
782 if (initLeft == null) {
783 result[0] = 0;
784 return (count > 0) ? addLeftChild(e, count) : this;
785 }
786
787 left = initLeft.setCount(comparator, e, count, result);
788
789 if (count == 0 && result[0] != 0) {
790 this.distinctElements--;
791 } else if (count > 0 && result[0] == 0) {
792 this.distinctElements++;
793 }
794
795 this.totalCount += count - result[0];
796 return rebalance();
797 } else if (cmp > 0) {
798 AvlNode<E> initRight = right;
799 if (initRight == null) {
800 result[0] = 0;
801 return (count > 0) ? addRightChild(e, count) : this;
802 }
803
804 right = initRight.setCount(comparator, e, count, result);
805
806 if (count == 0 && result[0] != 0) {
807 this.distinctElements--;
808 } else if (count > 0 && result[0] == 0) {
809 this.distinctElements++;
810 }
811
812 this.totalCount += count - result[0];
813 return rebalance();
814 }
815
816 // setting my count
817 result[0] = elemCount;
818 if (count == 0) {
819 return deleteMe();
820 }
821 this.totalCount += count - elemCount;
822 this.elemCount = count;
823 return this;
824 }
825
826 @CheckForNull
827 AvlNode<E> setCount(
828 Comparator<? super E> comparator,
829 @ParametricNullness E e,
830 int expectedCount,
831 int newCount,
832 int[] result) {
833 int cmp = comparator.compare(e, getElement());
834 if (cmp < 0) {
835 AvlNode<E> initLeft = left;
836 if (initLeft == null) {
837 result[0] = 0;
838 if (expectedCount == 0 && newCount > 0) {
839 return addLeftChild(e, newCount);
840 }
841 return this;
842 }
843
844 left = initLeft.setCount(comparator, e, expectedCount, newCount, result);
845
846 if (result[0] == expectedCount) {
847 if (newCount == 0 && result[0] != 0) {
848 this.distinctElements--;
849 } else if (newCount > 0 && result[0] == 0) {
850 this.distinctElements++;
851 }
852 this.totalCount += newCount - result[0];
853 }
854 return rebalance();
855 } else if (cmp > 0) {
856 AvlNode<E> initRight = right;
857 if (initRight == null) {
858 result[0] = 0;
859 if (expectedCount == 0 && newCount > 0) {
860 return addRightChild(e, newCount);
861 }
862 return this;
863 }
864
865 right = initRight.setCount(comparator, e, expectedCount, newCount, result);
866
867 if (result[0] == expectedCount) {
868 if (newCount == 0 && result[0] != 0) {
869 this.distinctElements--;
870 } else if (newCount > 0 && result[0] == 0) {
871 this.distinctElements++;
872 }
873 this.totalCount += newCount - result[0];
874 }
875 return rebalance();
876 }
877
878 // setting my count
879 result[0] = elemCount;
880 if (expectedCount == elemCount) {
881 if (newCount == 0) {
882 return deleteMe();
883 }
884 this.totalCount += newCount - elemCount;
885 this.elemCount = newCount;
886 }
887 return this;
888 }
889
890 @CheckForNull
891 private AvlNode<E> deleteMe() {
892 int oldElemCount = this.elemCount;
893 this.elemCount = 0;
894 successor(pred(), succ());
895 if (left == null) {
896 return right;
897 } else if (right == null) {
898 return left;
899 } else if (left.height >= right.height) {
900 AvlNode<E> newTop = pred();
901 // newTop is the maximum node in my left subtree
902 newTop.left = left.removeMax(newTop);
903 newTop.right = right;
904 newTop.distinctElements = distinctElements - 1;
905 newTop.totalCount = totalCount - oldElemCount;
906 return newTop.rebalance();
907 } else {
908 AvlNode<E> newTop = succ();
909 newTop.right = right.removeMin(newTop);
910 newTop.left = left;
911 newTop.distinctElements = distinctElements - 1;
912 newTop.totalCount = totalCount - oldElemCount;
913 return newTop.rebalance();
914 }
915 }
916
917 // Removes the minimum node from this subtree to be reused elsewhere
918 @CheckForNull
919 private AvlNode<E> removeMin(AvlNode<E> node) {
920 if (left == null) {
921 return right;
922 } else {
923 left = left.removeMin(node);
924 distinctElements--;
925 totalCount -= node.elemCount;
926 return rebalance();
927 }
928 }
929
930 // Removes the maximum node from this subtree to be reused elsewhere
931 @CheckForNull
932 private AvlNode<E> removeMax(AvlNode<E> node) {
933 if (right == null) {
934 return left;
935 } else {
936 right = right.removeMax(node);
937 distinctElements--;
938 totalCount -= node.elemCount;
939 return rebalance();
940 }
941 }
942
943 private void recomputeMultiset() {
944 this.distinctElements =
945 1 + TreeMultiset.distinctElements(left) + TreeMultiset.distinctElements(right);
946 this.totalCount = elemCount + totalCount(left) + totalCount(right);
947 }
948
949 private void recomputeHeight() {
950 this.height = 1 + Math.max(height(left), height(right));
951 }
952
953 private void recompute() {
954 recomputeMultiset();
955 recomputeHeight();
956 }
957
958 private AvlNode<E> rebalance() {
959 switch (balanceFactor()) {
960 case -2:
961 // requireNonNull is safe because right must exist in order to get a negative factor.
962 requireNonNull(right);
963 if (right.balanceFactor() > 0) {
964 right = right.rotateRight();
965 }
966 return rotateLeft();
967 case 2:
968 // requireNonNull is safe because left must exist in order to get a positive factor.
969 requireNonNull(left);
970 if (left.balanceFactor() < 0) {
971 left = left.rotateLeft();
972 }
973 return rotateRight();
974 default:
975 recomputeHeight();
976 return this;
977 }
978 }
979
980 private int balanceFactor() {
981 return height(left) - height(right);
982 }
983
984 private AvlNode<E> rotateLeft() {
985 checkState(right != null);
986 AvlNode<E> newTop = right;
987 this.right = newTop.left;
988 newTop.left = this;
989 newTop.totalCount = this.totalCount;
990 newTop.distinctElements = this.distinctElements;
991 this.recompute();
992 newTop.recomputeHeight();
993 return newTop;
994 }
995
996 private AvlNode<E> rotateRight() {
997 checkState(left != null);
998 AvlNode<E> newTop = left;
999 this.left = newTop.right;
1000 newTop.right = this;
1001 newTop.totalCount = this.totalCount;
1002 newTop.distinctElements = this.distinctElements;
1003 this.recompute();
1004 newTop.recomputeHeight();
1005 return newTop;
1006 }
1007
1008 private static long totalCount(@CheckForNull AvlNode<?> node) {
1009 return (node == null) ? 0 : node.totalCount;
1010 }
1011
1012 private static int height(@CheckForNull AvlNode<?> node) {
1013 return (node == null) ? 0 : node.height;
1014 }
1015
1016 @CheckForNull
1017 private AvlNode<E> ceiling(Comparator<? super E> comparator, @ParametricNullness E e) {
1018 int cmp = comparator.compare(e, getElement());
1019 if (cmp < 0) {
1020 return (left == null) ? this : MoreObjects.firstNonNull(left.ceiling(comparator, e), this);
1021 } else if (cmp == 0) {
1022 return this;
1023 } else {
1024 return (right == null) ? null : right.ceiling(comparator, e);
1025 }
1026 }
1027
1028 @CheckForNull
1029 private AvlNode<E> floor(Comparator<? super E> comparator, @ParametricNullness E e) {
1030 int cmp = comparator.compare(e, getElement());
1031 if (cmp > 0) {
1032 return (right == null) ? this : MoreObjects.firstNonNull(right.floor(comparator, e), this);
1033 } else if (cmp == 0) {
1034 return this;
1035 } else {
1036 return (left == null) ? null : left.floor(comparator, e);
1037 }
1038 }
1039
1040 @ParametricNullness
1041 E getElement() {
1042 // For discussion of this cast, see the comment on the elem field.
1043 return uncheckedCastNullableTToT(elem);
1044 }
1045
1046 int getCount() {
1047 return elemCount;
1048 }
1049
1050 @Override
1051 public String toString() {
1052 return Multisets.immutableEntry(getElement(), getCount()).toString();
1053 }
1054 }
1055
1056 private static <T extends @Nullable Object> void successor(AvlNode<T> a, AvlNode<T> b) {
1057 a.succ = b;
1058 b.pred = a;
1059 }
1060
1061 private static <T extends @Nullable Object> void successor(
1062 AvlNode<T> a, AvlNode<T> b, AvlNode<T> c) {
1063 successor(a, b);
1064 successor(b, c);
1065 }
1066
1067 /*
1068 * TODO(jlevy): Decide whether entrySet() should return entries with an equals() method that
1069 * calls the comparator to compare the two keys. If that change is made,
1070 * AbstractMultiset.equals() can simply check whether two multisets have equal entry sets.
1071 */
1072
1073 /**
1074 * @serialData the comparator, the number of distinct elements, the first element, its count, the
1075 * second element, its count, and so on
1076 */
1077 @GwtIncompatible // java.io.ObjectOutputStream
1078 private void writeObject(ObjectOutputStream stream) throws IOException {
1079 stream.defaultWriteObject();
1080 stream.writeObject(elementSet().comparator());
1081 Serialization.writeMultiset(this, stream);
1082 }
1083
1084 @GwtIncompatible // java.io.ObjectInputStream
1085 private void readObject(ObjectInputStream stream) throws IOException, ClassNotFoundException {
1086 stream.defaultReadObject();
1087 @SuppressWarnings("unchecked")
1088 // reading data stored by writeObject
1089 Comparator<? super E> comparator = (Comparator<? super E>) stream.readObject();
1090 Serialization.getFieldSetter(AbstractSortedMultiset.class, "comparator").set(this, comparator);
1091 Serialization.getFieldSetter(TreeMultiset.class, "range")
1092 .set(this, GeneralRange.all(comparator));
1093 Serialization.getFieldSetter(TreeMultiset.class, "rootReference")
1094 .set(this, new Reference<AvlNode<E>>());
1095 AvlNode<E> header = new AvlNode<>();
1096 Serialization.getFieldSetter(TreeMultiset.class, "header").set(this, header);
1097 successor(header, header);
1098 Serialization.populateMultiset(this, stream);
1099 }
1100
1101 @GwtIncompatible // not needed in emulated source
1102 private static final long serialVersionUID = 1;
1103 }