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