ClassBasedMap

import java.util.Arrays;
import java.util.Collection;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

public class ClassBasedMap<K, V> {
	private final Map<Class<? extends K>, V> _cache = new HashMap<Class<? extends K>, V>();
	private final Map<Class<? extends K>, V> _map = new HashMap<Class<? extends K>, V>();

	protected V findCandidate(Class<?> objectType)
	{
		final LinkedList<Class<?>> queue = new LinkedList<Class<?>>();
		queue.add(objectType);
		while (!queue.isEmpty()) {
			final Class<?> candidateType = queue.removeFirst();
			final V candidate = _map.get(candidateType);
			if (candidate != null) {
				return candidate;
			} else {
				final List<Class<?>> candidates = new LinkedList<Class<?>>();
				if ((candidateType.getSuperclass() != null) && (_map.get(candidateType.getSuperclass()) != null)) {
					candidates.add(candidateType.getSuperclass());
				}
				for (Class<?> iface : candidateType.getInterfaces()) {
					if (_map.get(iface) != null) {
						candidates.add(iface);
					}
				}
				switch (candidates.size()) {
				case 0:
					break;

				case 1:
					assert candidates.get(0) != null;
					assert _map.get(candidates.get(0)) != null;
					return _map.get(candidates.get(0));

				default:
					throw new IllegalArgumentException("Multiple drawer candidates found " + candidates);
				}
			}

			if (candidateType.getSuperclass() != null) {
				queue.add(candidateType.getSuperclass());
			}
			queue.addAll(Arrays.asList(candidateType.getInterfaces()));
		}
		return null;
	}

	private void removeFromCache(Class<?> objectType)
	{
		_cache.remove(objectType);
		if (objectType.getSuperclass() != null) {
			removeFromCache(objectType.getSuperclass());
		}
		for (Class<?> iface : objectType.getInterfaces()) {
			removeFromCache(iface);
		}
	}

	public V get(final K obj)
	{
		final Class<?> objectType = obj.getClass();
		return get(objectType);
	}

	public V get(final Class<?> objectType)
	{
		V value = _cache.get(objectType);
		if (value == null) {
			value = findCandidate(objectType);
		}
		return value;
	}

	public boolean add(final Class<? extends K> objectType, final V value)
	{
		if (!_map.containsKey(objectType)) {
			_map.put(objectType, value);
			removeFromCache(objectType);
			return true;
		} else {
			return false;
		}
	}

	public boolean remove(final Class<? extends K> objectType)
	{
		if (_map.remove(objectType) != null) {
			removeFromCache(objectType);
			return true;
		} else {
			return false;
		}
	}

	public V getDirect(final Class<? extends K> objectType)
	{
		return _map.get(objectType);
	}

	public Collection<Class<? extends K>> getRegisteredClasses()
	{
		return _map.keySet();
	}

}