LinearProbingHashMap.java

import java.util.AbstractCollection;
import java.util.AbstractSet;
import java.util.Arrays;
import java.util.Collection;
import java.util.Iterator;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.Set;

public class LinearProbingHashMap<K, V> implements Map<K, V> {

   public interface EEntry<K, V> {
      K key();
      V value();
      void setValue(V value);
   }


   public static final int DEFAULT_CAPACITY = 8;

   public static final float DEFAULT_LOAD_FACTOR = 0.5f;

   private static final Object NULL_VALUE = new Object();

   private int maxSize;

   private final float loadFactor;

   private K[] keys;
   private V[] values;
   private int size;
   private int mask;

   private final Set<K> keySet = new KeySet();
   private final Set<Entry<K, V>> entrySet = new EntrySet();
   private final Iterable<EEntry<K, V>> entries = new Iterable<EEntry<K, V>>() {
      @Override
      public Iterator<EEntry<K, V>> iterator() {
         return new EIterator();
      }
   };

   public LinearProbingHashMap() {
      this(DEFAULT_CAPACITY, DEFAULT_LOAD_FACTOR);
   }

   public LinearProbingHashMap(int initialCapacity) {
      this(initialCapacity, DEFAULT_LOAD_FACTOR);
   }

   public LinearProbingHashMap(int initialCapacity, float loadFactor) {
      if (loadFactor <= 0.0f || loadFactor > 1.0f) {
         // Cannot exceed 1 because we can never store more than capacity elements;
         // using a bigger loadFactor would trigger rehashing before the desired load is reached.
         throw new IllegalArgumentException("loadFactor must be > 0 and <= 1");
      }

      this.loadFactor = loadFactor;

      // Adjust the initial capacity if necessary.
      int capacity;
      if (initialCapacity <= 0) {
         capacity = 1;
      } else if (initialCapacity >= 0x40000000) {
         capacity = 0x40000000;
      } else {
         capacity = 1 << (32 - Integer.numberOfLeadingZeros(initialCapacity - 1));
      }
      mask = capacity - 1;

      // Allocate the arrays.
      keys = (K[]) new Object[capacity];
      values = (V[]) new Object[capacity];

      // Initialize the maximum size value.
      maxSize = calcMaxSize(capacity);
   }

   private static <T> T toExternal(T value) {
      assert value != null : "null is not a legitimate internal value. Concurrent Modification?";
      return value == NULL_VALUE ? null : value;
   }

   @SuppressWarnings("unchecked")
   private static <T> T toInternal(T value) {
      return value == null ? (T) NULL_VALUE : value;
   }

   @Override
   public V get(Object key) {
      int index = indexOf(key);
      return index == -1 ? null : toExternal(values[index]);
   }

   @Override
   public V put(K key, V value) {
      int startIndex = hashIndex(key);
      int index = startIndex;

      for (; ; ) {
         if (values[index] == null) {
            // Found empty slot, use it.
            keys[index] = key;
            values[index] = toInternal(value);
            growSize();
            return null;
         }
         if (keys[index] == key) {
            // Found existing entry with this key, just replace the value.
            V previousValue = values[index];
            values[index] = toInternal(value);
            return toExternal(previousValue);
         }

         // Conflict, keep probing ...
         if ((index = probeNext(index)) == startIndex) {
            // Can only happen if the map was full at MAX_ARRAY_SIZE and couldn't grow.
            throw new IllegalStateException("Unable to insert");
         }
      }
   }

   @Override
   public void putAll(Map<? extends K, ? extends V> sourceMap) {
      if (sourceMap instanceof LinearProbingHashMap) {
         // Optimization - iterate through the arrays.
         @SuppressWarnings("unchecked")
         LinearProbingHashMap<K, V> source = (LinearProbingHashMap<K, V>) sourceMap;
         for (int i = 0; i < source.values.length; ++i) {
            V sourceValue = source.values[i];
            if (sourceValue != null) {
               put(source.keys[i], sourceValue);
            }
         }
         return;
      }

      // Otherwise, just add each entry.
      for (Entry<? extends K, ? extends V> entry : sourceMap.entrySet()) {
         put(entry.getKey(), entry.getValue());
      }
   }

