View Javadoc
1   /*
2    * Copyright (C) 2007 The Guava Authors
3    *
4    * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
5    * in compliance with the License. You may obtain a copy of the License at
6    *
7    * http://www.apache.org/licenses/LICENSE-2.0
8    *
9    * Unless required by applicable law or agreed to in writing, software distributed under the License
10   * is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
11   * or implied. See the License for the specific language governing permissions and limitations under
12   * the License.
13   */
14  
15  package com.google.common.collect;
16  
17  import static com.google.common.base.Preconditions.checkArgument;
18  import static com.google.common.base.Preconditions.checkNotNull;
19  import static com.google.common.collect.CollectPreconditions.checkNonnegative;
20  import static com.google.common.collect.CollectPreconditions.checkRemove;
21  
22  import com.google.common.annotations.GwtCompatible;
23  import com.google.common.annotations.GwtIncompatible;
24  import com.google.common.primitives.Ints;
25  import com.google.errorprone.annotations.CanIgnoreReturnValue;
26  import java.io.IOException;
27  import java.io.ObjectInputStream;
28  import java.io.ObjectOutputStream;
29  import java.io.Serializable;
30  import java.util.Arrays;
31  import java.util.Iterator;
32  import java.util.NoSuchElementException;
33  import java.util.Set;
34  import javax.annotation.Nullable;
35  
36  /**
37   * Multiset implementation specialized for enum elements, supporting all single-element operations
38   * in O(1).
39   *
40   * <p>See the Guava User Guide article on <a href=
41   * "https://github.com/google/guava/wiki/NewCollectionTypesExplained#multiset"> {@code
42   * Multiset}</a>.
43   *
44   * @author Jared Levy
45   * @since 2.0
46   */
47  @GwtCompatible(emulated = true)
48  public final class EnumMultiset<E extends Enum<E>> extends AbstractMultiset<E>
49      implements Serializable {
50    /** Creates an empty {@code EnumMultiset}. */
51    public static <E extends Enum<E>> EnumMultiset<E> create(Class<E> type) {
52      return new EnumMultiset<E>(type);
53    }
54  
55    /**
56     * Creates a new {@code EnumMultiset} containing the specified elements.
57     *
58     * <p>This implementation is highly efficient when {@code elements} is itself a {@link Multiset}.
59     *
60     * @param elements the elements that the multiset should contain
61     * @throws IllegalArgumentException if {@code elements} is empty
62     */
63    public static <E extends Enum<E>> EnumMultiset<E> create(Iterable<E> elements) {
64      Iterator<E> iterator = elements.iterator();
65      checkArgument(iterator.hasNext(), "EnumMultiset constructor passed empty Iterable");
66      EnumMultiset<E> multiset = new EnumMultiset<>(iterator.next().getDeclaringClass());
67      Iterables.addAll(multiset, elements);
68      return multiset;
69    }
70  
71    /**
72     * Returns a new {@code EnumMultiset} instance containing the given elements. Unlike {@link
73     * EnumMultiset#create(Iterable)}, this method does not produce an exception on an empty iterable.
74     *
75     * @since 14.0
76     */
77    public static <E extends Enum<E>> EnumMultiset<E> create(Iterable<E> elements, Class<E> type) {
78      EnumMultiset<E> result = create(type);
79      Iterables.addAll(result, elements);
80      return result;
81    }
82  
83    private transient Class<E> type;
84    private transient E[] enumConstants;
85    private transient int[] counts;
86    private transient int distinctElements;
87    private transient long size;
88  
89    /** Creates an empty {@code EnumMultiset}. */
90    private EnumMultiset(Class<E> type) {
91      this.type = type;
92      checkArgument(type.isEnum());
93      this.enumConstants = type.getEnumConstants();
94      this.counts = new int[enumConstants.length];
95    }
96    
97    private boolean isActuallyE(@Nullable Object o) {
98      if (o instanceof Enum) {
99        Enum<?> e = (Enum<?>) o;
100       int index = e.ordinal();
101       return index < enumConstants.length && enumConstants[index] == e;
102     }
103     return false;
104   }
105 
106   /**
107    * Returns {@code element} cast to {@code E}, if it actually is a nonnull E.
108    * Otherwise, throws either a NullPointerException or a ClassCastException as appropriate.
109    */
110   @SuppressWarnings("unchecked")
111   void checkIsE(@Nullable Object element) {
112     checkNotNull(element);
113     if (!isActuallyE(element)) {
114       throw new ClassCastException("Expected an " + type + " but got " + element);
115     }
116   }
117 
118   @Override
119   int distinctElements() {
120     return distinctElements;
121   }
122 
123   @Override
124   public int size() {
125     return Ints.saturatedCast(size);
126   }
127 
128   @Override
129   public int count(@Nullable Object element) {
130     if (element == null || !isActuallyE(element)) {
131       return 0;
132     }
133     Enum<?> e = (Enum<?>) element;
134     return counts[e.ordinal()];
135   }
136 
137   // Modification Operations
138   @CanIgnoreReturnValue
139   @Override
140   public int add(E element, int occurrences) {
141     checkIsE(element);
142     checkNonnegative(occurrences, "occurrences");
143     if (occurrences == 0) {
144       return count(element);
145     }
146     int index = element.ordinal();
147     int oldCount = counts[index];
148     long newCount = (long) oldCount + occurrences;
149     checkArgument(newCount <= Integer.MAX_VALUE, "too many occurrences: %s", newCount);
150     counts[index] = (int) newCount;
151     if (oldCount == 0) {
152       distinctElements++;
153     }
154     size += occurrences;
155     return oldCount;
156   }
157 
158   // Modification Operations
159   @CanIgnoreReturnValue
160   @Override
161   public int remove(@Nullable Object element, int occurrences) {
162     if (element == null || !isActuallyE(element)) {
163       return 0;
164     }
165     Enum<?> e = (Enum<?>) element;
166     checkNonnegative(occurrences, "occurrences");
167     if (occurrences == 0) {
168       return count(element);
169     }
170     int index = e.ordinal();
171     int oldCount = counts[index];
172     if (oldCount == 0) {
173       return 0;
174     } else if (oldCount <= occurrences) {
175       counts[index] = 0;
176       distinctElements--;
177       size -= oldCount;
178     } else {
179       counts[index] = oldCount - occurrences;
180       size -= occurrences;
181     }
182     return oldCount;
183   }
184 
185   // Modification Operations
186   @CanIgnoreReturnValue
187   @Override
188   public int setCount(E element, int count) {
189     checkIsE(element);
190     checkNonnegative(count, "count");
191     int index = element.ordinal();
192     int oldCount = counts[index];
193     counts[index] = count;
194     size += count - oldCount;
195     if (oldCount == 0 && count > 0) {
196       distinctElements++;
197     } else if (oldCount > 0 && count == 0) {
198       distinctElements--;
199     }
200     return oldCount;
201   }
202 
203   @Override
204   public void clear() {
205     Arrays.fill(counts, 0);
206     size = 0;
207     distinctElements = 0;
208   }
209 
210   abstract class Itr<T> implements Iterator<T> {
211     int index = 0;
212     int toRemove = -1;
213 
214     abstract T output(int index);
215 
216     @Override
217     public boolean hasNext() {
218       for (; index < enumConstants.length; index++) {
219         if (counts[index] > 0) {
220           return true;
221         }
222       }
223       return false;
224     }
225 
226     @Override
227     public T next() {
228       if (!hasNext()) {
229         throw new NoSuchElementException();
230       }
231       T result = output(index);
232       toRemove = index;
233       index++;
234       return result;
235     }
236 
237     @Override
238     public void remove() {
239       checkRemove(toRemove >= 0);
240       if (counts[toRemove] > 0) {
241         distinctElements--;
242         size -= counts[toRemove];
243         counts[toRemove] = 0;
244       }
245       toRemove = -1;
246     }
247   }
248 
249   @Override
250   Set<E> createElementSet() {
251     return new ElementSet() {
252 
253       @Override
254       public Iterator<E> iterator() {
255         return new Itr<E>() {
256           @Override
257           E output(int index) {
258             return enumConstants[index];
259           }
260         };
261       }
262     };
263   }
264 
265   @Override
266   Iterator<Entry<E>> entryIterator() {
267     return new Itr<Entry<E>>() {
268       @Override
269       Entry<E> output(final int index) {
270         return new Multisets.AbstractEntry<E>() {
271           @Override
272           public E getElement() {
273             return enumConstants[index];
274           }
275 
276           @Override
277           public int getCount() {
278             return counts[index];
279           }
280         };
281       }
282     };
283   }
284 
285   @GwtIncompatible // java.io.ObjectOutputStream
286   private void writeObject(ObjectOutputStream stream) throws IOException {
287     stream.defaultWriteObject();
288     stream.writeObject(type);
289     Serialization.writeMultiset(this, stream);
290   }
291 
292   /**
293    * @serialData the {@code Class<E>} for the enum type, the number of distinct elements, the first
294    *     element, its count, the second element, its count, and so on
295    */
296   @GwtIncompatible // java.io.ObjectInputStream
297   private void readObject(ObjectInputStream stream) throws IOException, ClassNotFoundException {
298     stream.defaultReadObject();
299     @SuppressWarnings("unchecked") // reading data stored by writeObject
300     Class<E> localType = (Class<E>) stream.readObject();
301     type = localType;
302     enumConstants = type.getEnumConstants();
303     counts = new int[enumConstants.length];
304     Serialization.populateMultiset(this, stream);
305   }
306 
307   @GwtIncompatible // Not needed in emulated source
308   private static final long serialVersionUID = 0;
309 }