ReferenceBag

package com.sci.cclj.computer.util;

import java.lang.ref.Reference;
import java.lang.ref.ReferenceQueue;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Set;
import java.util.function.Function;

public final class ReferenceBag<T, R extends Reference<T>> {
    private final Function<T, R> constructor;
    private final Set<Reference<T>> references;
    private final ReferenceQueue<R> referenceQueue;

    public ReferenceBag(final Function<T, R> constructor) {
        this.constructor = constructor;
        this.references = new HashSet<>();
        this.referenceQueue = new ReferenceQueue<>();
    }

    public void add(final T obj) {
        this.references.add(this.constructor.apply(obj));
    }

    public T poll() {
        final Iterator<Reference<T>> it = this.references.iterator();
        while(it.hasNext()) {
            final Reference<T> ref = it.next();
            final T obj = ref.get();
            if(obj != null) {
                it.remove();
                return obj;
            }
        }
        return null;
    }

    public void clean() {
        for(;;) {
            final Object ref = this.referenceQueue.poll();
            if(ref == null) return;
            this.references.remove(ref);
        }
    }
}