com.google.common.collect.EnumMultiset Maven / Gradle / Ivy
/*
* Copyright (C) 2007 The Guava Authors
*
* Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
* in compliance with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software distributed under the License
* is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
* or implied. See the License for the specific language governing permissions and limitations under
* the License.
*/
package com.google.common.collect;
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.collect.CollectPreconditions.checkNonnegative;
import static com.google.common.collect.CollectPreconditions.checkRemove;
import static java.util.Objects.requireNonNull;
import com.google.common.annotations.GwtCompatible;
import com.google.common.annotations.GwtIncompatible;
import com.google.common.annotations.J2ktIncompatible;
import com.google.common.primitives.Ints;
import com.google.errorprone.annotations.CanIgnoreReturnValue;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.Serializable;
import java.util.Arrays;
import java.util.Iterator;
import java.util.NoSuchElementException;
import java.util.function.ObjIntConsumer;
import javax.annotation.CheckForNull;
/**
* Multiset implementation specialized for enum elements, supporting all single-element operations
* in O(1).
*
* See the Guava User Guide article on {@code Multiset}.
*
* @author Jared Levy
* @since 2.0
*/
@GwtCompatible(emulated = true)
@J2ktIncompatible
@ElementTypesAreNonnullByDefault
public final class EnumMultiset> extends AbstractMultiset
implements Serializable {
/** Creates an empty {@code EnumMultiset}. */
public static > EnumMultiset create(Class type) {
return new EnumMultiset<>(type);
}
/**
* Creates a new {@code EnumMultiset} containing the specified elements.
*
* This implementation is highly efficient when {@code elements} is itself a {@link Multiset}.
*
* @param elements the elements that the multiset should contain
* @throws IllegalArgumentException if {@code elements} is empty
*/
public static > EnumMultiset create(Iterable elements) {
Iterator iterator = elements.iterator();
checkArgument(iterator.hasNext(), "EnumMultiset constructor passed empty Iterable");
EnumMultiset multiset = new EnumMultiset<>(iterator.next().getDeclaringClass());
Iterables.addAll(multiset, elements);
return multiset;
}
/**
* Returns a new {@code EnumMultiset} instance containing the given elements. Unlike {@link
* EnumMultiset#create(Iterable)}, this method does not produce an exception on an empty iterable.
*
* @since 14.0
*/
public static > EnumMultiset create(Iterable elements, Class type) {
EnumMultiset result = create(type);
Iterables.addAll(result, elements);
return result;
}
private transient Class type;
private transient E[] enumConstants;
private transient int[] counts;
private transient int distinctElements;
private transient long size;
/** Creates an empty {@code EnumMultiset}. */
private EnumMultiset(Class type) {
this.type = type;
checkArgument(type.isEnum());
this.enumConstants = type.getEnumConstants();
this.counts = new int[enumConstants.length];
}
private boolean isActuallyE(@CheckForNull Object o) {
if (o instanceof Enum) {
Enum e = (Enum) o;
int index = e.ordinal();
return index < enumConstants.length && enumConstants[index] == e;
}
return false;
}
/**
* Returns {@code element} cast to {@code E}, if it actually is a nonnull E. Otherwise, throws
* either a NullPointerException or a ClassCastException as appropriate.
*/
private void checkIsE(Object element) {
checkNotNull(element);
if (!isActuallyE(element)) {
throw new ClassCastException("Expected an " + type + " but got " + element);
}
}
@Override
int distinctElements() {
return distinctElements;
}
@Override
public int size() {
return Ints.saturatedCast(size);
}
@Override
public int count(@CheckForNull Object element) {
// isActuallyE checks for null, but we check explicitly to help nullness checkers.
if (element == null || !isActuallyE(element)) {
return 0;
}
Enum e = (Enum) element;
return counts[e.ordinal()];
}
// Modification Operations
@CanIgnoreReturnValue
@Override
public int add(E element, int occurrences) {
checkIsE(element);
checkNonnegative(occurrences, "occurrences");
if (occurrences == 0) {
return count(element);
}
int index = element.ordinal();
int oldCount = counts[index];
long newCount = (long) oldCount + occurrences;
checkArgument(newCount <= Integer.MAX_VALUE, "too many occurrences: %s", newCount);
counts[index] = (int) newCount;
if (oldCount == 0) {
distinctElements++;
}
size += occurrences;
return oldCount;
}
// Modification Operations
@CanIgnoreReturnValue
@Override
public int remove(@CheckForNull Object element, int occurrences) {
// isActuallyE checks for null, but we check explicitly to help nullness checkers.
if (element == null || !isActuallyE(element)) {
return 0;
}
Enum e = (Enum) element;
checkNonnegative(occurrences, "occurrences");
if (occurrences == 0) {
return count(element);
}
int index = e.ordinal();
int oldCount = counts[index];
if (oldCount == 0) {
return 0;
} else if (oldCount <= occurrences) {
counts[index] = 0;
distinctElements--;
size -= oldCount;
} else {
counts[index] = oldCount - occurrences;
size -= occurrences;
}
return oldCount;
}
// Modification Operations
@CanIgnoreReturnValue
@Override
public int setCount(E element, int count) {
checkIsE(element);
checkNonnegative(count, "count");
int index = element.ordinal();
int oldCount = counts[index];
counts[index] = count;
size += count - oldCount;
if (oldCount == 0 && count > 0) {
distinctElements++;
} else if (oldCount > 0 && count == 0) {
distinctElements--;
}
return oldCount;
}
@Override
public void clear() {
Arrays.fill(counts, 0);
size = 0;
distinctElements = 0;
}
abstract class Itr implements Iterator {
int index = 0;
int toRemove = -1;
abstract T output(int index);
@Override
public boolean hasNext() {
for (; index < enumConstants.length; index++) {
if (counts[index] > 0) {
return true;
}
}
return false;
}
@Override
public T next() {
if (!hasNext()) {
throw new NoSuchElementException();
}
T result = output(index);
toRemove = index;
index++;
return result;
}
@Override
public void remove() {
checkRemove(toRemove >= 0);
if (counts[toRemove] > 0) {
distinctElements--;
size -= counts[toRemove];
counts[toRemove] = 0;
}
toRemove = -1;
}
}
@Override
Iterator elementIterator() {
return new Itr() {
@Override
E output(int index) {
return enumConstants[index];
}
};
}
@Override
Iterator> entryIterator() {
return new Itr>() {
@Override
Entry output(final int index) {
return new Multisets.AbstractEntry() {
@Override
public E getElement() {
return enumConstants[index];
}
@Override
public int getCount() {
return counts[index];
}
};
}
};
}
@Override
public void forEachEntry(ObjIntConsumer action) {
checkNotNull(action);
for (int i = 0; i < enumConstants.length; i++) {
if (counts[i] > 0) {
action.accept(enumConstants[i], counts[i]);
}
}
}
@Override
public Iterator iterator() {
return Multisets.iteratorImpl(this);
}
@GwtIncompatible // java.io.ObjectOutputStream
private void writeObject(ObjectOutputStream stream) throws IOException {
stream.defaultWriteObject();
stream.writeObject(type);
Serialization.writeMultiset(this, stream);
}
/**
* @serialData the {@code Class} for the enum type, the number of distinct elements, the first
* element, its count, the second element, its count, and so on
*/
@GwtIncompatible // java.io.ObjectInputStream
private void readObject(ObjectInputStream stream) throws IOException, ClassNotFoundException {
stream.defaultReadObject();
@SuppressWarnings("unchecked") // reading data stored by writeObject
Class localType = (Class) requireNonNull(stream.readObject());
type = localType;
enumConstants = type.getEnumConstants();
counts = new int[enumConstants.length];
Serialization.populateMultiset(this, stream);
}
@GwtIncompatible // Not needed in emulated source
private static final long serialVersionUID = 0;
}