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.collect.Multisets.checkNonnegative;
021    
022    import com.google.common.annotations.Beta;
023    import com.google.common.annotations.VisibleForTesting;
024    import com.google.common.collect.Serialization.FieldSetter;
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.AbstractSet;
032    import java.util.Iterator;
033    import java.util.List;
034    import java.util.Map;
035    import java.util.Set;
036    import java.util.concurrent.ConcurrentHashMap;
037    import java.util.concurrent.ConcurrentMap;
038    
039    import 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     */
049    public 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    }