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<E>(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}