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