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;
021import static java.util.Objects.requireNonNull;
022
023import com.google.common.annotations.GwtCompatible;
024import com.google.common.annotations.GwtIncompatible;
025import com.google.common.annotations.J2ktIncompatible;
026import com.google.common.primitives.Ints;
027import com.google.errorprone.annotations.CanIgnoreReturnValue;
028import java.io.IOException;
029import java.io.ObjectInputStream;
030import java.io.ObjectOutputStream;
031import java.io.Serializable;
032import java.util.Arrays;
033import java.util.Iterator;
034import java.util.NoSuchElementException;
035import java.util.function.ObjIntConsumer;
036import javax.annotation.CheckForNull;
037
038/**
039 * Multiset implementation specialized for enum elements, supporting all single-element operations
040 * in O(1).
041 *
042 * <p>See the Guava User Guide article on <a href=
043 * "https://github.com/google/guava/wiki/NewCollectionTypesExplained#multiset">{@code Multiset}</a>.
044 *
045 * @author Jared Levy
046 * @since 2.0
047 */
048@GwtCompatible(emulated = true)
049@J2ktIncompatible
050@ElementTypesAreNonnullByDefault
051public final class EnumMultiset<E extends Enum<E>> extends AbstractMultiset<E>
052    implements Serializable {
053  /** Creates an empty {@code EnumMultiset}. */
054  public static <E extends Enum<E>> EnumMultiset<E> create(Class<E> type) {
055    return new EnumMultiset<>(type);
056  }
057
058  /**
059   * Creates a new {@code EnumMultiset} containing the specified elements.
060   *
061   * <p>This implementation is highly efficient when {@code elements} is itself a {@link Multiset}.
062   *
063   * @param elements the elements that the multiset should contain
064   * @throws IllegalArgumentException if {@code elements} is empty
065   */
066  public static <E extends Enum<E>> EnumMultiset<E> create(Iterable<E> elements) {
067    Iterator<E> iterator = elements.iterator();
068    checkArgument(iterator.hasNext(), "EnumMultiset constructor passed empty Iterable");
069    EnumMultiset<E> multiset = new EnumMultiset<>(iterator.next().getDeclaringClass());
070    Iterables.addAll(multiset, elements);
071    return multiset;
072  }
073
074  /**
075   * Returns a new {@code EnumMultiset} instance containing the given elements. Unlike {@link
076   * EnumMultiset#create(Iterable)}, this method does not produce an exception on an empty iterable.
077   *
078   * @since 14.0
079   */
080  public static <E extends Enum<E>> EnumMultiset<E> create(Iterable<E> elements, Class<E> type) {
081    EnumMultiset<E> result = create(type);
082    Iterables.addAll(result, elements);
083    return result;
084  }
085
086  private transient Class<E> type;
087  private transient E[] enumConstants;
088  private transient int[] counts;
089  private transient int distinctElements;
090  private transient long size;
091
092  /** Creates an empty {@code EnumMultiset}. */
093  private EnumMultiset(Class<E> type) {
094    this.type = type;
095    checkArgument(type.isEnum());
096    this.enumConstants = type.getEnumConstants();
097    this.counts = new int[enumConstants.length];
098  }
099
100  private boolean isActuallyE(@CheckForNull Object o) {
101    if (o instanceof Enum) {
102      Enum<?> e = (Enum<?>) o;
103      int index = e.ordinal();
104      return index < enumConstants.length && enumConstants[index] == e;
105    }
106    return false;
107  }
108
109  /**
110   * Returns {@code element} cast to {@code E}, if it actually is a nonnull E. Otherwise, throws
111   * either a NullPointerException or a ClassCastException as appropriate.
112   */
113  private void checkIsE(Object element) {
114    checkNotNull(element);
115    if (!isActuallyE(element)) {
116      throw new ClassCastException("Expected an " + type + " but got " + element);
117    }
118  }
119
120  @Override
121  int distinctElements() {
122    return distinctElements;
123  }
124
125  @Override
126  public int size() {
127    return Ints.saturatedCast(size);
128  }
129
130  @Override
131  public int count(@CheckForNull Object element) {
132    // isActuallyE checks for null, but we check explicitly to help nullness checkers.
133    if (element == null || !isActuallyE(element)) {
134      return 0;
135    }
136    Enum<?> e = (Enum<?>) element;
137    return counts[e.ordinal()];
138  }
139
140  // Modification Operations
141  @CanIgnoreReturnValue
142  @Override
143  public int add(E element, int occurrences) {
144    checkIsE(element);
145    checkNonnegative(occurrences, "occurrences");
146    if (occurrences == 0) {
147      return count(element);
148    }
149    int index = element.ordinal();
150    int oldCount = counts[index];
151    long newCount = (long) oldCount + occurrences;
152    checkArgument(newCount <= Integer.MAX_VALUE, "too many occurrences: %s", newCount);
153    counts[index] = (int) newCount;
154    if (oldCount == 0) {
155      distinctElements++;
156    }
157    size += occurrences;
158    return oldCount;
159  }
160
161  // Modification Operations
162  @CanIgnoreReturnValue
163  @Override
164  public int remove(@CheckForNull Object element, int occurrences) {
165    // isActuallyE checks for null, but we check explicitly to help nullness checkers.
166    if (element == null || !isActuallyE(element)) {
167      return 0;
168    }
169    Enum<?> e = (Enum<?>) element;
170    checkNonnegative(occurrences, "occurrences");
171    if (occurrences == 0) {
172      return count(element);
173    }
174    int index = e.ordinal();
175    int oldCount = counts[index];
176    if (oldCount == 0) {
177      return 0;
178    } else if (oldCount <= occurrences) {
179      counts[index] = 0;
180      distinctElements--;
181      size -= oldCount;
182    } else {
183      counts[index] = oldCount - occurrences;
184      size -= occurrences;
185    }
186    return oldCount;
187  }
188
189  // Modification Operations
190  @CanIgnoreReturnValue
191  @Override
192  public int setCount(E element, int count) {
193    checkIsE(element);
194    checkNonnegative(count, "count");
195    int index = element.ordinal();
196    int oldCount = counts[index];
197    counts[index] = count;
198    size += count - oldCount;
199    if (oldCount == 0 && count > 0) {
200      distinctElements++;
201    } else if (oldCount > 0 && count == 0) {
202      distinctElements--;
203    }
204    return oldCount;
205  }
206
207  @Override
208  public void clear() {
209    Arrays.fill(counts, 0);
210    size = 0;
211    distinctElements = 0;
212  }
213
214  abstract class Itr<T> implements Iterator<T> {
215    int index = 0;
216    int toRemove = -1;
217
218    abstract T output(int index);
219
220    @Override
221    public boolean hasNext() {
222      for (; index < enumConstants.length; index++) {
223        if (counts[index] > 0) {
224          return true;
225        }
226      }
227      return false;
228    }
229
230    @Override
231    public T next() {
232      if (!hasNext()) {
233        throw new NoSuchElementException();
234      }
235      T result = output(index);
236      toRemove = index;
237      index++;
238      return result;
239    }
240
241    @Override
242    public void remove() {
243      checkRemove(toRemove >= 0);
244      if (counts[toRemove] > 0) {
245        distinctElements--;
246        size -= counts[toRemove];
247        counts[toRemove] = 0;
248      }
249      toRemove = -1;
250    }
251  }
252
253  @Override
254  Iterator<E> elementIterator() {
255    return new Itr<E>() {
256      @Override
257      E output(int index) {
258        return enumConstants[index];
259      }
260    };
261  }
262
263  @Override
264  Iterator<Entry<E>> entryIterator() {
265    return new Itr<Entry<E>>() {
266      @Override
267      Entry<E> output(final int index) {
268        return new Multisets.AbstractEntry<E>() {
269          @Override
270          public E getElement() {
271            return enumConstants[index];
272          }
273
274          @Override
275          public int getCount() {
276            return counts[index];
277          }
278        };
279      }
280    };
281  }
282
283  @Override
284  public void forEachEntry(ObjIntConsumer<? super E> action) {
285    checkNotNull(action);
286    for (int i = 0; i < enumConstants.length; i++) {
287      if (counts[i] > 0) {
288        action.accept(enumConstants[i], counts[i]);
289      }
290    }
291  }
292
293  @Override
294  public Iterator<E> iterator() {
295    return Multisets.iteratorImpl(this);
296  }
297
298  @GwtIncompatible // java.io.ObjectOutputStream
299  private void writeObject(ObjectOutputStream stream) throws IOException {
300    stream.defaultWriteObject();
301    stream.writeObject(type);
302    Serialization.writeMultiset(this, stream);
303  }
304
305  /**
306   * @serialData the {@code Class<E>} for the enum type, the number of distinct elements, the first
307   *     element, its count, the second element, its count, and so on
308   */
309  @GwtIncompatible // java.io.ObjectInputStream
310  private void readObject(ObjectInputStream stream) throws IOException, ClassNotFoundException {
311    stream.defaultReadObject();
312    @SuppressWarnings("unchecked") // reading data stored by writeObject
313    Class<E> localType = (Class<E>) requireNonNull(stream.readObject());
314    type = localType;
315    enumConstants = type.getEnumConstants();
316    counts = new int[enumConstants.length];
317    Serialization.populateMultiset(this, stream);
318  }
319
320  @GwtIncompatible // Not needed in emulated source
321  private static final long serialVersionUID = 0;
322}