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 }