   @Override
   public V remove(Object key) {
      int index = indexOf(key);
      if (index == -1) {
         return null;
      }

      V prev = values[index];
      removeAt(index);
      return toExternal(prev);
   }

   @Override
   public int size() {
      return size;
   }

   @Override
   public boolean isEmpty() {
      return size == 0;
   }

   @Override
   public void clear() {
      Arrays.fill(keys, null);
      Arrays.fill(values, null);
      size = 0;
   }

   @Override
   public boolean containsKey(Object key) {
      return indexOf(key) >= 0;
   }

   @Override
   public boolean containsValue(Object value) {
      @SuppressWarnings("unchecked")
      V v1 = toInternal((V) value);
      for (V v2 : values) {
         // The map supports null values; this will be matched as NULL_VALUE.equals(NULL_VALUE).
         if (v2 != null && v2.equals(v1)) {
            return true;
         }
      }
      return false;
   }

   public Iterable<EEntry<K, V>> entries() {
      return entries;
   }

   @Override
   public Collection<V> values() {
      return new AbstractCollection<V>() {
         @Override
         public Iterator<V> iterator() {
            return new Iterator<V>() {
               final EIterator iter = new EIterator();

               @Override
               public boolean hasNext() {
                  return iter.hasNext();
               }

               @Override
               public V next() {
                  return iter.next().value();
               }

               @Override
               public void remove() {
                  throw new UnsupportedOperationException();
               }
            };
         }

         @Override
         public int size() {
            return size;
         }
      };
   }

   @Override
   public int hashCode() {
      // Hashcode is based on all non-zero, valid keys. We have to scan the whole keys
      // array, which may have different lengths for two maps of same size(), so the
      // capacity cannot be used as input for hashing but the size can.
      int hash = size;
      for (K key : keys) {
         // 0 can be a valid key or unused slot, but won't impact the hashcode in either case.
         // This way we can use a cheap loop without conditionals, or hard-to-unroll operations,
         // or the devastatingly bad memory locality of visiting value objects.
         // Also, it's important to use a hash function that does not depend on the ordering
         // of terms, only their values; since the map is an unordered collection and
         // entries can end up in different positions in different maps that have the same
         // elements, but with different history of puts/removes, due to conflicts.
         hash ^= hashIndex(key);
      }
      return hash;
   }

   @Override
   public boolean equals(Object obj) {
      if (this == obj) {
         return true;
      }
      if (!(obj instanceof LinearProbingHashMap)) {
         return false;
      }
      @SuppressWarnings("rawtypes")
      LinearProbingHashMap other = (LinearProbingHashMap) obj;
      if (size != other.size()) {
         return false;
      }
      for (int i = 0; i < values.length; ++i) {
         V value = values[i];
         if (value != null) {
            K key = keys[i];
            Object otherValue = other.get(key);
            if (value == NULL_VALUE) {
               if (otherValue != null) {
                  return false;
               }
            } else if (!value.equals(otherValue)) {
               return false;
            }
         }
      }
      return true;
   }


   @Override
   public Set<K> keySet() {
      return keySet;
   }

   @Override
   public Set<Entry<K, V>> entrySet() {
      return entrySet;
   }

   /**
    * Locates the index for the given key. This method probes using double hashing.
    *
    * @param key the key for an entry in the map.
    * @return the index where the key was found, or {@code -1} if no entry is found for that key.
    */
   private int indexOf(Object key) {
      int startIndex = hashIndex(key);
      int index = startIndex;

      for (; ; ) {
         if (values[index] == null) {
            // It's available, so no chance that this value exists anywhere in the map.
            return -1;
         }
         if (key == keys[index]) {
            return index;
         }

         // Conflict, keep probing ...
         if ((index = probeNext(index)) == startIndex) {
            return -1;
         }
      }
   }

   /**
    * Returns the hashed index for the given key.
    */
   private int hashIndex(Object key) {
      int h;
      return ((key == null) ? 0 : (h = key.hashCode()) ^ (h >>> 16)) & mask;
   }

   /**
    * Get the next sequential index after {@code index} and wraps if necessary.
    */
   private int probeNext(int index) {
      // The array lengths are always a power of two, so we can use a bitmask to stay inside the array bounds.
      return (index + 1) & mask;
   }

