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.checkNotNull;
021    import static com.google.common.base.Preconditions.checkState;
022    import static com.google.common.collect.BstSide.LEFT;
023    import static com.google.common.collect.BstSide.RIGHT;
024    
025    import com.google.common.annotations.GwtCompatible;
026    import com.google.common.annotations.GwtIncompatible;
027    import com.google.common.base.Optional;
028    import com.google.common.primitives.Ints;
029    
030    import java.io.IOException;
031    import java.io.ObjectInputStream;
032    import java.io.ObjectOutputStream;
033    import java.io.Serializable;
034    import java.util.Comparator;
035    import java.util.ConcurrentModificationException;
036    import java.util.Iterator;
037    
038    import javax.annotation.Nullable;
039    
040    /**
041     * A multiset which maintains the ordering of its elements, according to either
042     * their natural order or an explicit {@link Comparator}. In all cases, this
043     * implementation uses {@link Comparable#compareTo} or {@link
044     * Comparator#compare} instead of {@link Object#equals} to determine
045     * equivalence of instances.
046     *
047     * <p><b>Warning:</b> The comparison must be <i>consistent with equals</i> as
048     * explained by the {@link Comparable} class specification. Otherwise, the
049     * resulting multiset will violate the {@link java.util.Collection} contract,
050     * which is specified in terms of {@link Object#equals}.
051     *
052     * <p>See the Guava User Guide article on <a href=
053     * "http://code.google.com/p/guava-libraries/wiki/NewCollectionTypesExplained#Multiset">
054     * {@code Multiset}</a>.
055     *
056     * @author Louis Wasserman
057     * @author Jared Levy
058     * @since 2.0 (imported from Google Collections Library)
059     */
060    @GwtCompatible(emulated = true)
061    public final class TreeMultiset<E> extends AbstractSortedMultiset<E>
062        implements Serializable {
063    
064      /**
065       * Creates a new, empty multiset, sorted according to the elements' natural
066       * order. All elements inserted into the multiset must implement the
067       * {@code Comparable} interface. Furthermore, all such elements must be
068       * <i>mutually comparable</i>: {@code e1.compareTo(e2)} must not throw a
069       * {@code ClassCastException} for any elements {@code e1} and {@code e2} in
070       * the multiset. If the user attempts to add an element to the multiset that
071       * violates this constraint (for example, the user attempts to add a string
072       * element to a set whose elements are integers), the {@code add(Object)}
073       * call will throw a {@code ClassCastException}.
074       *
075       * <p>The type specification is {@code <E extends Comparable>}, instead of the
076       * more specific {@code <E extends Comparable<? super E>>}, to support
077       * classes defined without generics.
078       */
079      public static <E extends Comparable> TreeMultiset<E> create() {
080        return new TreeMultiset<E>(Ordering.natural());
081      }
082    
083      /**
084       * Creates a new, empty multiset, sorted according to the specified
085       * comparator. All elements inserted into the multiset must be <i>mutually
086       * comparable</i> by the specified comparator: {@code comparator.compare(e1,
087       * e2)} must not throw a {@code ClassCastException} for any elements {@code
088       * e1} and {@code e2} in the multiset. If the user attempts to add an element
089       * to the multiset that violates this constraint, the {@code add(Object)} call
090       * will throw a {@code ClassCastException}.
091       *
092       * @param comparator the comparator that will be used to sort this multiset. A
093       *     null value indicates that the elements' <i>natural ordering</i> should
094       *     be used.
095       */
096      @SuppressWarnings("unchecked")
097      public static <E> TreeMultiset<E> create(
098          @Nullable Comparator<? super E> comparator) {
099        return (comparator == null)
100               ? new TreeMultiset<E>((Comparator) Ordering.natural())
101               : new TreeMultiset<E>(comparator);
102      }
103    
104      /**
105       * Creates an empty multiset containing the given initial elements, sorted
106       * according to the elements' natural order.
107       *
108       * <p>This implementation is highly efficient when {@code elements} is itself
109       * a {@link Multiset}.
110       *
111       * <p>The type specification is {@code <E extends Comparable>}, instead of the
112       * more specific {@code <E extends Comparable<? super E>>}, to support
113       * classes defined without generics.
114       */
115      public static <E extends Comparable> TreeMultiset<E> create(
116          Iterable<? extends E> elements) {
117        TreeMultiset<E> multiset = create();
118        Iterables.addAll(multiset, elements);
119        return multiset;
120      }
121    
122      /**
123       * Returns an iterator over the elements contained in this collection.
124       */
125      @Override
126      public Iterator<E> iterator() {
127        // Needed to avoid Javadoc bug.
128        return super.iterator();
129      }
130    
131      private TreeMultiset(Comparator<? super E> comparator) {
132        super(comparator);
133        this.range = GeneralRange.all(comparator);
134        this.rootReference = new Reference<Node<E>>();
135      }
136    
137      private TreeMultiset(GeneralRange<E> range, Reference<Node<E>> root) {
138        super(range.comparator());
139        this.range = range;
140        this.rootReference = root;
141      }
142    
143      E checkElement(Object o) {
144        @SuppressWarnings("unchecked")
145        E cast = (E) o;
146        // Make sure the object is accepted by the comparator (e.g., the right type, possibly non-null).
147        comparator.compare(cast, cast);
148        return cast;
149      }
150    
151      private transient final GeneralRange<E> range;
152    
153      private transient final Reference<Node<E>> rootReference;
154    
155      static final class Reference<T> {
156        T value;
157    
158        public Reference() {}
159    
160        public T get() {
161          return value;
162        }
163    
164        public boolean compareAndSet(T expected, T newValue) {
165          if (value == expected) {
166            value = newValue;
167            return true;
168          }
169          return false;
170        }
171      }
172    
173      @Override
174      int distinctElements() {
175        Node<E> root = rootReference.get();
176        return Ints.checkedCast(BstRangeOps.totalInRange(distinctAggregate(), range, root));
177      }
178    
179      @Override
180      public int size() {
181        Node<E> root = rootReference.get();
182        return Ints.saturatedCast(BstRangeOps.totalInRange(sizeAggregate(), range, root));
183      }
184    
185      @Override
186      public int count(@Nullable Object element) {
187        try {
188          E e = checkElement(element);
189          if (range.contains(e)) {
190            Node<E> node = BstOperations.seek(comparator(), rootReference.get(), e);
191            return countOrZero(node);
192          }
193          return 0;
194        } catch (ClassCastException e) {
195          return 0;
196        } catch (NullPointerException e) {
197          return 0;
198        }
199      }
200    
201      private int mutate(@Nullable E e, MultisetModifier modifier) {
202        BstMutationRule<E, Node<E>> mutationRule = BstMutationRule.createRule(
203            modifier,
204            BstCountBasedBalancePolicies.
205              <E, Node<E>>singleRebalancePolicy(distinctAggregate()),
206            nodeFactory());
207        BstMutationResult<E, Node<E>> mutationResult =
208            BstOperations.mutate(comparator(), mutationRule, rootReference.get(), e);
209        if (!rootReference.compareAndSet(
210            mutationResult.getOriginalRoot(), mutationResult.getChangedRoot())) {
211          throw new ConcurrentModificationException();
212        }
213        Node<E> original = mutationResult.getOriginalTarget();
214        return countOrZero(original);
215      }
216    
217      @Override
218      public int add(E element, int occurrences) {
219        checkElement(element);
220        if (occurrences == 0) {
221          return count(element);
222        }
223        checkArgument(range.contains(element));
224        return mutate(element, new AddModifier(occurrences));
225      }
226    
227      @Override
228      public int remove(@Nullable Object element, int occurrences) {
229        if (occurrences == 0) {
230          return count(element);
231        }
232        try {
233          E e = checkElement(element);
234          return range.contains(e) ? mutate(e, new RemoveModifier(occurrences)) : 0;
235        } catch (ClassCastException e) {
236          return 0;
237        } catch (NullPointerException e) {
238          return 0;
239        }
240      }
241    
242      @Override
243      public boolean setCount(E element, int oldCount, int newCount) {
244        checkElement(element);
245        checkArgument(range.contains(element));
246        return mutate(element, new ConditionalSetCountModifier(oldCount, newCount))
247            == oldCount;
248      }
249    
250      @Override
251      public int setCount(E element, int count) {
252        checkElement(element);
253        checkArgument(range.contains(element));
254        return mutate(element, new SetCountModifier(count));
255      }
256    
257      private BstPathFactory<Node<E>, BstInOrderPath<Node<E>>> pathFactory() {
258        return BstInOrderPath.inOrderFactory();
259      }
260    
261      @Override
262      Iterator<Entry<E>> entryIterator() {
263        Node<E> root = rootReference.get();
264        final BstInOrderPath<Node<E>> startingPath =
265            BstRangeOps.furthestPath(range, LEFT, pathFactory(), root);
266        return iteratorInDirection(startingPath, RIGHT);
267      }
268    
269      @Override
270      Iterator<Entry<E>> descendingEntryIterator() {
271        Node<E> root = rootReference.get();
272        final BstInOrderPath<Node<E>> startingPath =
273            BstRangeOps.furthestPath(range, RIGHT, pathFactory(), root);
274        return iteratorInDirection(startingPath, LEFT);
275      }
276    
277      private Iterator<Entry<E>> iteratorInDirection(
278          @Nullable BstInOrderPath<Node<E>> start, final BstSide direction) {
279        final Iterator<BstInOrderPath<Node<E>>> pathIterator =
280            new AbstractSequentialIterator<BstInOrderPath<Node<E>>>(start) {
281              @Override
282              protected BstInOrderPath<Node<E>> computeNext(BstInOrderPath<Node<E>> previous) {
283                if (!previous.hasNext(direction)) {
284                  return null;
285                }
286                BstInOrderPath<Node<E>> next = previous.next(direction);
287                // TODO(user): only check against one side
288                return range.contains(next.getTip().getKey()) ? next : null;
289              }
290            };
291        return new Iterator<Entry<E>>() {
292          final ToRemove<E> toRemove = new ToRemove<E>();
293    
294          @Override
295          public boolean hasNext() {
296            return pathIterator.hasNext();
297          }
298    
299          @Override
300          public Entry<E> next() {
301            BstInOrderPath<Node<E>> path = pathIterator.next();
302            return new LiveEntry(
303                toRemove.setAndGet(path.getTip().getKey()), path.getTip().elemCount());
304          }
305    
306          @Override
307          public void remove() {
308            setCount(toRemove.getAndClear(), 0);
309          }
310        };
311      }
312    
313      // If we were ever to resurrect AbstractRemovableIterator, we could use it instead.
314      private static final class ToRemove<E> {
315        @Nullable Optional<E> element;
316    
317        E setAndGet(@Nullable E element) {
318          this.element = Optional.fromNullable(element);
319          return element;
320        }
321    
322        E getAndClear() {
323          checkState(element != null);
324          E returnValue = element.orNull();
325          element = null;
326          return returnValue;
327        }
328      }
329    
330      class LiveEntry extends Multisets.AbstractEntry<E> {
331        private Node<E> expectedRoot;
332        private final E element;
333        private int count;
334    
335        private LiveEntry(E element, int count) {
336          this.expectedRoot = rootReference.get();
337          this.element = element;
338          this.count = count;
339        }
340    
341        @Override
342        public E getElement() {
343          return element;
344        }
345    
346        @Override
347        public int getCount() {
348          if (rootReference.get() == expectedRoot) {
349            return count;
350          } else {
351            // check for updates
352            expectedRoot = rootReference.get();
353            return count = TreeMultiset.this.count(element);
354          }
355        }
356      }
357    
358      @Override
359      public void clear() {
360        Node<E> root = rootReference.get();
361        Node<E> cleared = BstRangeOps.minusRange(range,
362            BstCountBasedBalancePolicies.<E, Node<E>>fullRebalancePolicy(distinctAggregate()),
363            nodeFactory(), root);
364        if (!rootReference.compareAndSet(root, cleared)) {
365          throw new ConcurrentModificationException();
366        }
367      }
368    
369      @Override
370      public SortedMultiset<E> headMultiset(E upperBound, BoundType boundType) {
371        checkNotNull(upperBound);
372        return new TreeMultiset<E>(
373            range.intersect(GeneralRange.upTo(comparator, upperBound, boundType)), rootReference);
374      }
375    
376      @Override
377      public SortedMultiset<E> tailMultiset(E lowerBound, BoundType boundType) {
378        checkNotNull(lowerBound);
379        return new TreeMultiset<E>(
380            range.intersect(GeneralRange.downTo(comparator, lowerBound, boundType)), rootReference);
381      }
382    
383      /**
384       * {@inheritDoc}
385       *
386       * @since 11.0
387       */
388      @Override
389      public Comparator<? super E> comparator() {
390        return super.comparator();
391      }
392    
393      private static final class Node<E> extends BstNode<E, Node<E>> implements Serializable {
394        private final long size;
395        private final int distinct;
396    
397        private Node(E key, int elemCount, @Nullable Node<E> left,
398            @Nullable Node<E> right) {
399          super(key, left, right);
400          checkArgument(elemCount > 0);
401          this.size = elemCount + sizeOrZero(left) + sizeOrZero(right);
402          this.distinct = 1 + distinctOrZero(left) + distinctOrZero(right);
403        }
404    
405        int elemCount() {
406          long result = size - sizeOrZero(childOrNull(LEFT))
407              - sizeOrZero(childOrNull(RIGHT));
408          return Ints.checkedCast(result);
409        }
410    
411        private Node(E key, int elemCount) {
412          this(key, elemCount, null, null);
413        }
414    
415        private static final long serialVersionUID = 0;
416      }
417    
418      private static long sizeOrZero(@Nullable Node<?> node) {
419        return (node == null) ? 0 : node.size;
420      }
421    
422      private static int distinctOrZero(@Nullable Node<?> node) {
423        return (node == null) ? 0 : node.distinct;
424      }
425    
426      private static int countOrZero(@Nullable Node<?> entry) {
427        return (entry == null) ? 0 : entry.elemCount();
428      }
429    
430      @SuppressWarnings("unchecked")
431      private BstAggregate<Node<E>> distinctAggregate() {
432        return (BstAggregate) DISTINCT_AGGREGATE;
433      }
434    
435      private static final BstAggregate<Node<Object>> DISTINCT_AGGREGATE =
436          new BstAggregate<Node<Object>>() {
437        @Override
438        public int entryValue(Node<Object> entry) {
439          return 1;
440        }
441    
442        @Override
443        public long treeValue(@Nullable Node<Object> tree) {
444          return distinctOrZero(tree);
445        }
446      };
447    
448      @SuppressWarnings("unchecked")
449      private BstAggregate<Node<E>> sizeAggregate() {
450        return (BstAggregate) SIZE_AGGREGATE;
451      }
452    
453      private static final BstAggregate<Node<Object>> SIZE_AGGREGATE =
454          new BstAggregate<Node<Object>>() {
455            @Override
456            public int entryValue(Node<Object> entry) {
457              return entry.elemCount();
458            }
459    
460            @Override
461            public long treeValue(@Nullable Node<Object> tree) {
462              return sizeOrZero(tree);
463            }
464          };
465    
466      @SuppressWarnings("unchecked")
467      private BstNodeFactory<Node<E>> nodeFactory() {
468        return (BstNodeFactory) NODE_FACTORY;
469      }
470    
471      private static final BstNodeFactory<Node<Object>> NODE_FACTORY =
472          new BstNodeFactory<Node<Object>>() {
473            @Override
474            public Node<Object> createNode(Node<Object> source, @Nullable Node<Object> left,
475                @Nullable Node<Object> right) {
476              return new Node<Object>(source.getKey(), source.elemCount(), left, right);
477            }
478          };
479    
480      private abstract class MultisetModifier implements BstModifier<E, Node<E>> {
481        abstract int newCount(int oldCount);
482    
483        @Nullable
484        @Override
485        public BstModificationResult<Node<E>> modify(E key, @Nullable Node<E> originalEntry) {
486          int oldCount = countOrZero(originalEntry);
487          int newCount = newCount(oldCount);
488          if (oldCount == newCount) {
489            return BstModificationResult.identity(originalEntry);
490          } else if (newCount == 0) {
491            return BstModificationResult.rebalancingChange(originalEntry, null);
492          } else if (oldCount == 0) {
493            return BstModificationResult.rebalancingChange(null, new Node<E>(key, newCount));
494          } else {
495            return BstModificationResult.rebuildingChange(originalEntry,
496                new Node<E>(originalEntry.getKey(), newCount));
497          }
498        }
499      }
500    
501      private final class AddModifier extends MultisetModifier {
502        private final int countToAdd;
503    
504        private AddModifier(int countToAdd) {
505          checkArgument(countToAdd > 0);
506          this.countToAdd = countToAdd;
507        }
508    
509        @Override
510        int newCount(int oldCount) {
511          checkArgument(countToAdd <= Integer.MAX_VALUE - oldCount, "Cannot add this many elements");
512          return oldCount + countToAdd;
513        }
514      }
515    
516      private final class RemoveModifier extends MultisetModifier {
517        private final int countToRemove;
518    
519        private RemoveModifier(int countToRemove) {
520          checkArgument(countToRemove > 0);
521          this.countToRemove = countToRemove;
522        }
523    
524        @Override
525        int newCount(int oldCount) {
526          return Math.max(0, oldCount - countToRemove);
527        }
528      }
529    
530      private final class SetCountModifier extends MultisetModifier {
531        private final int countToSet;
532    
533        private SetCountModifier(int countToSet) {
534          checkArgument(countToSet >= 0);
535          this.countToSet = countToSet;
536        }
537    
538        @Override
539        int newCount(int oldCount) {
540          return countToSet;
541        }
542      }
543    
544      private final class ConditionalSetCountModifier extends MultisetModifier {
545        private final int expectedCount;
546        private final int setCount;
547    
548        private ConditionalSetCountModifier(int expectedCount, int setCount) {
549          checkArgument(setCount >= 0 & expectedCount >= 0);
550          this.expectedCount = expectedCount;
551          this.setCount = setCount;
552        }
553    
554        @Override
555        int newCount(int oldCount) {
556          return (oldCount == expectedCount) ? setCount : oldCount;
557        }
558      }
559    
560      /*
561       * TODO(jlevy): Decide whether entrySet() should return entries with an
562       * equals() method that calls the comparator to compare the two keys. If that
563       * change is made, AbstractMultiset.equals() can simply check whether two
564       * multisets have equal entry sets.
565       */
566    
567      /**
568       * @serialData the comparator, the number of distinct elements, the first
569       *     element, its count, the second element, its count, and so on
570       */
571      @GwtIncompatible("java.io.ObjectOutputStream")
572      private void writeObject(ObjectOutputStream stream) throws IOException {
573        stream.defaultWriteObject();
574        stream.writeObject(elementSet().comparator());
575        Serialization.writeMultiset(this, stream);
576      }
577    
578      @GwtIncompatible("java.io.ObjectInputStream")
579      private void readObject(ObjectInputStream stream)
580          throws IOException, ClassNotFoundException {
581        stream.defaultReadObject();
582        @SuppressWarnings("unchecked") // reading data stored by writeObject
583        Comparator<? super E> comparator = (Comparator<? super E>) stream.readObject();
584        Serialization.getFieldSetter(AbstractSortedMultiset.class, "comparator").set(this, comparator);
585        Serialization.getFieldSetter(TreeMultiset.class, "range").set(this,
586            GeneralRange.all(comparator));
587        Serialization.getFieldSetter(TreeMultiset.class, "rootReference").set(this,
588            new Reference<Node<E>>());
589        Serialization.populateMultiset(this, stream);
590      }
591    
592      @GwtIncompatible("not needed in emulated source")
593      private static final long serialVersionUID = 1;
594    }