CustomHashMap.java

package com.sayandip.kvstore;

import java.util.ArrayList;
import java.util.List;

/**
 * A hash map built from scratch: separate chaining for collision
 * resolution, dynamic resizing once the load factor is exceeded.
 *
 * This exists to actually understand what java.util.HashMap does under
 * the hood -- bucket indexing off hashCode(), collision handling via
 * chaining, resize/rehash -- not to replace it. KVStore uses this
 * instead of java.util.HashMap on purpose.
 */
public class CustomHashMap<K, V> {

    private static final int DEFAULT_CAPACITY = 16;
    private static final double LOAD_FACTOR = 0.75;

    private static class Entry<K, V> {
        final K key;
        V value;
        Entry<K, V> next; // next entry in this bucket's chain

        Entry(K key, V value) {
            this.key = key;
            this.value = value;
        }
    }

    private Entry<K, V>[] buckets;
    private int size;

    @SuppressWarnings("unchecked")
    public CustomHashMap() {
        buckets = new Entry[DEFAULT_CAPACITY];
        size = 0;
    }

    private int indexFor(K key, int tableLength) {
        int h = (key == null) ? 0 : key.hashCode();
        // spread the high bits into the low bits so hash codes that
        // only differ up top don't all collide -- same trick
        // java.util.HashMap uses (h ^ (h >>> 16))
        h ^= (h >>> 16);
        return (tableLength - 1) & h; // tableLength is always a power of two
    }

    public void put(K key, V value) {
        int idx = indexFor(key, buckets.length);
        Entry<K, V> curr = buckets[idx];
        while (curr != null) {
            if (keysEqual(curr.key, key)) {
                curr.value = value;
                return;
            }
            curr = curr.next;
        }
        Entry<K, V> newEntry = new Entry<>(key, value);
        newEntry.next = buckets[idx];
        buckets[idx] = newEntry;
        size++;

        if ((double) size / buckets.length > LOAD_FACTOR) {
            resize();
        }
    }

    public V get(K key) {
        int idx = indexFor(key, buckets.length);
        Entry<K, V> curr = buckets[idx];
        while (curr != null) {
            if (keysEqual(curr.key, key)) return curr.value;
            curr = curr.next;
        }
        return null;
    }

    public V remove(K key) {
        int idx = indexFor(key, buckets.length);
        Entry<K, V> curr = buckets[idx];
        Entry<K, V> prev = null;
        while (curr != null) {
            if (keysEqual(curr.key, key)) {
                if (prev == null) buckets[idx] = curr.next;
                else prev.next = curr.next;
                size--;
                return curr.value;
            }
            prev = curr;
            curr = curr.next;
        }
        return null;
    }

    public int size() {
        return size;
    }

    private boolean keysEqual(K a, K b) {
        return (a == null && b == null) || (a != null && a.equals(b));
    }

    @SuppressWarnings("unchecked")
    private void resize() {
        Entry<K, V>[] oldBuckets = buckets;
        buckets = new Entry[oldBuckets.length * 2];
        size = 0;
        for (Entry<K, V> head : oldBuckets) {
            Entry<K, V> curr = head;
            while (curr != null) {
                put(curr.key, curr.value); // rehash into the new, bigger table
                curr = curr.next;
            }
        }
    }

    /** Exposed for tests/inspection only -- not part of the real API. */
    int bucketCount() {
        return buckets.length;
    }

    public List<K> keys() {
        List<K> result = new ArrayList<>();
        for (Entry<K, V> head : buckets) {
            for (Entry<K, V> curr = head; curr != null; curr = curr.next) {
                result.add(curr.key);
            }
        }
        return result;
    }
}