1/*
2 * Copyright (C) 2007 The Guava Authors
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17package com.google.common.collect;
18
19import static com.google.common.base.Preconditions.checkArgument;
20import static com.google.common.base.Preconditions.checkNotNull;
21import static com.google.common.base.Preconditions.checkState;
22import static com.google.common.collect.Multisets.checkNonnegative;
23
24import com.google.common.annotations.GwtCompatible;
25import com.google.common.annotations.GwtIncompatible;
26import com.google.common.primitives.Ints;
27
28import java.io.InvalidObjectException;
29import java.io.ObjectStreamException;
30import java.io.Serializable;
31import java.util.Collection;
32import java.util.ConcurrentModificationException;
33import java.util.Iterator;
34import java.util.Map;
35import java.util.Set;
36
37import javax.annotation.Nullable;
38
39/**
40 * Basic implementation of {@code Multiset<E>} backed by an instance of {@code
41 * Map<E, AtomicInteger>}.
42 *
43 * <p>For serialization to work, the subclass must specify explicit {@code
44 * readObject} and {@code writeObject} methods.
45 *
46 * @author Kevin Bourrillion
47 */
48@GwtCompatible(emulated = true)
49abstract class AbstractMapBasedMultiset<E> extends AbstractMultiset<E>
50    implements Serializable {
51
52  private transient Map<E, Count> backingMap;
53
54  /*
55   * Cache the size for efficiency. Using a long lets us avoid the need for
56   * overflow checking and ensures that size() will function correctly even if
57   * the multiset had once been larger than Integer.MAX_VALUE.
58   */
59  private transient long size;
60
61  /** Standard constructor. */
62  protected AbstractMapBasedMultiset(Map<E, Count> backingMap) {
63    this.backingMap = checkNotNull(backingMap);
64    this.size = super.size();
65  }
66
67  Map<E, Count> backingMap() {
68    return backingMap;
69  }
70
71  /** Used during deserialization only. The backing map must be empty. */
72  void setBackingMap(Map<E, Count> backingMap) {
73    this.backingMap = backingMap;
74  }
75
76  // Required Implementations
77
78  /**
79   * {@inheritDoc}
80   *
81   * <p>Invoking {@link Multiset.Entry#getCount} on an entry in the returned
82   * set always returns the current count of that element in the multiset, as
83   * opposed to the count at the time the entry was retrieved.
84   */
85  @Override
86  public Set<Multiset.Entry<E>> entrySet() {
87    return super.entrySet();
88  }
89
90  @Override
91  Iterator<Entry<E>> entryIterator() {
92    final Iterator<Map.Entry<E, Count>> backingEntries =
93        backingMap.entrySet().iterator();
94    return new Iterator<Multiset.Entry<E>>() {
95      Map.Entry<E, Count> toRemove;
96
97      @Override
98      public boolean hasNext() {
99        return backingEntries.hasNext();
100      }
101
102      @Override
103      public Multiset.Entry<E> next() {
104        final Map.Entry<E, Count> mapEntry = backingEntries.next();
105        toRemove = mapEntry;
106        return new Multisets.AbstractEntry<E>() {
107          @Override
108          public E getElement() {
109            return mapEntry.getKey();
110          }
111          @Override
112          public int getCount() {
113            int count = mapEntry.getValue().get();
114            if (count == 0) {
115              Count frequency = backingMap.get(getElement());
116              if (frequency != null) {
117                count = frequency.get();
118              }
119            }
120            return count;
121          }
122        };
123      }
124
125      @Override
126      public void remove() {
127        checkState(toRemove != null,
128            "no calls to next() since the last call to remove()");
129        size -= toRemove.getValue().getAndSet(0);
130        backingEntries.remove();
131        toRemove = null;
132      }
133    };
134  }
135
136  @Override
137  public void clear() {
138    for (Count frequency : backingMap.values()) {
139      frequency.set(0);
140    }
141    backingMap.clear();
142    size = 0L;
143  }
144
145  @Override
146  int distinctElements() {
147    return backingMap.size();
148  }
149
150  // Optimizations - Query Operations
151
152  @Override public int size() {
153    return Ints.saturatedCast(size);
154  }
155
156  @Override public Iterator<E> iterator() {
157    return new MapBasedMultisetIterator();
158  }
159
160  /*
161   * Not subclassing AbstractMultiset$MultisetIterator because next() needs to
162   * retrieve the Map.Entry<E, AtomicInteger> entry, which can then be used for
163   * a more efficient remove() call.
164   */
165  private class MapBasedMultisetIterator implements Iterator<E> {
166    final Iterator<Map.Entry<E, Count>> entryIterator;
167    Map.Entry<E, Count> currentEntry;
168    int occurrencesLeft;
169    boolean canRemove;
170
171    MapBasedMultisetIterator() {
172      this.entryIterator = backingMap.entrySet().iterator();
173    }
174
175    @Override
176    public boolean hasNext() {
177      return occurrencesLeft > 0 || entryIterator.hasNext();
178    }
179
180    @Override
181    public E next() {
182      if (occurrencesLeft == 0) {
183        currentEntry = entryIterator.next();
184        occurrencesLeft = currentEntry.getValue().get();
185      }
186      occurrencesLeft--;
187      canRemove = true;
188      return currentEntry.getKey();
189    }
190
191    @Override
192    public void remove() {
193      checkState(canRemove,
194          "no calls to next() since the last call to remove()");
195      int frequency = currentEntry.getValue().get();
196      if (frequency <= 0) {
197        throw new ConcurrentModificationException();
198      }
199      if (currentEntry.getValue().addAndGet(-1) == 0) {
200        entryIterator.remove();
201      }
202      size--;
203      canRemove = false;
204    }
205  }
206
207  @Override public int count(@Nullable Object element) {
208    try {
209      Count frequency = backingMap.get(element);
210      return (frequency == null) ? 0 : frequency.get();
211    } catch (NullPointerException e) {
212      return 0;
213    } catch (ClassCastException e) {
214      return 0;
215    }
216  }
217
218  // Optional Operations - Modification Operations
219
220  /**
221   * {@inheritDoc}
222   *
223   * @throws IllegalArgumentException if the call would result in more than
224   *     {@link Integer#MAX_VALUE} occurrences of {@code element} in this
225   *     multiset.
226   */
227  @Override public int add(@Nullable E element, int occurrences) {
228    if (occurrences == 0) {
229      return count(element);
230    }
231    checkArgument(
232        occurrences > 0, "occurrences cannot be negative: %s", occurrences);
233    Count frequency = backingMap.get(element);
234    int oldCount;
235    if (frequency == null) {
236      oldCount = 0;
237      backingMap.put(element, new Count(occurrences));
238    } else {
239      oldCount = frequency.get();
240      long newCount = (long) oldCount + (long) occurrences;
241      checkArgument(newCount <= Integer.MAX_VALUE,
242          "too many occurrences: %s", newCount);
243      frequency.getAndAdd(occurrences);
244    }
245    size += occurrences;
246    return oldCount;
247  }
248
249  @Override public int remove(@Nullable Object element, int occurrences) {
250    if (occurrences == 0) {
251      return count(element);
252    }
253    checkArgument(
254        occurrences > 0, "occurrences cannot be negative: %s", occurrences);
255    Count frequency = backingMap.get(element);
256    if (frequency == null) {
257      return 0;
258    }
259
260    int oldCount = frequency.get();
261
262    int numberRemoved;
263    if (oldCount > occurrences) {
264      numberRemoved = occurrences;
265    } else {
266      numberRemoved = oldCount;
267      backingMap.remove(element);
268    }
269
270    frequency.addAndGet(-numberRemoved);
271    size -= numberRemoved;
272    return oldCount;
273  }
274
275  // Roughly a 33% performance improvement over AbstractMultiset.setCount().
276  @Override public int setCount(E element, int count) {
277    checkNonnegative(count, "count");
278
279    Count existingCounter;
280    int oldCount;
281    if (count == 0) {
282      existingCounter = backingMap.remove(element);
283      oldCount = getAndSet(existingCounter, count);
284    } else {
285      existingCounter = backingMap.get(element);
286      oldCount = getAndSet(existingCounter, count);
287
288      if (existingCounter == null) {
289        backingMap.put(element, new Count(count));
290      }
291    }
292
293    size += (count - oldCount);
294    return oldCount;
295  }
296
297  private static int getAndSet(Count i, int count) {
298    if (i == null) {
299      return 0;
300    }
301
302    return i.getAndSet(count);
303  }
304
305  private int removeAllOccurrences(@Nullable Object element,
306      Map<E, Count> map) {
307    Count frequency = map.remove(element);
308    if (frequency == null) {
309      return 0;
310    }
311    int numberRemoved = frequency.getAndSet(0);
312    size -= numberRemoved;
313    return numberRemoved;
314  }
315
316  // Views
317
318  @Override Set<E> createElementSet() {
319    return new MapBasedElementSet(backingMap);
320  }
321
322  // TODO(user): once TreeMultiset is replaced with a SortedMultiset
323  // implementation, replace this with a subclass of Multisets.ElementSet.
324  class MapBasedElementSet extends ForwardingSet<E> {
325
326    // This mapping is the usually the same as 'backingMap', but can be a
327    // submap in some implementations.
328    private final Map<E, Count> map;
329    private final Set<E> delegate;
330
331    MapBasedElementSet(Map<E, Count> map) {
332      this.map = map;
333      delegate = map.keySet();
334    }
335
336    @Override protected Set<E> delegate() {
337      return delegate;
338    }
339
340    @Override public Iterator<E> iterator() {
341      final Iterator<Map.Entry<E, Count>> entries
342          = map.entrySet().iterator();
343      return new Iterator<E>() {
344        Map.Entry<E, Count> toRemove;
345
346        @Override
347        public boolean hasNext() {
348          return entries.hasNext();
349        }
350
351        @Override
352        public E next() {
353          toRemove = entries.next();
354          return toRemove.getKey();
355        }
356
357        @Override
358        public void remove() {
359          checkState(toRemove != null,
360              "no calls to next() since the last call to remove()");
361          size -= toRemove.getValue().getAndSet(0);
362          entries.remove();
363          toRemove = null;
364        }
365      };
366    }
367
368    @Override public boolean remove(Object element) {
369      return removeAllOccurrences(element, map) != 0;
370    }
371
372    @Override public boolean removeAll(Collection<?> elementsToRemove) {
373      return Iterators.removeAll(iterator(), elementsToRemove);
374    }
375
376    @Override public boolean retainAll(Collection<?> elementsToRetain) {
377      return Iterators.retainAll(iterator(), elementsToRetain);
378    }
379
380    @Override public void clear() {
381      if (map == backingMap) {
382        AbstractMapBasedMultiset.this.clear();
383      } else {
384        Iterator<E> i = iterator();
385        while (i.hasNext()) {
386          i.next();
387          i.remove();
388        }
389      }
390    }
391
392    public Map<E, Count> getMap() {
393      return map;
394    }
395  }
396
397  // Don't allow default serialization.
398  @GwtIncompatible("java.io.ObjectStreamException")
399  @SuppressWarnings("unused") // actually used during deserialization
400  private void readObjectNoData() throws ObjectStreamException {
401    throw new InvalidObjectException("Stream data required");
402  }
403
404  @GwtIncompatible("not needed in emulated source.")
405  private static final long serialVersionUID = -2250766705698539974L;
406}
407