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