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 org.checkerframework.checker.nullness.qual.Nullable;
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)
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(@Nullable 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  @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  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 void forEachEntry(ObjIntConsumer<? super E> action) {
281    checkNotNull(action);
282    for (int i = 0; i < enumConstants.length; i++) {
283      if (counts[i] > 0) {
284        action.accept(enumConstants[i], counts[i]);
285      }
286    }
287  }
288
289  @Override
290  public Iterator<E> iterator() {
291    return Multisets.iteratorImpl(this);
292  }
293
294  @GwtIncompatible // java.io.ObjectOutputStream
295  private void writeObject(ObjectOutputStream stream) throws IOException {
296    stream.defaultWriteObject();
297    stream.writeObject(type);
298    Serialization.writeMultiset(this, stream);
299  }
300
301  /**
302   * @serialData the {@code Class<E>} for the enum type, the number of distinct elements, the first
303   *     element, its count, the second element, its count, and so on
304   */
305  @GwtIncompatible // java.io.ObjectInputStream
306  private void readObject(ObjectInputStream stream) throws IOException, ClassNotFoundException {
307    stream.defaultReadObject();
308    @SuppressWarnings("unchecked") // reading data stored by writeObject
309    Class<E> localType = (Class<E>) stream.readObject();
310    type = localType;
311    enumConstants = type.getEnumConstants();
312    counts = new int[enumConstants.length];
313    Serialization.populateMultiset(this, stream);
314  }
315
316  @GwtIncompatible // Not needed in emulated source
317  private static final long serialVersionUID = 0;
318}