package org.infinispan.commons.util; import java.io.IOException; import java.io.ObjectInput; import java.io.ObjectOutput; import java.util.BitSet; import java.util.Collection; import java.util.PrimitiveIterator; import java.util.Set; import org.infinispan.commons.io.UnsignedNumeric; /** * Represent a set of integers (e.g. segments) as a {@code BitSet}. * * Memory usage depends on the highest element, as in {@link BitSet} and unlike in other collections such as * {@link java.util.HashSet}. * * @author Dan Berindei * @since 9.0 */ public class SmallIntSet implements Set<Integer> { private final BitSet bitSet; public static SmallIntSet of(int i1) { SmallIntSet set = new SmallIntSet(); set.set(i1); return set; } public static SmallIntSet of(int i1, int i2) { SmallIntSet set = new SmallIntSet(); set.set(i1); set.set(i2); return set; } public static SmallIntSet of(int i1, int i2, int i3) { SmallIntSet set = new SmallIntSet(); set.set(i1); set.set(i2); set.set(i3); return set; } public static SmallIntSet of(int... elements) { SmallIntSet set = new SmallIntSet(); for (int i : elements) { set.set(i); } return set; } public SmallIntSet() { bitSet = new BitSet(); } /** * Create a new {@code IntSet} and pre-allocate space for elements {@code 0..initialRange-1}. */ public SmallIntSet(int initialRange) { bitSet = new BitSet(initialRange); } public SmallIntSet(SmallIntSet other) { bitSet = new BitSet(other.bitSet.size()); bitSet.or(other.bitSet); } public SmallIntSet(Set<Integer> set) { bitSet = new BitSet(); set.forEach(bitSet::set); } @Override public int size() { return bitSet.cardinality(); } public int capacity() { return bitSet.size(); } @Override public boolean isEmpty() { return bitSet.isEmpty(); } @Override public boolean contains(Object o) { return (o instanceof Integer) && contains((Integer) o); } public boolean contains(Integer o) { return bitSet.get(o); } /** * Check if the set contains an integer without boxing the parameter. */ public boolean contains(int i) { return bitSet.get(i); } @Override public PrimitiveIterator.OfInt iterator() { return bitSet.stream().iterator(); } @Override public Object[] toArray() { int size = size(); Integer[] dest = new Integer[size]; copyToArray(size, dest); return dest; } @Override public <T> T[] toArray(T[] a) { if (!(a instanceof Integer[])) { throw new IllegalArgumentException("Only Integer arrays are supported"); } int size = size(); Integer[] dest = a.length < size ? new Integer[size] : (Integer[]) a; copyToArray(size, dest); return (T[]) dest; } private void copyToArray(int size, Integer[] dest) { int lastSetBit = -1; for (int i = 0; i < size; i++) { lastSetBit = bitSet.nextSetBit(lastSetBit + 1); dest[i] = lastSetBit; } } @Override public boolean add(Integer i) { return add((int) i); } /** * Add an integer to the set without boxing the parameter. */ public boolean add(int i) { boolean wasSet = bitSet.get(i); if (!wasSet) { bitSet.set(i); return true; } return false; } /** * Add an integer to the set without boxing the parameter or checking if the integer was already present in the set. */ public void set(int i) { bitSet.set(i); } /** * If {@code value} is {@code true}, add the integer to the set, otherwise remove the integer from the set. */ public void set(int i, boolean value) { bitSet.set(i, value); } @Override public boolean remove(Object o) { if (!(o instanceof Integer)) { return false; } return remove((int) o); } /** * Remove an integer from the set without boxing. */ public boolean remove(int i) { boolean wasSet = bitSet.get(i); if (wasSet) { bitSet.clear(i); return true; } return false; } @Override public boolean containsAll(Collection<?> c) { return c.stream().allMatch(this::contains); } @Override public boolean addAll(Collection<? extends Integer> c) { boolean modified = false; for (Integer integer : c) { modified |= add(integer); } return modified; } @Override public boolean removeAll(Collection<?> c) { boolean modified = false; for (Object integer : c) { modified |= remove(integer); } return modified; } @Override public boolean retainAll(Collection<?> c) { boolean modified = false; for (int i = bitSet.nextSetBit(0); i >= 0; i = bitSet.nextSetBit(i + 1)) { if (!c.contains(i)) { bitSet.clear(i); modified = true; } } return modified; } @Override public void clear() { bitSet.clear(); } @Override public boolean equals(Object o) { if (this == o) return true; if (o == null || !(o instanceof Set)) return false; if (o instanceof SmallIntSet) { SmallIntSet integers = (SmallIntSet) o; return bitSet.equals(integers.bitSet); } else { Set set = (Set) o; return size() == set.size() && containsAll(set); } } @Override public int hashCode() { return bitSet.hashCode(); } @Override public String toString() { StringBuilder sb = new StringBuilder("{"); for (int i = bitSet.nextSetBit(0); i >= 0; i = bitSet.nextSetBit(i + 1)) { if (sb.length() > "{".length()) { sb.append(' '); } int runStart = i; while (bitSet.get(i + 1)) { i++; } if (i == runStart) { sb.append(i); } else { sb.append(runStart).append('-').append(i); } } sb.append('}'); return sb.toString(); } public static void writeTo(ObjectOutput output, SmallIntSet set) throws IOException { UnsignedNumeric.writeUnsignedInt(output, set.capacity()); UnsignedNumeric.writeUnsignedInt(output, set.size()); for (int element : set) { UnsignedNumeric.writeUnsignedInt(output, element); } } public static SmallIntSet readFrom(ObjectInput input) throws IOException { int capacity = UnsignedNumeric.readUnsignedInt(input); int size = UnsignedNumeric.readUnsignedInt(input); SmallIntSet set = new SmallIntSet(capacity); for (int i = 0; i < size; i++) { int element = UnsignedNumeric.readUnsignedInt(input); set.set(element); } return set; } }