001/*
002 * Copyright (C) 2007 The Guava Authors
003 *
004 * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
005 * in compliance with the License. You may obtain a copy of the License at
006 *
007 * http://www.apache.org/licenses/LICENSE-2.0
008 *
009 * Unless required by applicable law or agreed to in writing, software distributed under the License
010 * is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
011 * or implied. See the License for the specific language governing permissions and limitations under
012 * the License.
013 */
014
015package com.google.common.collect;
016
017import static com.google.common.base.Preconditions.checkArgument;
018import static com.google.common.base.Preconditions.checkNotNull;
019import static com.google.common.collect.CollectPreconditions.checkNonnegative;
020import static com.google.common.collect.CollectPreconditions.checkRemove;
021
022import com.google.common.annotations.GwtCompatible;
023import com.google.common.annotations.GwtIncompatible;
024import com.google.common.primitives.Ints;
025import com.google.errorprone.annotations.CanIgnoreReturnValue;
026import java.io.IOException;
027import java.io.ObjectInputStream;
028import java.io.ObjectOutputStream;
029import java.io.Serializable;
030import java.util.Arrays;
031import java.util.Iterator;
032import java.util.NoSuchElementException;
033import java.util.function.ObjIntConsumer;
034import javax.annotation.CheckForNull;
035
036/**
037 * Multiset implementation specialized for enum elements, supporting all single-element operations
038 * in O(1).
039 *
040 * <p>See the Guava User Guide article on <a href=
041 * "https://github.com/google/guava/wiki/NewCollectionTypesExplained#multiset">{@code Multiset}</a>.
042 *
043 * @author Jared Levy
044 * @since 2.0
045 */
046@GwtCompatible(emulated = true)
047@ElementTypesAreNonnullByDefault
048public final class EnumMultiset<E extends Enum<E>> extends AbstractMultiset<E>
049    implements Serializable {
050  /** Creates an empty {@code EnumMultiset}. */
051  public static <E extends Enum<E>> EnumMultiset<E> create(Class<E> type) {
052    return new EnumMultiset<E>(type);
053  }
054
055  /**
056   * Creates a new {@code EnumMultiset} containing the specified elements.
057   *
058   * <p>This implementation is highly efficient when {@code elements} is itself a {@link Multiset}.
059   *
060   * @param elements the elements that the multiset should contain
061   * @throws IllegalArgumentException if {@code elements} is empty
062   */
063  public static <E extends Enum<E>> EnumMultiset<E> create(Iterable<E> elements) {
064    Iterator<E> iterator = elements.iterator();
065    checkArgument(iterator.hasNext(), "EnumMultiset constructor passed empty Iterable");
066    EnumMultiset<E> multiset = new EnumMultiset<>(iterator.next().getDeclaringClass());
067    Iterables.addAll(multiset, elements);
068    return multiset;
069  }
070
071  /**
072   * Returns a new {@code EnumMultiset} instance containing the given elements. Unlike {@link
073   * EnumMultiset#create(Iterable)}, this method does not produce an exception on an empty iterable.
074   *
075   * @since 14.0
076   */
077  public static <E extends Enum<E>> EnumMultiset<E> create(Iterable<E> elements, Class<E> type) {
078    EnumMultiset<E> result = create(type);
079    Iterables.addAll(result, elements);
080    return result;
081  }
082
083  private transient Class<E> type;
084  private transient E[] enumConstants;
085  private transient int[] counts;
086  private transient int distinctElements;
087  private transient long size;
088
089  /** Creates an empty {@code EnumMultiset}. */
090  private EnumMultiset(Class<E> type) {
091    this.type = type;
092    checkArgument(type.isEnum());
093    this.enumConstants = type.getEnumConstants();
094    this.counts = new int[enumConstants.length];
095  }
096
097  private boolean isActuallyE(@CheckForNull Object o) {
098    if (o instanceof Enum) {
099      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. Otherwise, throws
108   * either a NullPointerException or a ClassCastException as appropriate.
109   */
110  private void checkIsE(Object element) {
111    checkNotNull(element);
112    if (!isActuallyE(element)) {
113      throw new ClassCastException("Expected an " + type + " but got " + element);
114    }
115  }
116
117  @Override
118  int distinctElements() {
119    return distinctElements;
120  }
121
122  @Override
123  public int size() {
124    return Ints.saturatedCast(size);
125  }
126
127  @Override
128  public int count(@CheckForNull Object element) {
129    // isActuallyE checks for null, but we check explicitly to help nullness checkers.
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(@CheckForNull Object element, int occurrences) {
162    // isActuallyE checks for null, but we check explicitly to help nullness checkers.
163    if (element == null || !isActuallyE(element)) {
164      return 0;
165    }
166    Enum<?> e = (Enum<?>) element;
167    checkNonnegative(occurrences, "occurrences");
168    if (occurrences == 0) {
169      return count(element);
170    }
171    int index = e.ordinal();
172    int oldCount = counts[index];
173    if (oldCount == 0) {
174      return 0;
175    } else if (oldCount <= occurrences) {
176      counts[index] = 0;
177      distinctElements--;
178      size -= oldCount;
179    } else {
180      counts[index] = oldCount - occurrences;
181      size -= occurrences;
182    }
183    return oldCount;
184  }
185
186  // Modification Operations
187  @CanIgnoreReturnValue
188  @Override
189  public int setCount(E element, int count) {
190    checkIsE(element);
191    checkNonnegative(count, "count");
192    int index = element.ordinal();
193    int oldCount = counts[index];
194    counts[index] = count;
195    size += count - oldCount;
196    if (oldCount == 0 && count > 0) {
197      distinctElements++;
198    } else if (oldCount > 0 && count == 0) {
199      distinctElements--;
200    }
201    return oldCount;
202  }
203
204  @Override
205  public void clear() {
206    Arrays.fill(counts, 0);
207    size = 0;
208    distinctElements = 0;
209  }
210
211  abstract class Itr<T> implements Iterator<T> {
212    int index = 0;
213    int toRemove = -1;
214
215    abstract T output(int index);
216
217    @Override
218    public boolean hasNext() {
219      for (; index < enumConstants.length; index++) {
220        if (counts[index] > 0) {
221          return true;
222        }
223      }
224      return false;
225    }
226
227    @Override
228    public T next() {
229      if (!hasNext()) {
230        throw new NoSuchElementException();
231      }
232      T result = output(index);
233      toRemove = index;
234      index++;
235      return result;
236    }
237
238    @Override
239    public void remove() {
240      checkRemove(toRemove >= 0);
241      if (counts[toRemove] > 0) {
242        distinctElements--;
243        size -= counts[toRemove];
244        counts[toRemove] = 0;
245      }
246      toRemove = -1;
247    }
248  }
249
250  @Override
251  Iterator<E> elementIterator() {
252    return new Itr<E>() {
253      @Override
254      E output(int index) {
255        return enumConstants[index];
256      }
257    };
258  }
259
260  @Override
261  Iterator<Entry<E>> entryIterator() {
262    return new Itr<Entry<E>>() {
263      @Override
264      Entry<E> output(final int index) {
265        return new Multisets.AbstractEntry<E>() {
266          @Override
267          public E getElement() {
268            return enumConstants[index];
269          }
270
271          @Override
272          public int getCount() {
273            return counts[index];
274          }
275        };
276      }
277    };
278  }
279
280  @Override
281  public void forEachEntry(ObjIntConsumer<? super E> action) {
282    checkNotNull(action);
283    for (int i = 0; i < enumConstants.length; i++) {
284      if (counts[i] > 0) {
285        action.accept(enumConstants[i], counts[i]);
286      }
287    }
288  }
289
290  @Override
291  public Iterator<E> iterator() {
292    return Multisets.iteratorImpl(this);
293  }
294
295  @GwtIncompatible // java.io.ObjectOutputStream
296  private void writeObject(ObjectOutputStream stream) throws IOException {
297    stream.defaultWriteObject();
298    stream.writeObject(type);
299    Serialization.writeMultiset(this, stream);
300  }
301
302  /**
303   * @serialData the {@code Class<E>} for the enum type, the number of distinct elements, the first
304   *     element, its count, the second element, its count, and so on
305   */
306  @GwtIncompatible // java.io.ObjectInputStream
307  private void readObject(ObjectInputStream stream) throws IOException, ClassNotFoundException {
308    stream.defaultReadObject();
309    @SuppressWarnings("unchecked") // reading data stored by writeObject
310    Class<E> localType = (Class<E>) stream.readObject();
311    type = localType;
312    enumConstants = type.getEnumConstants();
313    counts = new int[enumConstants.length];
314    Serialization.populateMultiset(this, stream);
315  }
316
317  @GwtIncompatible // Not needed in emulated source
318  private static final long serialVersionUID = 0;
319}