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}