Map where two keys map to one value

DoubleKeyHashMap.java

package net.darkseraphim.dev

import lombok.EqualsAndHashCode;
import lombok.RequiredArgsConstructor;
import org.apache.commons.lang.Validate;

import java.util.AbstractSet;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @author DarkSeraphim
 */
public class DoubleKeyHashMap<K, L, V> {

    private final Class<K> k;
    private final Class<L> l;
    // They should provide the ordering such that if I put a keys-value pair, it should be at
    // the same position in the Set<Entry>
    private final Map<K, Entry<K, L, V>> kv = new ConcurrentHashMap<K, Entry<K, L, V>>();
    private final Map<L, Entry<K, L, V>> lv = new ConcurrentHashMap<L, Entry<K, L, V>>();
    private Entry<K, L, V> head;

    private DoubleKeyHashMap(Class<K> k, Class<L> l) {
        this.k = k;
        this.l = l;
    }

    public static <K, L, V> DoubleKeyHashMap<K, L, V> create(Class<K> k, Class<L> l, Class<V> v) {
        if (k == l) {
            throw new IllegalArgumentException("K cannot be the same as L");
        }
        return new DoubleKeyHashMap<K, L, V>(k, l);
    }

    synchronized public V get(Object key) {
        Validate.notNull(key, "null keys not supported");
        if (key.getClass().isAssignableFrom(this.k)) {
            return this.getK((K) key);

        } else if (key.getClass().isAssignableFrom(this.l)) {
            return this.getL((L) key);
        }
        throw new ClassCastException(String.format("Cannot cast %s to %s nor %s", key.getClass().toString(), this.k.toString(), this.l.toString()));
    }

    private V getK(K key) {
        Entry<K, L, V> e = this.kv.get(key);
        return e != null ? (V) e.getValue() : null;
    }

    private V getL(L key) {
        Entry<K, L, V> e = this.lv.get(key);
        return e != null ? (V) e.getValue() : null;
    }

    synchronized public V put(K key, L key2, V val) {
        Entry<K, L, V> entry = new Entry<K, L, V>(key, key2, val);
        if (this.head != null) {
            if (this.head.equals(entry)) {
                V old = (V) this.head.getValue();
                this.head.setValue(val);
                return old;
            }
            entry.next = this.head;
            entry.prev = null;
        }
        this.head = entry;
        Entry<K, L, V> a = this.kv.put(key, entry);
        Entry<K, L, V> b = this.lv.put(key2, entry);
        if (a != b) {
            throw new IllegalStateException("Contract violated");
        }
        // Link the next and prev Entries
        fixEntry(a);
        return a != null ? (V) a.getValue() : null;
    }

    synchronized public V remove(Object key) {
        Validate.notNull(key, "null keys not supported");
        if (key.getClass().isAssignableFrom(this.k)) {
            return this.removeK((K) key);
        } else if (key.getClass().isAssignableFrom(this.l)) {
            return this.removeL((L) key);
        }
        throw new ClassCastException(String.format("Cannot cast %s to %s nor %s", key.getClass().toString(), this.k.toString(), this.l.toString()));
    }

    private V removeK(K key) {
        // Get entry, remove entry
        Entry<K, L, V> a = this.kv.remove(key);
        Entry<K, L, V> b = null;
        if (a != null)
            b = this.lv.remove(a.getSecondKey());
        if (a != b) {
            throw new IllegalStateException("Contract violated");
        }
        fixEntry(a);
        return a != null ? (V) a.getValue() : null;
    }

    private V removeL(L key) {
        // Get entry, remove entry
        Entry<K, L, V> a = this.lv.remove(key);
        Entry<K, L, V> b = null;
        if (a != null)
            b = this.kv.remove(a.getKey());
        if (a != b) {
            throw new IllegalStateException("Contract violated");
        }
        fixEntry(a);
        return a != null ? (V) a.getValue() : null;
    }

    synchronized public boolean containsKey(Object key) {
        Validate.notNull(key, "null keys not supported");
        if (key.getClass().isAssignableFrom(this.k)) {
            return this.containsKeyK((K) key);
        } else if (key.getClass().isAssignableFrom(this.l)) {
            return this.containsKeyL((L) key);
        }
        throw new ClassCastException(String.format("Cannot cast %s to %s nor %s", key.getClass().toString(), this.k.toString(), this.l.toString()));
    }

    synchronized public boolean containsKeyK(K key) {
        Entry<K, L, V> entry = this.kv.get(key);
        if (entry != null && this.lv.get(entry.getSecondKey()) == null) {
            throw new IllegalStateException("Contract violated");
        }
        return entry != null;
    }

    synchronized public boolean containsKeyL(L key) {
        Entry<K, L, V> entry = this.lv.get(key);
        if (entry != null && this.kv.get(entry.getKey()) == null) {
            throw new IllegalStateException("Contract violated");
        }
        return entry != null;
    }

    private void fixEntry(Entry<K, L, V> entry) {
        if (entry != null) {
            if (entry.next != null) {
                entry.next.prev = entry.prev;
            }
            if (entry.prev != null) {
                entry.prev.next = entry.next;
            }
        }
    }

    public Set<Entry<K, L, V>> entrySet() {
        return new AbstractSet<Entry<K, L, V>>() {
            @Override
            public Iterator<Entry<K, L, V>> iterator() {
                return new EntryIterator(DoubleKeyHashMap.this);
            }

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

    public boolean isEmpty() {
        boolean ret = this.kv.isEmpty();
        if (this.lv.isEmpty() != ret)
            throw new IllegalStateException("Contract violated");
        return ret;
    }

    @EqualsAndHashCode
    public class Entry<K, L, V> implements Map.Entry {

        private final K k;
        private final L l;
        private transient Entry<K, L, V> prev;
        private transient Entry<K, L, V> next;
        private transient volatile V v;

        private Entry(K k, L l, V v) {
            this.k = k;
            this.l = l;
            this.v = v;
        }

        /**
         * Returns key K for the Entry
         *
         * @return
         */
        @Override
        public Object getKey() {
            return this.k;
        }

        /**
         * Returns key L for the entry
         *
         * @return
         */
        public Object getSecondKey() {
            return this.l;
        }

        @Override
        public Object getValue() {
            return this.v;
        }

        @Override
        public Object setValue(Object value) {
            V old = this.v;
            this.v = (V) value;
            return old;
        }

        private void remove(Map<K, ?> a, Map<L, ?> b) {
            a.remove(this.getKey());
            b.remove(this.getSecondKey());
        }
    }

    @RequiredArgsConstructor
    class EntryIterator implements Iterator<Entry<K, L, V>> {

        private final DoubleKeyHashMap backend;
        Entry<K, L, V> current;
        private boolean before = true;

        @Override
        public boolean hasNext() {
            return before ? backend.head != null : current.next != null;
        }

        @Override
        public Entry<K, L, V> next() {
            if (!hasNext())
                throw new IllegalStateException("No next entry found");
            current = before ? backend.head : current.next;
            if (before)
                before = false;
            return current;
        }

        @Override
        public void remove() {
            current.remove(backend.kv, backend.lv);
        }
    }
}