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