1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17 package org.apache.commons.collections4.multiset;
18
19 import java.io.IOException;
20 import java.io.ObjectInputStream;
21 import java.io.ObjectOutputStream;
22 import java.lang.reflect.Array;
23 import java.util.ConcurrentModificationException;
24 import java.util.Iterator;
25 import java.util.Map;
26
27 import org.apache.commons.collections4.MultiSet;
28 import org.apache.commons.collections4.iterators.AbstractIteratorDecorator;
29
30
31
32
33
34
35
36
37
38
39
40
41
42 public abstract class AbstractMapMultiSet<E> extends AbstractMultiSet<E> {
43
44
45
46
47 protected static class EntrySetIterator<E> implements Iterator<Entry<E>> {
48
49
50 protected final AbstractMapMultiSet<E> parent;
51
52 protected final Iterator<Map.Entry<E, MutableInteger>> decorated;
53
54
55 protected Entry<E> last;
56
57
58 protected boolean canRemove;
59
60
61
62
63
64
65 protected EntrySetIterator(final Iterator<Map.Entry<E, MutableInteger>> iterator,
66 final AbstractMapMultiSet<E> parent) {
67 this.decorated = iterator;
68 this.parent = parent;
69 }
70
71 @Override
72 public boolean hasNext() {
73 return decorated.hasNext();
74 }
75
76 @Override
77 public Entry<E> next() {
78 last = new MultiSetEntry<>(decorated.next());
79 canRemove = true;
80 return last;
81 }
82
83 @Override
84 public void remove() {
85 if (!canRemove) {
86 throw new IllegalStateException("Iterator remove() can only be called once after next()");
87 }
88 decorated.remove();
89 last = null;
90 canRemove = false;
91 }
92 }
93
94
95
96 private static final class MapBasedMultiSetIterator<E> implements Iterator<E> {
97 private final AbstractMapMultiSet<E> parent;
98 private final Iterator<Map.Entry<E, MutableInteger>> entryIterator;
99 private Map.Entry<E, MutableInteger> current;
100 private int itemCount;
101 private final int mods;
102 private boolean canRemove;
103
104
105
106
107
108
109 MapBasedMultiSetIterator(final AbstractMapMultiSet<E> parent) {
110 this.parent = parent;
111 this.entryIterator = parent.map.entrySet().iterator();
112 this.current = null;
113 this.mods = parent.modCount;
114 this.canRemove = false;
115 }
116
117
118 @Override
119 public boolean hasNext() {
120 return itemCount > 0 || entryIterator.hasNext();
121 }
122
123
124 @Override
125 public E next() {
126 if (parent.modCount != mods) {
127 throw new ConcurrentModificationException();
128 }
129 if (itemCount == 0) {
130 current = entryIterator.next();
131 itemCount = current.getValue().value;
132 }
133 canRemove = true;
134 itemCount--;
135 return current.getKey();
136 }
137
138
139 @Override
140 public void remove() {
141 if (parent.modCount != mods) {
142 throw new ConcurrentModificationException();
143 }
144 if (!canRemove) {
145 throw new IllegalStateException();
146 }
147 final MutableInteger mut = current.getValue();
148 if (mut.value > 1) {
149 mut.value--;
150 } else {
151 entryIterator.remove();
152 }
153 parent.size--;
154 canRemove = false;
155 }
156 }
157
158
159
160 protected static class MultiSetEntry<E> extends AbstractEntry<E> {
161
162 protected final Map.Entry<E, MutableInteger> parentEntry;
163
164
165
166
167
168 protected MultiSetEntry(final Map.Entry<E, MutableInteger> parentEntry) {
169 this.parentEntry = parentEntry;
170 }
171
172 @Override
173 public int getCount() {
174 return parentEntry.getValue().value;
175 }
176
177 @Override
178 public E getElement() {
179 return parentEntry.getKey();
180 }
181 }
182
183
184
185
186 protected static class MutableInteger {
187
188 protected int value;
189
190
191
192
193
194 MutableInteger(final int value) {
195 this.value = value;
196 }
197
198 @Override
199 public boolean equals(final Object obj) {
200 if (!(obj instanceof MutableInteger)) {
201 return false;
202 }
203 return ((MutableInteger) obj).value == value;
204 }
205
206 @Override
207 public int hashCode() {
208 return value;
209 }
210 }
211
212
213
214
215 protected static class UniqueSetIterator<E> extends AbstractIteratorDecorator<E> {
216
217
218 protected final AbstractMapMultiSet<E> parent;
219
220
221 protected E lastElement;
222
223
224 protected boolean canRemove;
225
226
227
228
229
230
231 protected UniqueSetIterator(final Iterator<E> iterator, final AbstractMapMultiSet<E> parent) {
232 super(iterator);
233 this.parent = parent;
234 }
235
236 @Override
237 public E next() {
238 lastElement = super.next();
239 canRemove = true;
240 return lastElement;
241 }
242
243 @Override
244 public void remove() {
245 if (!canRemove) {
246 throw new IllegalStateException("Iterator remove() can only be called once after next()");
247 }
248 final int count = parent.getCount(lastElement);
249 super.remove();
250 parent.remove(lastElement, count);
251 lastElement = null;
252 canRemove = false;
253 }
254 }
255
256
257 private transient Map<E, MutableInteger> map;
258
259
260 private transient int size;
261
262
263 private transient int modCount;
264
265
266
267
268 protected AbstractMapMultiSet() {
269 }
270
271
272
273
274
275
276
277 protected AbstractMapMultiSet(final Map<E, MutableInteger> map) {
278 this.map = map;
279 }
280
281 @Override
282 public int add(final E object, final int occurrences) {
283 if (occurrences < 0) {
284 throw new IllegalArgumentException("Occurrences must not be negative.");
285 }
286
287 final MutableInteger mut = map.get(object);
288 final int oldCount = mut != null ? mut.value : 0;
289
290 if (occurrences > 0) {
291 modCount++;
292 size += occurrences;
293 if (mut == null) {
294 map.put(object, new MutableInteger(occurrences));
295 } else {
296 mut.value += occurrences;
297 }
298 }
299 return oldCount;
300 }
301
302
303
304
305 @Override
306 public void clear() {
307 modCount++;
308 map.clear();
309 size = 0;
310 }
311
312
313
314
315
316
317
318
319 @Override
320 public boolean contains(final Object object) {
321 return map.containsKey(object);
322 }
323
324 @Override
325 protected Iterator<Entry<E>> createEntrySetIterator() {
326 return new EntrySetIterator<>(map.entrySet().iterator(), this);
327 }
328
329 @Override
330 protected Iterator<E> createUniqueSetIterator() {
331 return new UniqueSetIterator<>(getMap().keySet().iterator(), this);
332 }
333
334
335
336
337
338
339
340
341 @Override
342 protected void doReadObject(final ObjectInputStream in)
343 throws IOException, ClassNotFoundException {
344 final int entrySize = in.readInt();
345 for (int i = 0; i < entrySize; i++) {
346 @SuppressWarnings("unchecked")
347 final E obj = (E) in.readObject();
348 final int count = in.readInt();
349 map.put(obj, new MutableInteger(count));
350 size += count;
351 }
352 }
353
354
355
356
357
358
359 @Override
360 protected void doWriteObject(final ObjectOutputStream out) throws IOException {
361 out.writeInt(map.size());
362 for (final Map.Entry<E, MutableInteger> entry : map.entrySet()) {
363 out.writeObject(entry.getKey());
364 out.writeInt(entry.getValue().value);
365 }
366 }
367
368 @Override
369 public boolean equals(final Object object) {
370 if (object == this) {
371 return true;
372 }
373 if (!(object instanceof MultiSet)) {
374 return false;
375 }
376 final MultiSet<?> other = (MultiSet<?>) object;
377 if (other.size() != size()) {
378 return false;
379 }
380 for (final E element : map.keySet()) {
381 if (other.getCount(element) != getCount(element)) {
382 return false;
383 }
384 }
385 return true;
386 }
387
388
389
390
391
392
393
394
395 @Override
396 public int getCount(final Object object) {
397 final MutableInteger count = map.get(object);
398 if (count != null) {
399 return count.value;
400 }
401 return 0;
402 }
403
404
405
406
407
408
409
410 protected Map<E, MutableInteger> getMap() {
411 return map;
412 }
413
414 @Override
415 public int hashCode() {
416 int total = 0;
417 for (final Map.Entry<E, MutableInteger> entry : map.entrySet()) {
418 final E element = entry.getKey();
419 final MutableInteger count = entry.getValue();
420 total += (element == null ? 0 : element.hashCode()) ^ count.value;
421 }
422 return total;
423 }
424
425
426
427
428
429
430 @Override
431 public boolean isEmpty() {
432 return map.isEmpty();
433 }
434
435
436
437
438
439
440
441 @Override
442 public Iterator<E> iterator() {
443 return new MapBasedMultiSetIterator<>(this);
444 }
445
446 @Override
447 public int remove(final Object object, final int occurrences) {
448 if (occurrences < 0) {
449 throw new IllegalArgumentException("Occurrences must not be negative.");
450 }
451
452 final MutableInteger mut = map.get(object);
453 if (mut == null) {
454 return 0;
455 }
456 final int oldCount = mut.value;
457 if (occurrences > 0) {
458 modCount++;
459 if (occurrences < mut.value) {
460 mut.value -= occurrences;
461 size -= occurrences;
462 } else {
463 map.remove(object);
464 size -= mut.value;
465 mut.value = 0;
466 }
467 }
468 return oldCount;
469 }
470
471
472
473
474
475
476
477
478 protected void setMap(final Map<E, MutableInteger> map) {
479 this.map = map;
480 }
481
482
483
484
485
486
487 @Override
488 public int size() {
489 return size;
490 }
491
492
493
494
495
496
497 @Override
498 public Object[] toArray() {
499 final Object[] result = new Object[size()];
500 int i = 0;
501 for (final Map.Entry<E, MutableInteger> entry : map.entrySet()) {
502 final E current = entry.getKey();
503 final MutableInteger count = entry.getValue();
504 for (int index = count.value; index > 0; index--) {
505 result[i++] = current;
506 }
507 }
508 return result;
509 }
510
511
512
513
514
515
516
517
518
519
520
521
522
523 @Override
524 public <T> T[] toArray(T[] array) {
525 final int size = size();
526 if (array.length < size) {
527 @SuppressWarnings("unchecked")
528 final T[] unchecked = (T[]) Array.newInstance(array.getClass().getComponentType(), size);
529 array = unchecked;
530 }
531
532 int i = 0;
533 for (final Map.Entry<E, MutableInteger> entry : map.entrySet()) {
534 final E current = entry.getKey();
535 final MutableInteger count = entry.getValue();
536 for (int index = count.value; index > 0; index--) {
537
538 @SuppressWarnings("unchecked")
539 final T unchecked = (T) current;
540 array[i++] = unchecked;
541 }
542 }
543 while (i < array.length) {
544 array[i++] = null;
545 }
546 return array;
547 }
548
549 @Override
550 protected int uniqueElements() {
551 return map.size();
552 }
553 }