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 javax.annotation.CheckForNull;
034
035/**
036 * Multiset implementation specialized for enum elements, supporting all single-element operations
037 * in O(1).
038 *
039 * <p>See the Guava User Guide article on <a href=
040 * "https://github.com/google/guava/wiki/NewCollectionTypesExplained#multiset">{@code Multiset}</a>.
041 *
042 * @author Jared Levy
043 * @since 2.0
044 */
045@GwtCompatible(emulated = true)
046@ElementTypesAreNonnullByDefault
047public final class EnumMultiset<E extends Enum<E>> extends AbstractMultiset<E>
048    implements Serializable {
049  /** Creates an empty {@code EnumMultiset}. */
050  public static <E extends Enum<E>> EnumMultiset<E> create(Class<E> type) {
051    return new EnumMultiset<E>(type);
052  }
053
054  /**
055   * Creates a new {@code EnumMultiset} containing the specified elements.
056   *
057   * <p>This implementation is highly efficient when {@code elements} is itself a {@link Multiset}.
058   *
059   * @param elements the elements that the multiset should contain
060   * @throws IllegalArgumentException if {@code elements} is empty
061   */
062  public static <E extends Enum<E>> EnumMultiset<E> create(Iterable<E> elements) {
063    Iterator<E> iterator = elements.iterator();
064    checkArgument(iterator.hasNext(), "EnumMultiset constructor passed empty Iterable");
065    EnumMultiset<E> multiset = new EnumMultiset<>(iterator.next().getDeclaringClass());
066    Iterables.addAll(multiset, elements);
067    return multiset;
068  }
069
070  /**
071   * Returns a new {@code EnumMultiset} instance containing the given elements. Unlike {@link
072   * EnumMultiset#create(Iterable)}, this method does not produce an exception on an empty iterable.
073   *
074   * @since 14.0
075   */
076  public static <E extends Enum<E>> EnumMultiset<E> create(Iterable<E> elements, Class<E> type) {
077    EnumMultiset<E> result = create(type);
078    Iterables.addAll(result, elements);
079    return result;
080  }
081
082  private transient Class<E> type;
083  private transient E[] enumConstants;
084  private transient int[] counts;
085  private transient int distinctElements;
086  private transient long size;
087
088  /** Creates an empty {@code EnumMultiset}. */
089  private EnumMultiset(Class<E> type) {
090    this.type = type;
091    checkArgument(type.isEnum());
092    this.enumConstants = type.getEnumConstants();
093    this.counts = new int[enumConstants.length];
094  }
095
096  private boolean isActuallyE(@CheckForNull Object o) {
097    if (o instanceof Enum) {
098      Enum<?> e = (Enum<?>) o;
099      int index = e.ordinal();
100      return index < enumConstants.length && enumConstants[index] == e;
101    }
102    return false;
103  }
104
105  /**
106   * Returns {@code element} cast to {@code E}, if it actually is a nonnull E. Otherwise, throws
107   * either a NullPointerException or a ClassCastException as appropriate.
108   */
109  private void checkIsE(Object element) {
110    checkNotNull(element);
111    if (!isActuallyE(element)) {
112      throw new ClassCastException("Expected an " + type + " but got " + element);
113    }
114  }
115
116  @Override
117  int distinctElements() {
118    return distinctElements;
119  }
120
121  @Override
122  public int size() {
123    return Ints.saturatedCast(size);
124  }
125
126  @Override
127  public int count(@CheckForNull Object element) {
128    // isActuallyE checks for null, but we check explicitly to help nullness checkers.
129    if (element == null || !isActuallyE(element)) {
130      return 0;
131    }
132    Enum<?> e = (Enum<?>) element;
133    return counts[e.ordinal()];
134  }
135
136  // Modification Operations
137  @CanIgnoreReturnValue
138  @Override
139  public int add(E element, int occurrences) {
140    checkIsE(element);
141    checkNonnegative(occurrences, "occurrences");
142    if (occurrences == 0) {
143      return count(element);
144    }
145    int index = element.ordinal();
146    int oldCount = counts[index];
147    long newCount = (long) oldCount + occurrences;
148    checkArgument(newCount <= Integer.MAX_VALUE, "too many occurrences: %s", newCount);
149    counts[index] = (int) newCount;
150    if (oldCount == 0) {
151      distinctElements++;
152    }
153    size += occurrences;
154    return oldCount;
155  }
156
157  // Modification Operations
158  @CanIgnoreReturnValue
159  @Override
160  public int remove(@CheckForNull Object element, int occurrences) {
161    // isActuallyE checks for null, but we check explicitly to help nullness checkers.
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  Iterator<E> elementIterator() {
251    return new Itr<E>() {
252      @Override
253      E output(int index) {
254        return enumConstants[index];
255      }
256    };
257  }
258
259  @Override
260  Iterator<Entry<E>> entryIterator() {
261    return new Itr<Entry<E>>() {
262      @Override
263      Entry<E> output(final int index) {
264        return new Multisets.AbstractEntry<E>() {
265          @Override
266          public E getElement() {
267            return enumConstants[index];
268          }
269
270          @Override
271          public int getCount() {
272            return counts[index];
273          }
274        };
275      }
276    };
277  }
278
279  @Override
280  public Iterator<E> iterator() {
281    return Multisets.iteratorImpl(this);
282  }
283
284  @GwtIncompatible // java.io.ObjectOutputStream
285  private void writeObject(ObjectOutputStream stream) throws IOException {
286    stream.defaultWriteObject();
287    stream.writeObject(type);
288    Serialization.writeMultiset(this, stream);
289  }
290
291  /**
292   * @serialData the {@code Class<E>} for the enum type, the number of distinct elements, the first
293   *     element, its count, the second element, its count, and so on
294   */
295  @GwtIncompatible // java.io.ObjectInputStream
296  private void readObject(ObjectInputStream stream) throws IOException, ClassNotFoundException {
297    stream.defaultReadObject();
298    @SuppressWarnings("unchecked") // reading data stored by writeObject
299    Class<E> localType = (Class<E>) stream.readObject();
300    type = localType;
301    enumConstants = type.getEnumConstants();
302    counts = new int[enumConstants.length];
303    Serialization.populateMultiset(this, stream);
304  }
305
306  @GwtIncompatible // Not needed in emulated source
307  private static final long serialVersionUID = 0;
308}