Map decorator that stores the state of map changes and doesn't change the original decorated map

DeltaStateMapDecorator.java

import java.io.Serializable;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.Map;
import java.util.Set;

import org.apache.commons.collections.collection.CompositeCollection;
import org.apache.commons.collections.keyvalue.AbstractMapEntry;
import org.apache.commons.collections.set.CompositeSet;
import org.apache.commons.collections.set.UnmodifiableSet;


/**
 * 
 * Map decorator that stores the state of map changes to only the instance of this class.
 * Doesn't change the original decorated map.
 * 
 * Map iterators are not protected against changes in all cases.
 * 
 * 
 * @author Lari Hotari
 */
public class DeltaStateMapDecorator<K, V> implements Map<K, V>, Serializable, Cloneable {
	private static final long serialVersionUID = 1L;
	protected final Map<K, V> wrapped;
    private Set<K> removedKeys;
    private Map<K, V> changes;
    private Map<K, V> adds;

    public DeltaStateMapDecorator(Map<K, V> wrapped) {
    	this.wrapped=wrapped;
    	this.removedKeys=new HashSet<K>();
    	this.changes=new HashMap<K, V>();
    	this.adds=new LinkedHashMap<K, V>();
    }

    public Map<K, V> getWrappedMap() {
        return wrapped;
    }

    /**
     * @see Map#clear()
     */
    public void clear() {
    	adds.clear();
    	changes.clear();
    	removedKeys.addAll(wrapped.keySet());
    }

    /**
     * @see Map#size()
     */
    public int size() {
    	return wrapped.size() - removedKeys.size() + adds.size();
    }

    /**
     * @see Map#isEmpty()
     */
    public boolean isEmpty() {
        return (size()==0);
    }

    /**
     * @see Map#containsKey(java.lang.Object)
     */
    public boolean containsKey(Object key) {
        return (adds.containsKey(key) || (!removedKeys.contains(key) && wrapped.containsKey(key)));
    }

    /**
     * @see Map#containsValue(java.lang.Object)
     */
    public boolean containsValue(Object value) {
    	if(changes.size()==0 && removedKeys.size()==0) {
    		return (adds.containsValue(value) || wrapped.containsValue(value));
    	} else if(adds.containsValue(value) || changes.containsValue(value)) {
    		return true;
    	} else {
    		for(Map.Entry<K, V> entry:wrapped.entrySet()) {
    			Object entryValue=entry.getValue();
    			if(!removedKeys.contains(entry.getKey()) && !changes.containsKey(entry.getKey()) && (entryValue==value || (entryValue != null && entryValue.equals(value)))) {
    				return true;
    			}
    		}
    		return false;
    	}
    }

    /**
     * @see Map#values()
     */
    public Collection<V> values() {
    	if(removedKeys.size() > 0 || changes.size() > 0) {
	        Collection<V> values = new ArrayList<V>(size());
			for(Map.Entry<K, V> entry:wrapped.entrySet()) {
				if(!removedKeys.contains(entry.getKey())) {
					if(changes.containsKey(entry.getKey())) {
						values.add(changes.get(entry.getKey()));
					} else {
						values.add(entry.getValue());
					}
				}
			}
			values.addAll(adds.values());
			return values;
    	} else if(adds.size() > 0) {
    		return (Collection<V>)new CompositeCollection(new Collection[]{wrapped.values(),adds.values()});
    	} else {
    		return wrapped.values();
    	}
    }

    /**
     * @see Map#putAll(java.util.Map)
     */
    public void putAll(Map<? extends K, ? extends V> map) {
    	for(Map.Entry<? extends K, ? extends V> entry:map.entrySet()) {
    		K key=entry.getKey();
    		if(removedKeys.remove(key) || wrapped.containsKey(key)) {
    			changes.put(key, entry.getValue());
    		} else {
    			adds.put(key, entry.getValue());
    		}
    	}
    }

    /**
     * @see Map#entrySet()
     */
    public Set<Map.Entry<K, V>> entrySet() {
    	if(removedKeys.size() > 0 || changes.size() > 0) {
	        Set<Map.Entry<K, V>> entries=new LinkedHashSet<Map.Entry<K,V>>(size());
			for(Map.Entry<K, V> entry:wrapped.entrySet()) {
				K key=entry.getKey();
				if(!removedKeys.contains(key)) {
					V val=(changes.containsKey(key)) ? changes.get(key) : entry.getValue();
					entries.add(new HashEntry(key, val));
				}
			}
			entries.addAll(adds.entrySet());
			return entries;
    	} else if(adds.size() > 0) {
    		return (Set<Map.Entry<K, V>>)new CompositeSet(new Set[]{wrapped.entrySet(),adds.entrySet()});
    	} else {
    		return (Set<Map.Entry<K, V>>)UnmodifiableSet.decorate(wrapped.entrySet());
    	}
    }

    /**
     * @see Map#keySet()
     */
    public Set<K> keySet() {
        if(removedKeys.size() > 0) {
        	Set<K> keys=new LinkedHashSet<K>(wrapped.keySet());
        	keys.removeAll(removedKeys);
        	keys.addAll(adds.keySet());
        	return keys;
        } else if (adds.size() > 0) {
    		return (Set<K>)new CompositeSet(new Set[]{wrapped.keySet(),adds.keySet()});
    	} else {
    		return (Set<K>)UnmodifiableSet.decorate(wrapped.keySet());
        }
    }

    /**
     * @see Map#get(java.lang.Object)
     */
    public V get(Object key) {
    	if(changes.containsKey(key)) {
    		return changes.get(key);
    	} else if(!removedKeys.contains(key) && wrapped.containsKey(key)) {
    		return wrapped.get(key);
    	} else {
    		return adds.get(key);
    	}
    }

    /**
     * @see Map#remove(java.lang.Object)
     */
    public V remove(Object key) {
    	if(changes.containsKey(key)) {
    		removedKeys.add((K)key);
    		return changes.remove(key);
    	} else if(!removedKeys.contains(key) && wrapped.containsKey(key)) {
    		removedKeys.add((K)key);
    		return wrapped.get(key);
    	} else if (adds.containsKey(key)) {
    		return adds.remove(key);
    	} else {
    		return null;
    	}
    }

    /**
     * @see Map#put(java.lang.Object, java.lang.Object)
     */
    public V put(K key, V value) {
    	if(wrapped.containsKey(key)) {
    		V oldval;
    		if(removedKeys.contains(key)) {
    			removedKeys.remove(key);
    			oldval=null;
    		} else if(changes.containsKey(key)) {
    			oldval=changes.get(key);
    		} else {
    			oldval=wrapped.get(key);
    		}
    		changes.put(key, value);
    		return oldval;
    	} else {
    		return adds.put(key, value);
    	}
    }

    protected static class HashEntry extends AbstractMapEntry {
        protected HashEntry(Object key, Object value) {
        	super(key, value);
        }
    }

	@Override
	protected Object clone() throws CloneNotSupportedException {
		DeltaStateMapDecorator<K,V> cloned=(DeltaStateMapDecorator<K,V>)super.clone();
    	cloned.removedKeys=new HashSet<K>(removedKeys);
    	cloned.changes=new HashMap<K, V>(changes);
    	cloned.adds=new LinkedHashMap<K, V>(adds);
		return cloned;
	}
}