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
042 * Multiset}</a>.
043 *
044 * @author Jared Levy
045 * @since 2.0
046 */
047@GwtCompatible(emulated = true)
048@ElementTypesAreNonnullByDefault
049public final class EnumMultiset<E extends Enum<E>> extends AbstractMultiset<E>
050    implements Serializable {
051  /** Creates an empty {@code EnumMultiset}. */
052  public static <E extends Enum<E>> EnumMultiset<E> create(Class<E> type) {
053    return new EnumMultiset<E>(type);
054  }
055
056  /**
057   * Creates a new {@code EnumMultiset} containing the specified elements.
058   *
059   * <p>This implementation is highly efficient when {@code elements} is itself a {@link Multiset}.
060   *
061   * @param elements the elements that the multiset should contain
062   * @throws IllegalArgumentException if {@code elements} is empty
063   */
064  public static <E extends Enum<E>> EnumMultiset<E> create(Iterable<E> elements) {
065    Iterator<E> iterator = elements.iterator();
066    checkArgument(iterator.hasNext(), "EnumMultiset constructor passed empty Iterable");
067    EnumMultiset<E> multiset = new EnumMultiset<>(iterator.next().getDeclaringClass());
068    Iterables.addAll(multiset, elements);
069    return multiset;
070  }
071
072  /**
073   * Returns a new {@code EnumMultiset} instance containing the given elements. Unlike {@link
074   * EnumMultiset#create(Iterable)}, this method does not produce an exception on an empty iterable.
075   *
076   * @since 14.0
077   */
078  public static <E extends Enum<E>> EnumMultiset<E> create(Iterable<E> elements, Class<E> type) {
079    EnumMultiset<E> result = create(type);
080    Iterables.addAll(result, elements);
081    return result;
082  }
083
084  private transient Class<E> type;
085  private transient E[] enumConstants;
086  private transient int[] counts;
087  private transient int distinctElements;
088  private transient long size;
089
090  /** Creates an empty {@code EnumMultiset}. */
091  private EnumMultiset(Class<E> type) {
092    this.type = type;
093    checkArgument(type.isEnum());
094    this.enumConstants = type.getEnumConstants();
095    this.counts = new int[enumConstants.length];
096  }
097
098  private boolean isActuallyE(@CheckForNull Object o) {
099    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}