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