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