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