   /**
    * Grows the map size after an insertion. If necessary, performs a rehash of the map.
    */
   private void growSize() {
      size++;

      if (size > maxSize) {
         if (keys.length == Integer.MAX_VALUE) {
            throw new IllegalStateException("Max capacity reached at size=" + size);
         }

         // Double the capacity.
         rehash(keys.length << 1);
      }
   }

   /**
    * Removes entry at the given index position. Also performs opportunistic, incremental rehashing
    * if necessary to not break conflict chains.
    *
    * @param index the index position of the element to remove.
    * @return {@code true} if the next item was moved back. {@code false} otherwise.
    */
   private boolean removeAt(final int index) {
      --size;
      // Clearing the key is not strictly necessary (for GC like in a regular collection),
      // but recommended for security. The memory location is still fresh in the cache anyway.
      keys[index] = null;
      values[index] = null;

      // In the interval from index to the next available entry, the arrays may have entries
      // that are displaced from their base position due to prior conflicts. Iterate these
      // entries and move them back if possible, optimizing future lookups.
      // Knuth Section 6.4 Algorithm R, also used by the JDK's IdentityLinearProbingHashMap.

      int nextFree = index;
      int i = probeNext(index);
      for (V value = values[i]; value != null; value = values[i = probeNext(i)]) {
         K key = keys[i];
         int bucket = hashIndex(key);
         if (i < bucket && (bucket <= nextFree || nextFree <= i) ||
                 bucket <= nextFree && nextFree <= i) {
            // Move the displaced entry "back" to the first available position.
            keys[nextFree] = key;
            values[nextFree] = value;
            // Put the first entry after the displaced entry
            keys[i] = null;
            values[i] = null;
            nextFree = i;
         }
      }
      return nextFree != index;
   }

   /**
    * Calculates the maximum size allowed before rehashing.
    */
   private int calcMaxSize(int capacity) {
      // Clip the upper bound so that there will always be at least one available slot.
      int upperBound = capacity - 1;
      return Math.min(upperBound, (int) (capacity * loadFactor));
   }

   /**
    * Rehashes the map for the given capacity.
    *
    * @param newCapacity the new capacity for the map.
    */
   private void rehash(int newCapacity) {
      K[] oldKeys = keys;
      V[] oldVals = values;

      keys = (K[]) new Object[newCapacity];
      values = (V[]) new Object[newCapacity];

      maxSize = calcMaxSize(newCapacity);
      mask = newCapacity - 1;

      // Insert to the new arrays.
      for (int i = 0; i < oldVals.length; ++i) {
         V oldVal = oldVals[i];
         if (oldVal != null) {
            // Inlined put(), but much simpler: we don't need to worry about
            // duplicated keys, growing/rehashing, or failing to insert.
            K oldKey = oldKeys[i];
            int index = hashIndex(oldKey);

            for (; ; ) {
               if (values[index] == null) {
                  keys[index] = oldKey;
                  values[index] = oldVal;
                  break;
               }

               // Conflict, keep probing. Can wrap around, but never reaches startIndex again.
               index = probeNext(index);
            }
         }
      }
   }

   @Override
   public String toString() {
      if (isEmpty()) {
         return "{}";
      }
      StringBuilder sb = new StringBuilder(4 * size);
      sb.append('{');
      boolean first = true;
      for (int i = 0; i < values.length; ++i) {
         V value = values[i];
         if (value != null) {
            if (!first) {
               sb.append(", ");
            }
            sb.append(keys[i].toString()).append('=').append(value == this ? "(this Map)" :
                    toExternal(value));
            first = false;
         }
      }
      return sb.append('}').toString();
   }

   /**
    * Set implementation for iterating over the entries of the map.
    */
   private final class EntrySet extends AbstractSet<Entry<K, V>> {
      @Override
      public Iterator<Entry<K, V>> iterator() {
         return new MapIterator();
      }

      @Override
      public int size() {
         return LinearProbingHashMap.this.size();
      }
   }

