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.function.ObjIntConsumer;
34 import javax.annotation.CheckForNull;
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 @ElementTypesAreNonnullByDefault
49 public final class EnumMultiset<E extends Enum<E>> extends AbstractMultiset<E>
50 implements Serializable {
51 /** Creates an empty {@code EnumMultiset}. */
52 public static <E extends Enum<E>> EnumMultiset<E> create(Class<E> type) {
53 return new EnumMultiset<E>(type);
54 }
55
56 /**
57 * Creates a new {@code EnumMultiset} containing the specified elements.
58 *
59 * <p>This implementation is highly efficient when {@code elements} is itself a {@link Multiset}.
60 *
61 * @param elements the elements that the multiset should contain
62 * @throws IllegalArgumentException if {@code elements} is empty
63 */
64 public static <E extends Enum<E>> EnumMultiset<E> create(Iterable<E> elements) {
65 Iterator<E> iterator = elements.iterator();
66 checkArgument(iterator.hasNext(), "EnumMultiset constructor passed empty Iterable");
67 EnumMultiset<E> multiset = new EnumMultiset<>(iterator.next().getDeclaringClass());
68 Iterables.addAll(multiset, elements);
69 return multiset;
70 }
71
72 /**
73 * Returns a new {@code EnumMultiset} instance containing the given elements. Unlike {@link
74 * EnumMultiset#create(Iterable)}, this method does not produce an exception on an empty iterable.
75 *
76 * @since 14.0
77 */
78 public static <E extends Enum<E>> EnumMultiset<E> create(Iterable<E> elements, Class<E> type) {
79 EnumMultiset<E> result = create(type);
80 Iterables.addAll(result, elements);
81 return result;
82 }
83
84 private transient Class<E> type;
85 private transient E[] enumConstants;
86 private transient int[] counts;
87 private transient int distinctElements;
88 private transient long size;
89
90 /** Creates an empty {@code EnumMultiset}. */
91 private EnumMultiset(Class<E> type) {
92 this.type = type;
93 checkArgument(type.isEnum());
94 this.enumConstants = type.getEnumConstants();
95 this.counts = new int[enumConstants.length];
96 }
97
98 private boolean isActuallyE(@CheckForNull Object o) {
99 if (o instanceof Enum) {
100 Enum<?> e = (Enum<?>) o;
101 int index = e.ordinal();
102 return index < enumConstants.length && enumConstants[index] == e;
103 }
104 return false;
105 }
106
107 /**
108 * Returns {@code element} cast to {@code E}, if it actually is a nonnull E. Otherwise, throws
109 * either a NullPointerException or a ClassCastException as appropriate.
110 */
111 private void checkIsE(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(@CheckForNull Object element) {
130 // isActuallyE checks for null, but we check explicitly to help nullness checkers.
131 if (element == null || !isActuallyE(element)) {
132 return 0;
133 }
134 Enum<?> e = (Enum<?>) element;
135 return counts[e.ordinal()];
136 }
137
138 // Modification Operations
139 @CanIgnoreReturnValue
140 @Override
141 public int add(E element, int occurrences) {
142 checkIsE(element);
143 checkNonnegative(occurrences, "occurrences");
144 if (occurrences == 0) {
145 return count(element);
146 }
147 int index = element.ordinal();
148 int oldCount = counts[index];
149 long newCount = (long) oldCount + occurrences;
150 checkArgument(newCount <= Integer.MAX_VALUE, "too many occurrences: %s", newCount);
151 counts[index] = (int) newCount;
152 if (oldCount == 0) {
153 distinctElements++;
154 }
155 size += occurrences;
156 return oldCount;
157 }
158
159 // Modification Operations
160 @CanIgnoreReturnValue
161 @Override
162 public int remove(@CheckForNull Object element, int occurrences) {
163 // isActuallyE checks for null, but we check explicitly to help nullness checkers.
164 if (element == null || !isActuallyE(element)) {
165 return 0;
166 }
167 Enum<?> e = (Enum<?>) element;
168 checkNonnegative(occurrences, "occurrences");
169 if (occurrences == 0) {
170 return count(element);
171 }
172 int index = e.ordinal();
173 int oldCount = counts[index];
174 if (oldCount == 0) {
175 return 0;
176 } else if (oldCount <= occurrences) {
177 counts[index] = 0;
178 distinctElements--;
179 size -= oldCount;
180 } else {
181 counts[index] = oldCount - occurrences;
182 size -= occurrences;
183 }
184 return oldCount;
185 }
186
187 // Modification Operations
188 @CanIgnoreReturnValue
189 @Override
190 public int setCount(E element, int count) {
191 checkIsE(element);
192 checkNonnegative(count, "count");
193 int index = element.ordinal();
194 int oldCount = counts[index];
195 counts[index] = count;
196 size += count - oldCount;
197 if (oldCount == 0 && count > 0) {
198 distinctElements++;
199 } else if (oldCount > 0 && count == 0) {
200 distinctElements--;
201 }
202 return oldCount;
203 }
204
205 @Override
206 public void clear() {
207 Arrays.fill(counts, 0);
208 size = 0;
209 distinctElements = 0;
210 }
211
212 abstract class Itr<T> implements Iterator<T> {
213 int index = 0;
214 int toRemove = -1;
215
216 abstract T output(int index);
217
218 @Override
219 public boolean hasNext() {
220 for (; index < enumConstants.length; index++) {
221 if (counts[index] > 0) {
222 return true;
223 }
224 }
225 return false;
226 }
227
228 @Override
229 public T next() {
230 if (!hasNext()) {
231 throw new NoSuchElementException();
232 }
233 T result = output(index);
234 toRemove = index;
235 index++;
236 return result;
237 }
238
239 @Override
240 public void remove() {
241 checkRemove(toRemove >= 0);
242 if (counts[toRemove] > 0) {
243 distinctElements--;
244 size -= counts[toRemove];
245 counts[toRemove] = 0;
246 }
247 toRemove = -1;
248 }
249 }
250
251 @Override
252 Iterator<E> elementIterator() {
253 return new Itr<E>() {
254 @Override
255 E output(int index) {
256 return enumConstants[index];
257 }
258 };
259 }
260
261 @Override
262 Iterator<Entry<E>> entryIterator() {
263 return new Itr<Entry<E>>() {
264 @Override
265 Entry<E> output(final int index) {
266 return new Multisets.AbstractEntry<E>() {
267 @Override
268 public E getElement() {
269 return enumConstants[index];
270 }
271
272 @Override
273 public int getCount() {
274 return counts[index];
275 }
276 };
277 }
278 };
279 }
280
281 @Override
282 public void forEachEntry(ObjIntConsumer<? super E> action) {
283 checkNotNull(action);
284 for (int i = 0; i < enumConstants.length; i++) {
285 if (counts[i] > 0) {
286 action.accept(enumConstants[i], counts[i]);
287 }
288 }
289 }
290
291 @Override
292 public Iterator<E> iterator() {
293 return Multisets.iteratorImpl(this);
294 }
295
296 @GwtIncompatible // java.io.ObjectOutputStream
297 private void writeObject(ObjectOutputStream stream) throws IOException {
298 stream.defaultWriteObject();
299 stream.writeObject(type);
300 Serialization.writeMultiset(this, stream);
301 }
302
303 /**
304 * @serialData the {@code Class<E>} for the enum type, the number of distinct elements, the first
305 * element, its count, the second element, its count, and so on
306 */
307 @GwtIncompatible // java.io.ObjectInputStream
308 private void readObject(ObjectInputStream stream) throws IOException, ClassNotFoundException {
309 stream.defaultReadObject();
310 @SuppressWarnings("unchecked") // reading data stored by writeObject
311 Class<E> localType = (Class<E>) stream.readObject();
312 type = localType;
313 enumConstants = type.getEnumConstants();
314 counts = new int[enumConstants.length];
315 Serialization.populateMultiset(this, stream);
316 }
317
318 @GwtIncompatible // Not needed in emulated source
319 private static final long serialVersionUID = 0;
320 }