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