   /**
    * Set implementation for iterating over the keys.
    */
   private final class KeySet extends AbstractSet<K> {
      @Override
      public int size() {
         return LinearProbingHashMap.this.size();
      }

      @Override
      public boolean contains(Object o) {
         return LinearProbingHashMap.this.containsKey(o);
      }

      @Override
      public boolean remove(Object o) {
         return LinearProbingHashMap.this.remove(o) != null;
      }

      @Override
      public boolean retainAll(Collection<?> retainedKeys) {
         boolean changed = false;
         for (Iterator<EEntry<K, V>> iter = entries().iterator(); iter.hasNext(); ) {
            EEntry<K, V> entry = iter.next();
            if (!retainedKeys.contains(entry.key())) {
               changed = true;
               iter.remove();
            }
         }
         return changed;
      }

      @Override
      public void clear() {
         LinearProbingHashMap.this.clear();
      }

      @Override
      public Iterator<K> iterator() {
         return new Iterator<K>() {
            private final Iterator<Entry<K, V>> iter = entrySet.iterator();

            @Override
            public boolean hasNext() {
               return iter.hasNext();
            }

            @Override
            public K next() {
               return iter.next().getKey();
            }

            @Override
            public void remove() {
               iter.remove();
            }
         };
      }
   }

   /**
    * Iterator over primitive entries. Entry key/values are overwritten by each call to {@link #next()}.
    */
   private final class EIterator implements Iterator<EEntry<K, V>>, EEntry<K, V> {
      private int prevIndex = -1;
      private int nextIndex = -1;
      private int entryIndex = -1;

      private void scanNext() {
         while (++nextIndex != values.length && values[nextIndex] == null) {
         }
      }

      @Override
      public boolean hasNext() {
         if (nextIndex == -1) {
            scanNext();
         }
         return nextIndex != values.length;
      }

      @Override
      public EEntry<K, V> next() {
         if (!hasNext()) {
            throw new NoSuchElementException();
         }

         prevIndex = nextIndex;
         scanNext();

         // Always return the same Entry object, just change its index each time.
         entryIndex = prevIndex;
         return this;
      }

      @Override
      public void remove() {
         if (prevIndex == -1) {
            throw new IllegalStateException("next must be called before each remove.");
         }
         if (removeAt(prevIndex)) {
            // removeAt may move elements "back" in the array if they have been displaced because their spot in the
            // array was occupied when they were inserted. If this occurs then the nextIndex is now invalid and
            // should instead point to the prevIndex which now holds an element which was "moved back".
            nextIndex = prevIndex;
         }
         prevIndex = -1;
      }

      // Entry implementation. Since this implementation uses a single Entry, we coalesce that
      // into the Iterator object (potentially making loop optimization much easier).

      @Override
      public K key() {
         return keys[entryIndex];
      }

      @Override
      public V value() {
         return toExternal(values[entryIndex]);
      }

      @Override
      public void setValue(V value) {
         values[entryIndex] = toInternal(value);
      }
   }

   /**
    * Iterator used by the {@link Map} interface.
    */
   private final class MapIterator implements Iterator<Entry<K, V>> {
      private final EIterator iter = new EIterator();

      @Override
      public boolean hasNext() {
         return iter.hasNext();
      }

      @Override
      public Entry<K, V> next() {
         if (!hasNext()) {
            throw new NoSuchElementException();
         }

         iter.next();

         return new MapEntry(iter.entryIndex);
      }

      @Override
      public void remove() {
         iter.remove();
      }
   }

   /**
    * A single entry in the map.
    */
   final class MapEntry implements Entry<K, V> {
      private final int entryIndex;

      MapEntry(int entryIndex) {
         this.entryIndex = entryIndex;
      }

      @Override
      public K getKey() {
         verifyExists();
         return keys[entryIndex];
      }

      @Override
      public V getValue() {
         verifyExists();
         return toExternal(values[entryIndex]);
      }

      @Override
      public V setValue(V value) {
         verifyExists();
         V prevValue = toExternal(values[entryIndex]);
         values[entryIndex] = toInternal(value);
         return prevValue;
      }

      private void verifyExists() {
         if (values[entryIndex] == null) {
            throw new IllegalStateException("The map entry has been removed");
         }
      }
   }
}