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