org.nd4j.common.collections.WeakIdentityHashMap Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of nd4j-common Show documentation
There is a newer version: 1.0.0-M2.1
package org.nd4j.common.collections;

import lombok.*;

import java.lang.ref.Reference;
import java.lang.ref.ReferenceQueue;
import java.lang.ref.WeakReference;
import java.util.*;

/**
 * A hash map implementation with weak identity keys.
 * For details, see {@link WeakHashMap} and {@link IdentityHashMap}
 *
 * @param  Key type
 * @param  Value type
 * @author Alex Black
 */
public class WeakIdentityHashMap implements Map {

    protected final Map, V> map;
    protected final ReferenceQueue refQueue;

    public WeakIdentityHashMap(){
        map = new HashMap<>();
        refQueue = new ReferenceQueue<>();
    }

    //Clear references to any map keys that have been GC'd
    protected void clearReferences(){
        Reference r;
        while((r = refQueue.poll()) != null){
            map.remove(r);
        }
    }

    @Override
    public int size() {
        clearReferences();
        return map.size();
    }

    @Override
    public boolean isEmpty() {
        clearReferences();
        return map.isEmpty();
    }

    @Override
    public boolean containsKey(Object key) {
        clearReferences();
        return map.containsKey(new KeyRef<>(key));
    }

    @Override
    public boolean containsValue(Object value) {
        clearReferences();
        return map.containsValue(value);
    }

    @Override
    public V get(Object key) {
        clearReferences();
        return map.get(new KeyRef<>(key));
    }

    @Override
    public V put(K key, V value) {
        clearReferences();
        map.put(new KeyRef<>(key), value);
        return value;
    }

    @Override
    public V remove(Object key) {
        clearReferences();
        return map.remove(new KeyRef<>(key));
    }

    @Override
    public void putAll(Map m) {
        clearReferences();
        for(Map.Entry e : m.entrySet()){
            map.put(new KeyRef<>(e.getKey()), e.getValue());
        }
    }

    @Override
    public void clear() {
        map.clear();
        clearReferences();
    }

    @Override
    public Set keySet() {
        clearReferences();
        Set ret = new HashSet<>();
        for(KeyRef k : map.keySet() ){
            K key = k.get();
            if(key != null)
                ret.add(key);
        }
        return ret;
    }

    @Override
    public Collection values() {
        clearReferences();
        return map.values();
    }

    @Override
    public Set> entrySet() {
        clearReferences();
        Set> ret = new HashSet<>();
        for(Map.Entry, V> e : map.entrySet()){
            K k = e.getKey().get();
            if(k != null){
                ret.add(new Entry(k, e.getValue()));
            }
        }
        return ret;
    }


    protected static class KeyRef extends WeakReference {
        private final int hash;
        public KeyRef(@NonNull K referent) {
            super(referent);
            this.hash = System.identityHashCode(referent);
        }

        @Override
        public int hashCode(){
            return hash;
        }

        @Override
        public boolean equals(Object o){
            if(this == o){
                return true;
            }
            if(o instanceof WeakReference){
                return this.get() == ((WeakReference) o).get();
            }
            return false;
        }
    }

    @Data
    @AllArgsConstructor
    protected static class Entry implements Map.Entry {
        protected K key;
        protected V value;

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