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  void checkIsE(@Nullable 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(@Nullable Object element) {
129    if (!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(@Nullable Object element, int occurrences) {
161    if (!isActuallyE(element)) {
162      return 0;
163    }
164    Enum<?> e = (Enum<?>) element;
165    checkNonnegative(occurrences, "occurrences");
166    if (occurrences == 0) {
167      return count(element);
168    }
169    int index = e.ordinal();
170    int oldCount = counts[index];
171    if (oldCount == 0) {
172      return 0;
173    } else if (oldCount <= occurrences) {
174      counts[index] = 0;
175      distinctElements--;
176      size -= oldCount;
177    } else {
178      counts[index] = oldCount - occurrences;
179      size -= occurrences;
180    }
181    return oldCount;
182  }
183
184  // Modification Operations
185  @CanIgnoreReturnValue
186  @Override
187  public int setCount(E element, int count) {
188    checkIsE(element);
189    checkNonnegative(count, "count");
190    int index = element.ordinal();
191    int oldCount = counts[index];
192    counts[index] = count;
193    size += count - oldCount;
194    if (oldCount == 0 && count > 0) {
195      distinctElements++;
196    } else if (oldCount > 0 && count == 0) {
197      distinctElements--;
198    }
199    return oldCount;
200  }
201
202  @Override
203  public void clear() {
204    Arrays.fill(counts, 0);
205    size = 0;
206    distinctElements = 0;
207  }
208
209  abstract class Itr<T> implements Iterator<T> {
210    int index = 0;
211    int toRemove = -1;
212
213    abstract T output(int index);
214
215    @Override
216    public boolean hasNext() {
217      for (; index < enumConstants.length; index++) {
218        if (counts[index] > 0) {
219          return true;
220        }
221      }
222      return false;
223    }
224
225    @Override
226    public T next() {
227      if (!hasNext()) {
228        throw new NoSuchElementException();
229      }
230      T result = output(index);
231      toRemove = index;
232      index++;
233      return result;
234    }
235
236    @Override
237    public void remove() {
238      checkRemove(toRemove >= 0);
239      if (counts[toRemove] > 0) {
240        distinctElements--;
241        size -= counts[toRemove];
242        counts[toRemove] = 0;
243      }
244      toRemove = -1;
245    }
246  }
247
248  @Override
249  Iterator<E> elementIterator() {
250    return new Itr<E>() {
251      @Override
252      E output(int index) {
253        return enumConstants[index];
254      }
255    };
256  }
257
258  @Override
259  Iterator<Entry<E>> entryIterator() {
260    return new Itr<Entry<E>>() {
261      @Override
262      Entry<E> output(final int index) {
263        return new Multisets.AbstractEntry<E>() {
264          @Override
265          public E getElement() {
266            return enumConstants[index];
267          }
268
269          @Override
270          public int getCount() {
271            return counts[index];
272          }
273        };
274      }
275    };
276  }
277
278  @Override
279  public void forEachEntry(ObjIntConsumer<? super E> action) {
280    checkNotNull(action);
281    for (int i = 0; i < enumConstants.length; i++) {
282      if (counts[i] > 0) {
283        action.accept(enumConstants[i], counts[i]);
284      }
285    }
286  }
287
288  @Override
289  public Iterator<E> iterator() {
290    return Multisets.iteratorImpl(this);
291  }
292
293  @GwtIncompatible // java.io.ObjectOutputStream
294  private void writeObject(ObjectOutputStream stream) throws IOException {
295    stream.defaultWriteObject();
296    stream.writeObject(type);
297    Serialization.writeMultiset(this, stream);
298  }
299
300  /**
301   * @serialData the {@code Class<E>} for the enum type, the number of distinct elements, the first
302   *     element, its count, the second element, its count, and so on
303   */
304  @GwtIncompatible // java.io.ObjectInputStream
305  private void readObject(ObjectInputStream stream) throws IOException, ClassNotFoundException {
306    stream.defaultReadObject();
307    @SuppressWarnings("unchecked") // reading data stored by writeObject
308    Class<E> localType = (Class<E>) stream.readObject();
309    type = localType;
310    enumConstants = type.getEnumConstants();
311    counts = new int[enumConstants.length];
312    Serialization.populateMultiset(this, stream);
313  }
314
315  @GwtIncompatible // Not needed in emulated source
316  private static final long serialVersionUID = 0;
317}