package mb.statix.concurrent.util;

import org.metaborg.util.collection.CapsuleUtil;

import com.google.common.collect.Streams;

import io.usethesource.capsule.Set;
import io.usethesource.capsule.SetMultimap;
import io.usethesource.capsule.util.stream.CapsuleCollectors;
import mb.nabl2.terms.ITermVar;
import mb.nabl2.terms.unification.u.IUnifier;

public class VarIndexedCollection<V> {

    private final SetMultimap.Transient<ITermVar, Entry> index;

    public VarIndexedCollection() {
        this.index = SetMultimap.Transient.of();
    }

    public Set.Immutable<Entry> get(ITermVar indexVar) {
        return index.get(indexVar);
    }

    public boolean put(V value, Iterable<ITermVar> indexVars, IUnifier unifier) {
        final Set<ITermVar> vars =
                Streams.stream(indexVars).flatMap(v -> unifier.getVars(v).stream()).collect(CapsuleCollectors.toSet());
        final Entry entry = new Entry(value);
        for(ITermVar var : vars) {
            index.__insert(var, entry.inc());
        }
        return !vars.isEmpty();
    }

    public Set.Immutable<V> update(Iterable<ITermVar> indexVars, IUnifier unifier) {
        final Set.Transient<V> done = CapsuleUtil.transientSet();
        for(ITermVar indexVar : indexVars) {
            update(indexVar, unifier, done);
        }
        return done.freeze();
    }

    public Set.Immutable<V> update(ITermVar indexVar, IUnifier unifier) {
        final Set.Transient<V> done = CapsuleUtil.transientSet();
        update(indexVar, unifier, done);
        return done.freeze();
    }

    private void update(ITermVar indexVar, IUnifier unifier, Set.Transient<V> done) {
        final Set<ITermVar> vars = unifier.getVars(indexVar);
        final Set<Entry> entries = index.get(indexVar);
        index.__remove(indexVar);
        for(Entry entry : entries) {
            entry.dec();
            for(ITermVar var : vars) {
                index.__insert(var, entry.inc());
            }
            if(entry.refcount() == 0) {
                done.__insert(entry.value);
            }
        }
    }

    public class Entry {
        private final V value;
        private int refcount;

        public Entry(V value) {
            this.value = value;
            this.refcount = 0;
        }

        public V value() {
            return value;
        }

        public int refcount() {
            return refcount;
        }

        public Entry dec() {
            refcount -= 1;
            return this;
        }

        public Entry inc() {
            refcount += 1;
            return this;
        }

    }

}
@@ -112,16 +112,106 @@
 collect.
+HashMultimap;%0Aimport com.google.common.collect.Multimap;%0Aimport com.google.common.collect.
 Streams;
@@ -250,52 +250,8 @@
 .Set
-;%0Aimport io.usethesource.capsule.SetMultimap
 ;%0Aim
@@ -447,27 +447,24 @@
 e final 
-Set
 Multimap
 .Transie
@@ -455,26 +455,16 @@
 Multimap
-.Transient
 %3CITermVa
@@ -538,19 +538,20 @@
 index = 
-Set
+Hash
 Multimap
@@ -555,121 +555,15 @@
 map.
-Transient.of();%0A    %7D%0A%0A    public Set.Immutable%3CEntry%3E get(ITermVar indexVar) %7B%0A        return index.get(indexVar
+create(
 );%0A 
@@ -896,39 +896,34 @@
           index.
-__inser
+pu
 t(var, entry.inc
@@ -1644,19 +1644,24 @@
   final 
-Set
+Iterable
 %3CEntry%3E 
@@ -1680,45 +1680,17 @@
 dex.
-get(indexVar);%0A        index.__
 remove
+All
 (ind
@@ -1820,23 +1820,18 @@
   index.
-__inser
+pu
 t(var, e
@@ -2493,8 +2493,9 @@
     %7D%0A%0A%7D
+%0A

Use more efficient multimap in temporary structure