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