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.collect.Multisets.checkNonnegative;
021
022import com.google.common.annotations.Beta;
023import com.google.common.annotations.VisibleForTesting;
024import com.google.common.collect.Serialization.FieldSetter;
025import com.google.common.primitives.Ints;
026
027import java.io.IOException;
028import java.io.ObjectInputStream;
029import java.io.ObjectOutputStream;
030import java.io.Serializable;
031import java.util.AbstractSet;
032import java.util.Iterator;
033import java.util.List;
034import java.util.Map;
035import java.util.Set;
036import java.util.concurrent.ConcurrentHashMap;
037import java.util.concurrent.ConcurrentMap;
038
039import javax.annotation.Nullable;
040
041/**
042 * A multiset that supports concurrent modifications and that provides atomic
043 * versions of most {@code Multiset} operations (exceptions where noted). Null
044 * elements are not supported.
045 *
046 * @author Cliff L. Biffle
047 * @since 2 (imported from Google Collections Library)
048 */
049public final class ConcurrentHashMultiset<E> extends AbstractMultiset<E>
050    implements Serializable {
051  /*
052   * The ConcurrentHashMultiset's atomic operations are implemented in terms of
053   * ConcurrentMap's atomic operations. Many of them, such as add(E, int), are
054   * read-modify-write sequences, and so are implemented as loops that wrap
055   * ConcurrentMap's compare-and-set operations (like putIfAbsent).
056   *
057   * Invariant: all entries have a positive value. In particular, there are no
058   * entries with zero value. Some operations would fail if this was not the
059   * case.
060   */
061
062  /** The number of occurrences of each element. */
063  private final transient ConcurrentMap<E, Integer> countMap;
064
065  // This constant allows the deserialization code to set a final field. This
066  // holder class makes sure it is not initialized unless an instance is
067  // deserialized.
068  private static class FieldSettersHolder {
069    @SuppressWarnings("unchecked")
070    // eclipse doesn't like the raw type here, but it's harmless
071    static final FieldSetter<ConcurrentHashMultiset> COUNT_MAP_FIELD_SETTER
072        = Serialization.getFieldSetter(
073            ConcurrentHashMultiset.class, "countMap");
074  }
075
076  /**
077   * Creates a new, empty {@code ConcurrentHashMultiset} using the default
078   * initial capacity, load factor, and concurrency settings.
079   */
080  public static <E> ConcurrentHashMultiset<E> create() {
081    return new ConcurrentHashMultiset<E>(new ConcurrentHashMap<E, Integer>());
082  }
083
084  /**
085   * Creates a new {@code ConcurrentHashMultiset} containing the specified
086   * elements, using the default initial capacity, load factor, and concurrency
087   * settings.
088   *
089   * @param elements the elements that the multiset should contain
090   */
091  public static <E> ConcurrentHashMultiset<E> create(
092      Iterable<? extends E> elements) {
093    ConcurrentHashMultiset<E> multiset = ConcurrentHashMultiset.create();
094    Iterables.addAll(multiset, elements);
095    return multiset;
096  }
097
098  /**
099   * Creates a new, empty {@code ConcurrentHashMultiset} using {@code mapMaker}
100   * to construct the internal backing map.
101   *
102   * <p>If this {@link MapMaker} is configured to use entry eviction of any
103   * kind, this eviction applies to all occurrences of a given element as a
104   * single unit.
105   *
106   * <p>The returned multiset is serializable but any serialization caveats
107   * given in {@code MapMaker} apply.
108   *
109   * <p>Finally, soft/weak values can be used but are not very useful.
110   * Soft/weak keys on the other hand can be useful in some scenarios.
111   * 
112   * @since 7
113   */
114  @Beta
115  public static <E> ConcurrentHashMultiset<E> create(
116      GenericMapMaker<? super E, ? super Number> mapMaker) {
117    return new ConcurrentHashMultiset<E>(mapMaker.<E, Integer>makeMap());
118  }
119
120  /**
121   * Creates an instance using {@code countMap} to store elements and their
122   * counts.
123   *
124   * <p>This instance will assume ownership of {@code countMap}, and other code
125   * should not maintain references to the map or modify it in any way.
126   *
127   * @param countMap backing map for storing the elements in the multiset and
128   *     their counts. It must be empty.
129   * @throws IllegalArgumentException if {@code countMap} is not empty
130   */
131  @VisibleForTesting ConcurrentHashMultiset(
132      ConcurrentMap<E, Integer> countMap) {
133    checkArgument(countMap.isEmpty());
134    this.countMap = countMap;
135  }
136
137  // Query Operations
138
139  /**
140   * Returns the number of occurrences of {@code element} in this multiset.
141   *
142   * @param element the element to look for
143   * @return the nonnegative number of occurrences of the element
144   */
145  @Override public int count(@Nullable Object element) {
146    try {
147      return unbox(countMap.get(element));
148    } catch (NullPointerException e) {
149      return 0;
150    } catch (ClassCastException e) {
151      return 0;
152    }
153  }
154
155  /**
156   * {@inheritDoc}
157   *
158   * <p>If the data in the multiset is modified by any other threads during this
159   * method, it is undefined which (if any) of these modifications will be
160   * reflected in the result.
161   */
162  @Override public int size() {
163    long sum = 0L;
164    for (Integer value : countMap.values()) {
165      sum += value;
166    }
167    return Ints.saturatedCast(sum);
168  }
169
170  /*
171   * Note: the superclass toArray() methods assume that size() gives a correct
172   * answer, which ours does not.
173   */
174
175  @Override public Object[] toArray() {
176    return snapshot().toArray();
177  }
178
179  @Override public <T> T[] toArray(T[] array) {
180    return snapshot().toArray(array);
181  }
182
183  /*
184   * We'd love to use 'new ArrayList(this)' or 'list.addAll(this)', but
185   * either of these would recurse back to us again!
186   */
187  private List<E> snapshot() {
188    List<E> list = Lists.newArrayListWithExpectedSize(size());
189    for (Multiset.Entry<E> entry : entrySet()) {
190      E element = entry.getElement();
191      for (int i = entry.getCount(); i > 0; i--) {
192        list.add(element);
193      }
194    }
195    return list;
196  }
197
198  // Modification Operations
199
200  /**
201   * Adds a number of occurrences of the specified element to this multiset.
202   *
203   * @param element the element to add
204   * @param occurrences the number of occurrences to add
205   * @return the previous count of the element before the operation; possibly
206   *     zero
207   * @throws IllegalArgumentException if {@code occurrences} is negative, or if
208   *     the resulting amount would exceed {@link Integer#MAX_VALUE}
209   */
210  @Override public int add(E element, int occurrences) {
211    if (occurrences == 0) {
212      return count(element);
213    }
214    checkArgument(occurrences > 0, "Invalid occurrences: %s", occurrences);
215
216    while (true) {
217      int current = count(element);
218      if (current == 0) {
219        if (countMap.putIfAbsent(element, occurrences) == null) {
220          return 0;
221        }
222      } else {
223        checkArgument(occurrences <= Integer.MAX_VALUE - current,
224            "Overflow adding %s occurrences to a count of %s",
225            occurrences, current);
226        int next = current + occurrences;
227        if (countMap.replace(element, current, next)) {
228          return current;
229        }
230      }
231      // If we're still here, there was a race, so just try again.
232    }
233  }
234
235  /**
236   * Removes a number of occurrences of the specified element from this
237   * multiset. If the multiset contains fewer than this number of occurrences to
238   * begin with, all occurrences will be removed.
239   *
240   * @param element the element whose occurrences should be removed
241   * @param occurrences the number of occurrences of the element to remove
242   * @return the count of the element before the operation; possibly zero
243   * @throws IllegalArgumentException if {@code occurrences} is negative
244   */
245  @Override public int remove(@Nullable Object element, int occurrences) {
246    if (occurrences == 0) {
247      return count(element);
248    }
249    checkArgument(occurrences > 0, "Invalid occurrences: %s", occurrences);
250
251    while (true) {
252      int current = count(element);
253      if (current == 0) {
254        return 0;
255      }
256      if (occurrences >= current) {
257        if (countMap.remove(element, current)) {
258          return current;
259        }
260      } else {
261        // We know it's an "E" because it already exists in the map.
262        @SuppressWarnings("unchecked")
263        E casted = (E) element;
264
265        if (countMap.replace(casted, current, current - occurrences)) {
266          return current;
267        }
268      }
269      // If we're still here, there was a race, so just try again.
270    }
271  }
272
273  /**
274   * Removes <b>all</b> occurrences of the specified element from this multiset.
275   * This method complements {@link Multiset#remove(Object)}, which removes only
276   * one occurrence at a time.
277   *
278   * @param element the element whose occurrences should all be removed
279   * @return the number of occurrences successfully removed, possibly zero
280   */
281  private int removeAllOccurrences(@Nullable Object element) {
282    try {
283      return unbox(countMap.remove(element));
284    } catch (NullPointerException e) {
285      return 0;
286    } catch (ClassCastException e) {
287      return 0;
288    }
289  }
290
291  /**
292   * Removes exactly the specified number of occurrences of {@code element}, or
293   * makes no change if this is not possible.
294   *
295   * <p>This method, in contrast to {@link #remove(Object, int)}, has no effect
296   * when the element count is smaller than {@code occurrences}.
297   *
298   * @param element the element to remove
299   * @param occurrences the number of occurrences of {@code element} to remove
300   * @return {@code true} if the removal was possible (including if {@code
301   *     occurrences} is zero)
302   */
303  public boolean removeExactly(@Nullable Object element, int occurrences) {
304    if (occurrences == 0) {
305      return true;
306    }
307    checkArgument(occurrences > 0, "Invalid occurrences: %s", occurrences);
308
309    while (true) {
310      int current = count(element);
311      if (occurrences > current) {
312        return false;
313      }
314      if (occurrences == current) {
315        if (countMap.remove(element, occurrences)) {
316          return true;
317        }
318      } else {
319        @SuppressWarnings("unchecked") // it's in the map, must be an "E"
320        E casted = (E) element;
321        if (countMap.replace(casted, current, current - occurrences)) {
322          return true;
323        }
324      }
325      // If we're still here, there was a race, so just try again.
326    }
327  }
328
329  /**
330   * Adds or removes occurrences of {@code element} such that the {@link #count}
331   * of the element becomes {@code count}.
332   *
333   * @return the count of {@code element} in the multiset before this call
334   * @throws IllegalArgumentException if {@code count} is negative
335   */
336  @Override public int setCount(E element, int count) {
337    checkNonnegative(count, "count");
338    return (count == 0)
339        ? removeAllOccurrences(element)
340        : unbox(countMap.put(element, count));
341  }
342
343  /**
344   * Sets the number of occurrences of {@code element} to {@code newCount}, but
345   * only if the count is currently {@code oldCount}. If {@code element} does
346   * not appear in the multiset exactly {@code oldCount} times, no changes will
347   * be made.
348   *
349   * @return {@code true} if the change was successful. This usually indicates
350   *     that the multiset has been modified, but not always: in the case that
351   *     {@code oldCount == newCount}, the method will return {@code true} if
352   *     the condition was met.
353   * @throws IllegalArgumentException if {@code oldCount} or {@code newCount} is
354   *     negative
355   */
356  @Override public boolean setCount(E element, int oldCount, int newCount) {
357    checkNonnegative(oldCount, "oldCount");
358    checkNonnegative(newCount, "newCount");
359    if (newCount == 0) {
360      if (oldCount == 0) {
361        // No change to make, but must return true if the element is not present
362        return !countMap.containsKey(element);
363      } else {
364        return countMap.remove(element, oldCount);
365      }
366    }
367    if (oldCount == 0) {
368      return countMap.putIfAbsent(element, newCount) == null;
369    }
370    return countMap.replace(element, oldCount, newCount);
371  }
372
373  // Views
374
375  @Override Set<E> createElementSet() {
376    final Set<E> delegate = countMap.keySet();
377    return new ForwardingSet<E>() {
378      @Override protected Set<E> delegate() {
379        return delegate;
380      }
381      @Override public boolean remove(Object object) {
382        try {
383          return delegate.remove(object);
384        } catch (NullPointerException e) {
385          return false;
386        } catch (ClassCastException e) {
387          return false;
388        }
389      }
390    };
391  }
392
393  private transient EntrySet entrySet;
394
395  @Override public Set<Multiset.Entry<E>> entrySet() {
396    EntrySet result = entrySet;
397    if (result == null) {
398      entrySet = result = new EntrySet();
399    }
400    return result;
401  }
402
403  private class EntrySet extends AbstractSet<Multiset.Entry<E>> {
404    @Override public int size() {
405      return countMap.size();
406    }
407
408    @Override public boolean isEmpty() {
409      return countMap.isEmpty();
410    }
411
412    @Override public boolean contains(Object object) {
413      if (object instanceof Multiset.Entry) {
414        Multiset.Entry<?> entry = (Multiset.Entry<?>) object;
415        Object element = entry.getElement();
416        int entryCount = entry.getCount();
417        return entryCount > 0 && count(element) == entryCount;
418      }
419      return false;
420    }
421
422    @Override public Iterator<Multiset.Entry<E>> iterator() {
423      final Iterator<Map.Entry<E, Integer>> backingIterator
424          = countMap.entrySet().iterator();
425      return new Iterator<Multiset.Entry<E>>() {
426        @Override
427        public boolean hasNext() {
428          return backingIterator.hasNext();
429        }
430
431        @Override
432        public Multiset.Entry<E> next() {
433          Map.Entry<E, Integer> backingEntry = backingIterator.next();
434          return Multisets.immutableEntry(
435              backingEntry.getKey(), backingEntry.getValue());
436        }
437
438        @Override
439        public void remove() {
440          backingIterator.remove();
441        }
442      };
443    }
444
445    /*
446     * Note: the superclass toArray() methods assume that size() gives a correct
447     * answer, which ours does not.
448     */
449
450    @Override public Object[] toArray() {
451      return snapshot().toArray();
452    }
453
454    @Override public <T> T[] toArray(T[] array) {
455      return snapshot().toArray(array);
456    }
457
458    /*
459     * We'd love to use 'new ArrayList(this)' or 'list.addAll(this)', but
460     * either of these would recurse back to us again!
461     */
462    private List<Multiset.Entry<E>> snapshot() {
463      List<Multiset.Entry<E>> list = Lists.newArrayListWithExpectedSize(size());
464      for (Multiset.Entry<E> entry : this) {
465        list.add(entry);
466      }
467      return list;
468    }
469
470    @Override public boolean remove(Object object) {
471      if (object instanceof Multiset.Entry) {
472        Multiset.Entry<?> entry = (Multiset.Entry<?>) object;
473        Object element = entry.getElement();
474        int entryCount = entry.getCount();
475        return countMap.remove(element, entryCount);
476      }
477      return false;
478    }
479
480    @Override public void clear() {
481      countMap.clear();
482    }
483
484    /**
485     * The hash code is the same as countMap's, though the objects aren't equal.
486     */
487    @Override public int hashCode() {
488      return countMap.hashCode();
489    }
490  }
491
492  /**
493   * We use a special form of unboxing that treats null as zero.
494   */
495  private static int unbox(@Nullable Integer i) {
496    return (i == null) ? 0 : i;
497  }
498
499  /**
500   * @serialData the ConcurrentMap of elements and their counts.
501   */
502  private void writeObject(ObjectOutputStream stream) throws IOException {
503    stream.defaultWriteObject();
504    stream.writeObject(countMap);
505  }
506
507  private void readObject(ObjectInputStream stream)
508      throws IOException, ClassNotFoundException {
509    stream.defaultReadObject();
510    @SuppressWarnings("unchecked") // reading data stored by writeObject
511    ConcurrentMap<E, Integer> deserializedCountMap =
512        (ConcurrentMap<E, Integer>) stream.readObject();
513    FieldSettersHolder.COUNT_MAP_FIELD_SETTER.set(this, deserializedCountMap);
514  }
515
516  private static final long serialVersionUID = 1;
517}