<commit_before>package mb.p_raffrayi.impl;

import static com.google.common.collect.Streams.stream;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.stream.Collectors;

import javax.annotation.Nullable;

import org.metaborg.util.Ref;
import org.metaborg.util.collection.CapsuleUtil;
import org.metaborg.util.collection.MultiSet;
import org.metaborg.util.collection.MultiSetMap;
import org.metaborg.util.functions.Function1;
import org.metaborg.util.functions.Function2;
import org.metaborg.util.future.AggregateFuture;
import org.metaborg.util.future.AggregateFuture.SC;
import org.metaborg.util.future.CompletableFuture;
import org.metaborg.util.future.ICompletableFuture;
import org.metaborg.util.future.IFuture;
import org.metaborg.util.iterators.Iterables2;
import org.metaborg.util.log.ILogger;
import org.metaborg.util.log.LoggerUtils;
import org.metaborg.util.tuple.Tuple2;
import org.metaborg.util.unit.Unit;

import com.google.common.collect.ImmutableMultimap;
import com.google.common.collect.Iterables;
import com.google.common.collect.Sets;

import io.usethesource.capsule.Set;
import mb.p_raffrayi.IIncrementalTypeCheckerContext;
import mb.p_raffrayi.IRecordedQuery;
import mb.p_raffrayi.IScopeGraphLibrary;
import mb.p_raffrayi.ITypeChecker;
import mb.p_raffrayi.ITypeChecker.IOutput;
import mb.p_raffrayi.ITypeChecker.IState;
import mb.p_raffrayi.IUnitResult;
import mb.p_raffrayi.IUnitResult.TransitionTrace;
import mb.p_raffrayi.actors.IActor;
import mb.p_raffrayi.actors.IActorRef;
import mb.p_raffrayi.impl.confirm.ConfirmResult;
import mb.p_raffrayi.impl.confirm.IConfirmation;
import mb.p_raffrayi.impl.confirm.IConfirmationContext;
import mb.p_raffrayi.impl.confirm.IConfirmationFactory;
import mb.p_raffrayi.impl.confirm.OptimisticConfirmation;
import mb.p_raffrayi.impl.diff.AddingDiffer;
import mb.p_raffrayi.impl.diff.IDifferContext;
import mb.p_raffrayi.impl.diff.IDifferDataOps;
import mb.p_raffrayi.impl.diff.IDifferOps;
import mb.p_raffrayi.impl.diff.IScopeGraphDiffer;
import mb.p_raffrayi.impl.diff.MatchingDiffer;
import mb.p_raffrayi.impl.diff.ScopeDiff;
import mb.p_raffrayi.impl.diff.ScopeGraphDiffer;
import mb.p_raffrayi.impl.diff.StaticDifferContext;
import mb.p_raffrayi.impl.envdiff.EnvDiffer;
import mb.p_raffrayi.impl.envdiff.IEnvDiff;
import mb.p_raffrayi.impl.envdiff.IEnvDiffer;
import mb.p_raffrayi.impl.envdiff.IEnvDifferContext;
import mb.p_raffrayi.impl.envdiff.IndexedEnvDiffer;
import mb.p_raffrayi.impl.tokens.CloseLabel;
import mb.p_raffrayi.impl.tokens.CloseScope;
import mb.p_raffrayi.impl.tokens.Confirm;
import mb.p_raffrayi.impl.tokens.DifferState;
import mb.p_raffrayi.impl.tokens.EnvDifferState;
import mb.p_raffrayi.impl.tokens.IWaitFor;
import mb.p_raffrayi.impl.tokens.InitScope;
import mb.p_raffrayi.impl.tokens.Query;
import mb.p_raffrayi.nameresolution.DataLeq;
import mb.p_raffrayi.nameresolution.DataWf;
import mb.p_raffrayi.nameresolution.IQuery;
import mb.p_raffrayi.nameresolution.NameResolutionQuery;
import mb.p_raffrayi.nameresolution.StateMachineQuery;
import mb.scopegraph.ecoop21.LabelOrder;
import mb.scopegraph.ecoop21.LabelWf;
import mb.scopegraph.oopsla20.IScopeGraph;
import mb.scopegraph.oopsla20.diff.BiMap;
import mb.scopegraph.oopsla20.path.IResolutionPath;
import mb.scopegraph.oopsla20.reference.EdgeOrData;
import mb.scopegraph.oopsla20.reference.ScopeGraph;
import mb.scopegraph.oopsla20.terms.newPath.ScopePath;
import mb.scopegraph.patching.IPatchCollection;
import mb.scopegraph.patching.PatchCollection;
import mb.scopegraph.patching.Patcher;
import mb.scopegraph.resolution.StateMachine;

class TypeCheckerUnit<S, L, D, R extends IOutput<S, L, D>, T extends IState<S, L, D>>
        extends AbstractUnit<S, L, D, Result<S, L, D, R, T>> implements IIncrementalTypeCheckerContext<S, L, D, R, T> {


    private static final ILogger logger = LoggerUtils.logger(TypeCheckerUnit.class);

    private final ITypeChecker<S, L, D, R, T> typeChecker;
    private final boolean changed;
    private final @Nullable IUnitResult<S, L, D, Result<S, L, D, R, T>> previousResult;

    private volatile UnitState state;

    private final IPatchCollection.Transient<S> matchedBySharing = PatchCollection.Transient.of();

    private final Ref<IScopeGraph.Immutable<S, L, D>> localScopeGraph = new Ref<>(ScopeGraph.Immutable.of());

    private final Set.Transient<String> addedUnitIds = CapsuleUtil.transientSet();

    private final Ref<StateCapture<S, L, D, T>> localCapture = new Ref<>();
    private final ICompletableFuture<Unit> whenActive = new CompletableFuture<>(); // activated when transitioning from unknown to active/released.
    private final ICompletableFuture<Unit> whenContextActivated = new CompletableFuture<>(); // activated when transitioning from local/unknown to active.

    private IEnvDiffer<S, L, D> envDiffer;

    private final IPatchCollection.Transient<S> externalMatches = PatchCollection.Transient.of();

    private final IConfirmationFactory<S, L, D> confirmation;
    private final ICompletableFuture<Optional<IPatchCollection.Immutable<S>>> confirmationResult =
            new CompletableFuture<>();

    // Intermediate confirmation result aggregations
    private final IPatchCollection.Transient<S> resultPatches = PatchCollection.Transient.of();
    private final IPatchCollection.Transient<S> globalPatches = PatchCollection.Transient.of();
    private final Set.Transient<IRecordedQuery<S, L, D>> addedQueries = CapsuleUtil.transientSet();
    private final Set.Transient<IRecordedQuery<S, L, D>> removedQueries = CapsuleUtil.transientSet();

    TypeCheckerUnit(IActor<? extends IUnit<S, L, D, Result<S, L, D, R, T>>> self,
            @Nullable IActorRef<? extends IUnit<S, L, D, ?>> parent, IUnitContext<S, L, D> context,
            ITypeChecker<S, L, D, R, T> unitChecker, Iterable<L> edgeLabels, boolean inputChanged,
            IUnitResult<S, L, D, Result<S, L, D, R, T>> previousResult) {
        super(self, parent, context, edgeLabels);
        this.typeChecker = unitChecker;
        this.changed = inputChanged;
        this.previousResult = previousResult;
        this.state = UnitState.INIT_UNIT;
        this.confirmation = OptimisticConfirmation.factory();
    }

    TypeCheckerUnit(IActor<? extends IUnit<S, L, D, Result<S, L, D, R, T>>> self,
            @Nullable IActorRef<? extends IUnit<S, L, D, ?>> parent, IUnitContext<S, L, D> context,
            ITypeChecker<S, L, D, R, T> unitChecker, Iterable<L> edgeLabels) {
        this(self, parent, context, unitChecker, edgeLabels, true, null);
    }

    private final MultiSetMap.Transient<D, ICompletableFuture<D>> pendingExternalDatums = MultiSetMap.Transient.of();

    @Override protected IFuture<D> getExternalDatum(D datum) {
        if(!changed && previousResult != null && previousResult.result().localState() != null) {
            // when previous result is present, reliable (i.e., !changed), and has local state
            // try to find datum from previous state
            final Optional<D> datumOpt =
                    previousResult.result().localState().typeCheckerState().tryGetExternalDatum(datum);
            if(datumOpt.isPresent()) {
                return CompletableFuture.completedFuture(datumOpt.get());
            }
        }
        return whenActive.compose((u, ex) -> {
            if(this.state.equals(UnitState.RELEASED) || this.state.equals(UnitState.DONE)
                    && this.stateTransitionTrace.equals(TransitionTrace.RELEASED)) {
                final D result = previousResult.result().analysis().getExternalRepresentation(datum);
                return CompletableFuture.completedFuture(context.substituteScopes(result, resultPatches.patches().invert()));
            }
            final IFuture<D> future = typeChecker.getExternalDatum(datum);
            if(future.isDone()) {
                return future;
            }
            final ICompletableFuture<D> result = new CompletableFuture<>();
            pendingExternalDatums.put(datum, result);
            future.whenComplete(result::complete);
            return result.whenComplete((__, ex2) -> {
                pendingExternalDatums.remove(datum, result);
            });
        });
    }

    @Override protected D getPreviousDatum(D datum) {
        assertPreviousResultProvided();
        return previousResult.result().analysis().getExternalRepresentation(datum);
    }

    ///////////////////////////////////////////////////////////////////////////
    // IBroker2UnitProtocol interface, called by IBroker implementations
    ///////////////////////////////////////////////////////////////////////////

    @Override public IFuture<IUnitResult<S, L, D, Result<S, L, D, R, T>>> _start(List<S> rootScopes) {
        assertInState(UnitState.INIT_UNIT);
        resume();

        doStart(rootScopes);
        if(isIncrementalEnabled() && previousResult != null) {
            if(previousResult.rootScopes().size() != rootScopes.size()) {
                throw new IllegalStateException("Root scope counts do not match.");
            }
            Iterables2.zip(rootScopes, previousResult.rootScopes(), Tuple2::of).forEach(match -> {
                matchedBySharing.put(match._1(), match._2());
                externalMatches.put(match._1(), match._2());
            });
        }

        final IFuture<Result<S, L, D, R, T>> result;
        try {
            state = UnitState.INIT_TC;
            result = this.typeChecker.run(this, rootScopes).whenComplete((r, ex) -> {
                if(state == UnitState.INIT_TC) {
                    stateTransitionTrace = TransitionTrace.INITIALLY_STARTED;
                }
                if(inLocalPhase()) {
                    doCapture();
                }
                state = UnitState.DONE;
                resume();
                tryFinish();
            }).thenApply(r -> {
                return Result.<S, L, D, R, T>of(r, localCapture.get(), localScopeGraph.get(),
                        matchedBySharing.patchRange());
            });
        } catch(Exception e) {
            logger.error("Exception starting type-checker {}.", e);
            return doFinish(CompletableFuture.completedExceptionally(e));
        }

        // Start phantom units for all units that have not yet been restarted
        if(isIncrementalEnabled() && previousResult != null) {
            for(String removedId : Sets.difference(previousResult.subUnitResults().keySet(), addedUnitIds)) {
                @SuppressWarnings("unchecked") final IUnitResult<S, L, D, ? extends IOutput<S, L, D>> subResult =
                        (IUnitResult<S, L, D, ? extends IOutput<S, L, D>>) previousResult.subUnitResults()
                                .get(removedId);
                this.<Unit>doAddSubUnit(removedId,
                        (subself, subcontext) -> new PhantomUnit<>(subself, self, subcontext, edgeLabels, subResult),
                        new ArrayList<>(), true);
            }
        }

        if(state == UnitState.INIT_TC) {
            // runIncremental not called, so start eagerly
            doRestart(false);
        } else if(state == UnitState.DONE) {
            // Completed synchronously
            whenActive.complete(Unit.unit);
            confirmationResult.complete(Optional.of(PatchCollection.Immutable.of()));
        }

        return doFinish(result);
    }

    ///////////////////////////////////////////////////////////////////////////
    // IUnit2UnitProtocol interface, called by IUnit implementations
    ///////////////////////////////////////////////////////////////////////////

    @Override public IFuture<ConfirmResult<S, L, D>> _confirm(ScopePath<S, L> path, LabelWf<L> labelWF,
            DataWf<S, L, D> dataWF, boolean prevEnvEmpty) {
        assertConfirmationEnabled();
        stats.incomingConfirmations++;
        final IActorRef<? extends IUnit<S, L, D, ?>> sender = self.sender(TYPE);
        final ICompletableFuture<ConfirmResult<S, L, D>> result = new CompletableFuture<>();
        whenActive.whenComplete((__, ex) -> {
            if(ex != null && ex != Release.instance) {
                result.completeExceptionally(ex);
            } else {
                // Cannot immediately confirm on `Release.instance` exception
                // because subunits might change edges in shared scopes.
                confirmation.getConfirmation(new ConfirmationContext(sender))
                        .confirm(path, labelWF, dataWF, prevEnvEmpty).whenComplete(result::complete);
            }
        });
        return result;
    }

    @Override public IFuture<IQueryAnswer<S, L, D>> _queryPrevious(ScopePath<S, L> path, IQuery<S, L, D> query,
            DataWf<S, L, D> dataWF, DataLeq<S, L, D> dataEquiv) {
        assertPreviousResultProvided();
        return doQueryPrevious(self.sender(TYPE), previousResult.scopeGraph(), path, query, dataWF, dataEquiv);
    }

    @Override public IFuture<Optional<S>> _match(S previousScope) {
        assertOwnScope(previousScope);
        if(matchedBySharing.identityPatches().contains(previousScope)) {
            return CompletableFuture.completedFuture(Optional.of(previousScope));
        }
        if(matchedBySharing.patchDomain().contains(previousScope)) {
            return CompletableFuture.completedFuture(Optional.of(matchedBySharing.patch(previousScope)));
        }
        if(!changed && previousResult.result().localState() != null
                && previousResult.result().localState().scopes().contains(previousScope)) {
            return CompletableFuture.completedFuture(Optional.of(previousScope));
        }
        return super._match(previousScope);
    }

    @Override public IFuture<StateSummary<S, L, D>> _state() {
        // resume();
        if(state.equals(UnitState.ACTIVE)
                || (state == UnitState.DONE && stateTransitionTrace != TransitionTrace.RELEASED)) {
            return CompletableFuture.completedFuture(StateSummary.restart(process, dependentSet()));
        }

        if(state.equals(UnitState.RELEASED)
                || (state == UnitState.DONE && stateTransitionTrace == TransitionTrace.RELEASED)) {
            return CompletableFuture.completedFuture(StateSummary.released(process, dependentSet()));
        }

        return CompletableFuture.completedFuture(StateSummary.release(process, dependentSet()));
    }

    @Override public void _restart() {
        assertIncrementalEnabled();
        if(doRestart(true)) {
            stateTransitionTrace = TransitionTrace.RESTARTED;
        }
    }

    @Override public void _release() {
        if(state == UnitState.UNKNOWN) {
            doRelease(CapsuleUtil.toSet(this.addedQueries), CapsuleUtil.toSet(this.removedQueries),
                    PatchCollection.Immutable.<S>of().putAll(this.resultPatches),
                    PatchCollection.Immutable.<S>of().putAll(globalPatches));
        }
    }

    ///////////////////////////////////////////////////////////////////////////
    // ITypeCheckerContext interface, called by ITypeChecker implementations
    ///////////////////////////////////////////////////////////////////////////

    // NB. Invoke methods via `local` so that we have the same scheduling & ordering
    // guarantees as for remote calls.

    @Override public String id() {
        return self.id();
    }

    @SuppressWarnings("unchecked") @Override public <Q extends IOutput<S, L, D>, U extends IState<S, L, D>>
            IFuture<IUnitResult<S, L, D, Result<S, L, D, Q, U>>>
            add(String id, ITypeChecker<S, L, D, Q, U> unitChecker, List<S> rootScopes, boolean changed) {
        assertActive();

        // No previous result for subunit
        if(this.previousResult == null || !this.previousResult.subUnitResults().containsKey(id)) {
            this.addedUnitIds.__insert(id);

            return ifActive(whenContextActive(this.<Result<S, L, D, Q, U>>doAddSubUnit(id, (subself, subcontext) -> {
                return new TypeCheckerUnit<>(subself, self, subcontext, unitChecker, edgeLabels);
            }, rootScopes, false)._2()));
        }

        final IUnitResult<S, L, D, Result<S, L, D, Q, U>> subUnitPreviousResult;
        if(isIncrementalEnabled()) {

            subUnitPreviousResult =
                    (IUnitResult<S, L, D, Result<S, L, D, Q, U>>) this.previousResult.subUnitResults().get(id);


            // When a scope is shared, the shares must be consistent.
            // Also, it is not necessary that shared scopes are reachable from the root scopes
            // (A unit started by the Broker does not even have root scopes)
            // Therefore we enforce here that the current root scopes and the previous ones match.
            List<S> previousRootScopes = subUnitPreviousResult.rootScopes();

            int pSize = previousRootScopes.size();
            int cSize = rootScopes.size();
            if(cSize != pSize) {
                logger.error("Unit {} adds subunit {} with initial state but with different root scope count.");
                throw new IllegalStateException("Different root scope count.");
            }

            final BiMap.Transient<S> req = BiMap.Transient.of();
            for(int i = 0; i < rootScopes.size(); i++) {
                final S cRoot = rootScopes.get(i);
                final S pRoot = previousRootScopes.get(i);
                req.put(cRoot, pRoot);
                if(isOwner(cRoot)) {
                    matchedBySharing.put(cRoot, pRoot);
                }
            }

            if(isDifferEnabled()) {
                whenDifferActivated.thenAccept(__ -> {
                    if(!differ.matchScopes(req.freeze())) {
                        logger.error("Unit {} adds subunit {} with initial state but with different root scope count.");
                        throw new IllegalStateException("Could not match.");
                    }
                });
            }

            this.addedUnitIds.__insert(id);
        } else {
            subUnitPreviousResult = null;
        }
        final IFuture<IUnitResult<S, L, D, Result<S, L, D, Q, U>>> result =
                this.<Result<S, L, D, Q, U>>doAddSubUnit(id, (subself, subcontext) -> {
                    return new TypeCheckerUnit<S, L, D, Q, U>(subself, self, subcontext, unitChecker, edgeLabels,
                            changed, subUnitPreviousResult);
                }, rootScopes, false)._2();

        return ifActive(whenContextActive(result));
    }

    @Override public IFuture<IUnitResult<S, L, D, Unit>> add(String id, IScopeGraphLibrary<S, L, D> library,
            List<S> rootScopes) {
        assertActive();

        final IFuture<IUnitResult<S, L, D, Unit>> result = this.<Unit>doAddSubUnit(id, (subself, subcontext) -> {
            IUnitResult<S, L, D, ?> previousResult = this.previousResult != null ?
                    this.previousResult.subUnitResults().get(id) : null;

            return new ScopeGraphLibraryUnit<>(subself, self, subcontext, edgeLabels, library, previousResult);
        }, rootScopes, false)._2();

        this.addedUnitIds.__insert(id);
        return ifActive(whenContextActive(result));
    }

    @Override public void initScope(S root, Iterable<L> labels, boolean sharing) {
        assertActive();

        final List<EdgeOrData<L>> edges = stream(labels).map(EdgeOrData::edge).collect(Collectors.toList());

        doInitShare(self, root, edges, sharing);
    }

    @Override public S freshScope(String baseName, Iterable<L> edgeLabels, boolean data, boolean sharing) {
        assertActive();
        doImplicitActivate();

        final S scope = doFreshScope(baseName, edgeLabels, data, sharing);

        return scope;
    }

    @Override public S stableFreshScope(String name, Iterable<L> edgeLabels, boolean data) {
        assertActive();

        final S scope = doStableFreshScope(name, edgeLabels, data);

        return scope;
    }

    @Override public void shareLocal(S scope) {
        assertActive();

        doAddShare(self, scope);
    }

    @Override public void setDatum(S scope, D datum) {
        assertActive();

        doSetDatum(scope, datum);
        localScopeGraph.set(localScopeGraph.get().setDatum(scope, datum));
    }

    @Override public void addEdge(S source, L label, S target) {
        assertActive();
        doImplicitActivate();

        doAddEdge(self, source, label, target);
        localScopeGraph.set(localScopeGraph.get().addEdge(source, label, target));
    }

    @Override public void closeEdge(S source, L label) {
        assertActive();

        doCloseLabel(self, source, EdgeOrData.edge(label));
    }

    @Override public void closeScope(S scope) {
        assertActive();

        doCloseScope(self, scope);
    }

    @Override public IFuture<? extends java.util.Set<IResolutionPath<S, L, D>>> query(S scope, IQuery<S, L, D> query,
            DataWf<S, L, D> dataWF, DataLeq<S, L, D> dataEquiv, DataWf<S, L, D> dataWfInternal,
            DataLeq<S, L, D> dataEquivInternal) {
        assertActive();
        doImplicitActivate();

        final ScopePath<S, L> path = new ScopePath<>(scope);
        final IFuture<IQueryAnswer<S, L, D>> result =
                doQuery(self, self, true, path, query, dataWF, dataEquiv, dataWfInternal, dataEquivInternal);
        final IFuture<IQueryAnswer<S, L, D>> ret;
        if(result.isDone()) {
            ret = result;
        } else {
            final Query<S, L, D> wf = Query.of(self, path, query, dataWF, dataEquiv, result);
            waitFor(wf, self);
            ret = result.whenComplete((env, ex) -> {
                granted(wf, self);
                resume();
            });
        }
        stats.localQueries += 1;
        return ifActive(ret.thenCompose(ans -> {
            final IFuture<IQueryAnswer<S, L, D>> res = CompletableFuture.completedFuture(ans);
            // TODO: Does the relation q.transitiveQueries().isEmpty() <=> q non-local always hold?
            if(ans.transitiveQueries().isEmpty() && ans.predicateQueries().isEmpty()) {
                return res;
            } else {
                return whenContextActive(res);
            }
        })).thenApply(ans -> {
            return CapsuleUtil.toSet(ans.env());
        });
    }

    @Override public IFuture<? extends java.util.Set<IResolutionPath<S, L, D>>> query(S scope, LabelWf<L> labelWF,
            LabelOrder<L> labelOrder, DataWf<S, L, D> dataWF, DataLeq<S, L, D> dataEquiv,
            @Nullable DataWf<S, L, D> dataWfInternal, @Nullable DataLeq<S, L, D> dataEquivInternal) {
        return query(scope, new NameResolutionQuery<>(labelWF, labelOrder, edgeLabels), dataWF, dataEquiv,
                dataWfInternal, dataEquivInternal);
    }

    @Override public IFuture<? extends java.util.Set<IResolutionPath<S, L, D>>> query(S scope,
            StateMachine<L> stateMachine, LabelWf<L> labelWf, DataWf<S, L, D> dataWF, DataLeq<S, L, D> dataEquiv,
            DataWf<S, L, D> dataWfInternal, DataLeq<S, L, D> dataEquivInternal) {
        return query(scope, new StateMachineQuery<>(stateMachine, labelWf), dataWF, dataEquiv, dataWfInternal,
                dataEquivInternal);
    }

    @Override public <Q> IFuture<R> runIncremental(Function1<Optional<T>, IFuture<Q>> runLocalTypeChecker,
            Function1<R, Q> extractLocal, Function2<Q, IPatchCollection.Immutable<S>, Q> patch,
            Function2<Q, Throwable, IFuture<R>> combine) {
        assertInState(UnitState.INIT_TC);
        state = UnitState.UNKNOWN;
        if(!isIncrementalEnabled() || changed) {
            logger.debug("Unit changed or no previous result was available.");
            stateTransitionTrace = TransitionTrace.INITIALLY_STARTED;
            doRestart(false);
            return runLocalTypeChecker.apply(Optional.empty()).compose(combine::apply);
        }

        // Invariant: added units are marked as changed.
        // Therefore, if unit is not changed, previousResult cannot be null.
        assertPreviousResultProvided();

        if(previousResult.result().localState() != null) {
            doRestore(previousResult.result().localState());
        }

        doConfirmQueries();

        return confirmationResult.thenCompose(patches -> {
            if(patches.isPresent()) {
                final Q previousLocalResult = extractLocal.apply(previousResult.result().analysis());
                return combine.apply(patch.apply(previousLocalResult, patches.get()), null);
            } else {
                final Optional<T> initialState =
                        Optional.ofNullable(previousResult.result().localState()).map(StateCapture::typeCheckerState);
                return runLocalTypeChecker.apply(initialState).compose(combine::apply);
            }
        });
    }

    ///////////////////////////////////////////////////////////////////////////
    // Implementation -- confirmation, restart and reuse
    ///////////////////////////////////////////////////////////////////////////

    private void doConfirmQueries() {
        logger.debug("Confirming queries from previous result.");
        assertInState(UnitState.UNKNOWN);
        assertIncrementalEnabled();
        assertPreviousResultProvided();
        resume();

        if(previousResult.queries().isEmpty()) {
            logger.debug("Releasing - no queries in previous result.");
            doRelease(CapsuleUtil.immutableSet(), CapsuleUtil.immutableSet(), PatchCollection.Immutable.of(),
                    PatchCollection.Immutable.of());
            return;
        }

        final IConfirmation<S, L, D> confirmaton = confirmation.getConfirmation(new ConfirmationContext(self));

        // @formatter:off
        final List<IFuture<SC<Boolean, Boolean>>> futures = previousResult.queries().stream()
                .map(confirmaton::confirm)
                .<IFuture<SC<Boolean, Boolean>>>map(intermediateFuture -> {
                    return intermediateFuture.thenApply(intermediate -> intermediate.match(
                        () -> SC.shortCircuit(false),
                        (addedQueries, removedQueries, resultPatches, globalPatches) -> {
                            // Store intermediate set of patches, to be used on deadlock.
                            this.resultPatches.putAll(resultPatches);
                            this.globalPatches.putAll(globalPatches);
                            this.addedQueries.__insertAll(addedQueries);
                            this.removedQueries.__insertAll(removedQueries);
                            return SC.<Boolean, Boolean>of(true);
                        }
                    ));
                }).collect(Collectors.toList());
        // @formatter:on

        // @formatter:off
        AggregateFuture.ofShortCircuitable(patchCollections -> true, futures).whenComplete((r, ex) -> {
                if(ex == Release.instance) {
                    logger.debug("Confirmation received release.");
                    // Do nothing, unit is already released by deadlock resolution.
                } else if(ex != null) {
                    logger.error("Failure in confirmation.", ex);
                    failures.add(ex);
                } else {
                    if(r) {
                        logger.debug("Queries confirmed - releasing.");
                        this.doRelease(addedQueries.freeze(), removedQueries.freeze(), resultPatches.freeze(), globalPatches.freeze());
                    } else {
                        // No confirmation, hence restart
                        logger.debug("Query confirmation denied - restarting.");
                        if(doRestart(true)) {
                            stateTransitionTrace = TransitionTrace.RESTARTED;
                        }
                    }
                }
            });
        // @formatter:on
    }

    private void doRelease(Set.Immutable<IRecordedQuery<S, L, D>> addedQueries,
            Set.Immutable<IRecordedQuery<S, L, D>> removedQueries, IPatchCollection.Immutable<S> resultPatches,
            IPatchCollection.Immutable<S> globalPatches) {
        assertIncrementalEnabled();
        if(state == UnitState.UNKNOWN) {
            logger.debug("{} releasing.", this);
            logger.trace("Patches: result: {}; global: {}.", resultPatches, globalPatches);
            assertPreviousResultProvided();

            // TODO When a unit has active subunits, the matching differ cannot be used, because these units may add/remove edges.
            // However, a full-blown scope graph differ may impose quite some overhead. We should invent a smarter solution for this.
            // - Perhaps only maintain local scope graphs (i.e. no edge sharing), and relay queries to subunits
            //   - Con: very similar to epsilon edges, which degraded performance a lot, and complicated shadowing
            // - Somehow inform differ of local edges and subunit edges, and perform diff based on those.

            // Copy scopes
            // No need to patch, because they can only be local, and local state is reused.
            this.scopes.__insertAll(previousResult.scopes());

            final IPatchCollection.Immutable<S> localPatches = matchedBySharing.freeze();
            // @formatter:off
            final IScopeGraphDiffer<S, L, D> differ = localPatches.isIdentity() && this.addedUnitIds.isEmpty() ?
                new MatchingDiffer<S, L, D>(differOps(), differContext(typeChecker::internalData), resultPatches.allPatches()) :
                new ScopeGraphDiffer<>(differContext(typeChecker::internalData), new StaticDifferContext<>(previousResult.scopeGraph(),
                        previousResult.scopes(), new DifferDataOps()), differOps(), edgeLabels);
            // @formatter:on

            final Collection<S> openScopes =
                    this.addedUnitIds.isEmpty() ? Collections.emptySet() : previousResult.rootScopes();
            initDiffer(differ, previousResult.scopeGraph(), previousResult.scopes(),
                    previousResult.result().sharedScopes(), localPatches, openScopes, ImmutableMultimap.of());
            if(isConfirmationEnabled()) {
                this.envDiffer = new IndexedEnvDiffer<>(new EnvDiffer<>(envDifferContext, differOps()));
            }

            logger.debug("Rebuilding scope graph.");
            // @formatter:off
            final Patcher<S, L, D> patcher = new Patcher.Builder<S, L, D>()
                .patchSources(localPatches)
                .patchEdgeTargets(resultPatches)
                .patchDatumSources(localPatches)
                .patchDatums(resultPatches, context::substituteScopes)
                .build();

            final IScopeGraph.Immutable<S, L, D> patchedLocalScopeGraph = patcher.<Boolean>apply(
                previousResult.result().scopeGraph(),
                (oldSource, newSource) -> isOwner(newSource),
                (oldSource, newSource, lbl, oldTarget, newTarget, sourceLocal) -> {
                    if(!sourceLocal) {
                        doAddEdge(self, newSource, lbl, newTarget);
                    }
                },
                Patcher.DataPatchCallback.noop()
            );

            scopeGraph.set(scopeGraph.get().addAll(patchedLocalScopeGraph));
            localScopeGraph.set(localScopeGraph.get().addAll(patchedLocalScopeGraph));

            // initialize all scopes that are pending, and close all open labels.
            // these should be set by the now reused scopegraph.
            logger.debug("Close pending tokens.");
            // @formatter:off
            final IWaitFor.Cases<S, L, D> cases = IWaitFor.cases(
                initScope -> doInitShare(self, initScope.scope(), CapsuleUtil.immutableSet(), false),
                closeScope -> {},
                closeLabel -> doCloseLabel(self, closeLabel.scope(), closeLabel.label()),
                query -> {},
                pQuery -> {},
                confirm -> {},
                match -> {},
                result -> {},
                typeCheckerState -> {},
                differResult -> {},
                differState -> {},
                envDifferState -> {},
                unitAdd -> {}
            );
            // @formatter:on
            for(IWaitFor<S, L, D> wf: ownTokens()) {
                wf.visit(cases);
            }

            pendingExternalDatums.asMap().forEach((d, futures) -> futures.elementSet().forEach(future -> {
                final D datum = previousResult.result().analysis().getExternalRepresentation(d);
                final D result = context.substituteScopes(datum, resultPatches.patches().invert());
                self.complete(future, result, null);
            }));

            final IPatchCollection.Immutable<S> allPatches = resultPatches.putAll(globalPatches);
            java.util.Set<IRecordedQuery<S, L, D>> newRecordedQueries =
                    previousResult.queries().stream().map(q -> q.patch(allPatches)).collect(Collectors.toSet());
            java.util.Set<IRecordedQuery<S, L, D>> patchedRemovedQueries =
                    removedQueries.stream().map(q -> q.patch(allPatches)).collect(Collectors.toSet());
            newRecordedQueries.removeAll(patchedRemovedQueries);
            newRecordedQueries.addAll(addedQueries);
            recordedQueries.addAll(newRecordedQueries);

            stateTransitionTrace = TransitionTrace.RELEASED;

            self.complete(confirmationResult, Optional.of(resultPatches), null);
            // Cancel all futures waiting for activation
            self.complete(whenActive, null, Release.instance);
            state = UnitState.RELEASED;

            resume();
            tryFinish(); // FIXME needed?
            logger.debug("{} released.", this);
        }
    }

    private boolean doRestart(boolean external) {
        if(state == UnitState.INIT_TC || state == UnitState.UNKNOWN) {
            logger.debug("{} restarting.", this);
            state = UnitState.ACTIVE;
            self.complete(confirmationResult, Optional.empty(), null);
            self.complete(whenActive, Unit.unit, null);
            // Not activating context here, as `doRestart` is also used before non-incremental runs.
            if(isDifferEnabled()) {
                final IDifferContext<S, L, D> context = differContext(typeChecker::internalData);
                final IDifferOps<S, L, D> differOps = differOps();
                if(previousResult != null) {
                    final IDifferContext<S, L, D> differContext = new StaticDifferContext<>(previousResult.scopeGraph(),
                            previousResult.scopes(), new DifferDataOps());

                    if(external) {
                        final StateCapture<S, L, D, T> capture = previousResult.result().localState();
                        final java.util.Set<S> openScopes = Sets.union(capture.openScopes().elementSet(),
                                capture.unInitializedScopes().elementSet());

                        initDiffer(new ScopeGraphDiffer<>(context, differContext, differOps, edgeLabels),
                                capture.scopeGraph(), capture.scopes(), previousResult.result().sharedScopes(),
                                matchedBySharing.freeze(), openScopes, capture.openEdges());
                    } else {
                        initDiffer(new ScopeGraphDiffer<>(context, differContext, differOps, edgeLabels),
                                this.rootScopes, previousResult.rootScopes());
                    }
                } else {
                    initDiffer(new AddingDiffer<>(context, differOps, edgeLabels), Collections.emptyList(),
                            Collections.emptyList());
                }
                if(isConfirmationEnabled()) {
                    this.envDiffer = new IndexedEnvDiffer<>(new EnvDiffer<>(envDifferContext, differOps()));
                }
            }
            pendingExternalDatums.asMap().forEach((d, futures) -> {
                typeChecker.getExternalDatum(d).whenComplete((d2, ex) -> futures.forEach(future -> {
                    future.complete(d2, ex);
                }));
            });

            resume();
            tryFinish(); // FIXME needed?
            logger.debug("{} restarted.", this);
            return true;
        }
        return false;
    }

    private void doImplicitActivate() {
        // If invoked synchronously, do a implicit transition to ACTIVE<commit_msg>Init env differ before SG differ<commit_after>package mb.p_raffrayi.impl;

import static com.google.common.collect.Streams.stream;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.stream.Collectors;

import javax.annotation.Nullable;

import org.metaborg.util.Ref;
import org.metaborg.util.collection.CapsuleUtil;
import org.metaborg.util.collection.MultiSet;
import org.metaborg.util.collection.MultiSetMap;
import org.metaborg.util.functions.Function1;
import org.metaborg.util.functions.Function2;
import org.metaborg.util.future.AggregateFuture;
import org.metaborg.util.future.AggregateFuture.SC;
import org.metaborg.util.future.CompletableFuture;
import org.metaborg.util.future.ICompletableFuture;
import org.metaborg.util.future.IFuture;
import org.metaborg.util.iterators.Iterables2;
import org.metaborg.util.log.ILogger;
import org.metaborg.util.log.LoggerUtils;
import org.metaborg.util.tuple.Tuple2;
import org.metaborg.util.unit.Unit;

import com.google.common.collect.ImmutableMultimap;
import com.google.common.collect.Iterables;
import com.google.common.collect.Sets;

import io.usethesource.capsule.Set;
import mb.p_raffrayi.IIncrementalTypeCheckerContext;
import mb.p_raffrayi.IRecordedQuery;
import mb.p_raffrayi.IScopeGraphLibrary;
import mb.p_raffrayi.ITypeChecker;
import mb.p_raffrayi.ITypeChecker.IOutput;
import mb.p_raffrayi.ITypeChecker.IState;
import mb.p_raffrayi.IUnitResult;
import mb.p_raffrayi.IUnitResult.TransitionTrace;
import mb.p_raffrayi.actors.IActor;
import mb.p_raffrayi.actors.IActorRef;
import mb.p_raffrayi.impl.confirm.ConfirmResult;
import mb.p_raffrayi.impl.confirm.IConfirmation;
import mb.p_raffrayi.impl.confirm.IConfirmationContext;
import mb.p_raffrayi.impl.confirm.IConfirmationFactory;
import mb.p_raffrayi.impl.confirm.OptimisticConfirmation;
import mb.p_raffrayi.impl.diff.AddingDiffer;
import mb.p_raffrayi.impl.diff.IDifferContext;
import mb.p_raffrayi.impl.diff.IDifferDataOps;
import mb.p_raffrayi.impl.diff.IDifferOps;
import mb.p_raffrayi.impl.diff.IScopeGraphDiffer;
import mb.p_raffrayi.impl.diff.MatchingDiffer;
import mb.p_raffrayi.impl.diff.ScopeDiff;
import mb.p_raffrayi.impl.diff.ScopeGraphDiffer;
import mb.p_raffrayi.impl.diff.StaticDifferContext;
import mb.p_raffrayi.impl.envdiff.EnvDiffer;
import mb.p_raffrayi.impl.envdiff.IEnvDiff;
import mb.p_raffrayi.impl.envdiff.IEnvDiffer;
import mb.p_raffrayi.impl.envdiff.IEnvDifferContext;
import mb.p_raffrayi.impl.envdiff.IndexedEnvDiffer;
import mb.p_raffrayi.impl.tokens.CloseLabel;
import mb.p_raffrayi.impl.tokens.CloseScope;
import mb.p_raffrayi.impl.tokens.Confirm;
import mb.p_raffrayi.impl.tokens.DifferState;
import mb.p_raffrayi.impl.tokens.EnvDifferState;
import mb.p_raffrayi.impl.tokens.IWaitFor;
import mb.p_raffrayi.impl.tokens.InitScope;
import mb.p_raffrayi.impl.tokens.Query;
import mb.p_raffrayi.nameresolution.DataLeq;
import mb.p_raffrayi.nameresolution.DataWf;
import mb.p_raffrayi.nameresolution.IQuery;
import mb.p_raffrayi.nameresolution.NameResolutionQuery;
import mb.p_raffrayi.nameresolution.StateMachineQuery;
import mb.scopegraph.ecoop21.LabelOrder;
import mb.scopegraph.ecoop21.LabelWf;
import mb.scopegraph.oopsla20.IScopeGraph;
import mb.scopegraph.oopsla20.diff.BiMap;
import mb.scopegraph.oopsla20.path.IResolutionPath;
import mb.scopegraph.oopsla20.reference.EdgeOrData;
import mb.scopegraph.oopsla20.reference.ScopeGraph;
import mb.scopegraph.oopsla20.terms.newPath.ScopePath;
import mb.scopegraph.patching.IPatchCollection;
import mb.scopegraph.patching.PatchCollection;
import mb.scopegraph.patching.Patcher;
import mb.scopegraph.resolution.StateMachine;

class TypeCheckerUnit<S, L, D, R extends IOutput<S, L, D>, T extends IState<S, L, D>>
        extends AbstractUnit<S, L, D, Result<S, L, D, R, T>> implements IIncrementalTypeCheckerContext<S, L, D, R, T> {


    private static final ILogger logger = LoggerUtils.logger(TypeCheckerUnit.class);

    private final ITypeChecker<S, L, D, R, T> typeChecker;
    private final boolean changed;
    private final @Nullable IUnitResult<S, L, D, Result<S, L, D, R, T>> previousResult;

    private volatile UnitState state;

    private final IPatchCollection.Transient<S> matchedBySharing = PatchCollection.Transient.of();

    private final Ref<IScopeGraph.Immutable<S, L, D>> localScopeGraph = new Ref<>(ScopeGraph.Immutable.of());

    private final Set.Transient<String> addedUnitIds = CapsuleUtil.transientSet();

    private final Ref<StateCapture<S, L, D, T>> localCapture = new Ref<>();
    private final ICompletableFuture<Unit> whenActive = new CompletableFuture<>(); // activated when transitioning from unknown to active/released.
    private final ICompletableFuture<Unit> whenContextActivated = new CompletableFuture<>(); // activated when transitioning from local/unknown to active.

    private IEnvDiffer<S, L, D> envDiffer;

    private final IPatchCollection.Transient<S> externalMatches = PatchCollection.Transient.of();

    private final IConfirmationFactory<S, L, D> confirmation;
    private final ICompletableFuture<Optional<IPatchCollection.Immutable<S>>> confirmationResult =
            new CompletableFuture<>();

    // Intermediate confirmation result aggregations
    private final IPatchCollection.Transient<S> resultPatches = PatchCollection.Transient.of();
    private final IPatchCollection.Transient<S> globalPatches = PatchCollection.Transient.of();
    private final Set.Transient<IRecordedQuery<S, L, D>> addedQueries = CapsuleUtil.transientSet();
    private final Set.Transient<IRecordedQuery<S, L, D>> removedQueries = CapsuleUtil.transientSet();

    TypeCheckerUnit(IActor<? extends IUnit<S, L, D, Result<S, L, D, R, T>>> self,
            @Nullable IActorRef<? extends IUnit<S, L, D, ?>> parent, IUnitContext<S, L, D> context,
            ITypeChecker<S, L, D, R, T> unitChecker, Iterable<L> edgeLabels, boolean inputChanged,
            IUnitResult<S, L, D, Result<S, L, D, R, T>> previousResult) {
        super(self, parent, context, edgeLabels);
        this.typeChecker = unitChecker;
        this.changed = inputChanged;
        this.previousResult = previousResult;
        this.state = UnitState.INIT_UNIT;
        this.confirmation = OptimisticConfirmation.factory();
    }

    TypeCheckerUnit(IActor<? extends IUnit<S, L, D, Result<S, L, D, R, T>>> self,
            @Nullable IActorRef<? extends IUnit<S, L, D, ?>> parent, IUnitContext<S, L, D> context,
            ITypeChecker<S, L, D, R, T> unitChecker, Iterable<L> edgeLabels) {
        this(self, parent, context, unitChecker, edgeLabels, true, null);
    }

    private final MultiSetMap.Transient<D, ICompletableFuture<D>> pendingExternalDatums = MultiSetMap.Transient.of();

    @Override protected IFuture<D> getExternalDatum(D datum) {
        if(!changed && previousResult != null && previousResult.result().localState() != null) {
            // when previous result is present, reliable (i.e., !changed), and has local state
            // try to find datum from previous state
            final Optional<D> datumOpt =
                    previousResult.result().localState().typeCheckerState().tryGetExternalDatum(datum);
            if(datumOpt.isPresent()) {
                return CompletableFuture.completedFuture(datumOpt.get());
            }
        }
        return whenActive.compose((u, ex) -> {
            if(this.state.equals(UnitState.RELEASED) || this.state.equals(UnitState.DONE)
                    && this.stateTransitionTrace.equals(TransitionTrace.RELEASED)) {
                final D result = previousResult.result().analysis().getExternalRepresentation(datum);
                return CompletableFuture.completedFuture(context.substituteScopes(result, resultPatches.patches().invert()));
            }
            final IFuture<D> future = typeChecker.getExternalDatum(datum);
            if(future.isDone()) {
                return future;
            }
            final ICompletableFuture<D> result = new CompletableFuture<>();
            pendingExternalDatums.put(datum, result);
            future.whenComplete(result::complete);
            return result.whenComplete((__, ex2) -> {
                pendingExternalDatums.remove(datum, result);
            });
        });
    }

    @Override protected D getPreviousDatum(D datum) {
        assertPreviousResultProvided();
        return previousResult.result().analysis().getExternalRepresentation(datum);
    }

    ///////////////////////////////////////////////////////////////////////////
    // IBroker2UnitProtocol interface, called by IBroker implementations
    ///////////////////////////////////////////////////////////////////////////

    @Override public IFuture<IUnitResult<S, L, D, Result<S, L, D, R, T>>> _start(List<S> rootScopes) {
        assertInState(UnitState.INIT_UNIT);
        resume();

        doStart(rootScopes);
        if(isIncrementalEnabled() && previousResult != null) {
            if(previousResult.rootScopes().size() != rootScopes.size()) {
                throw new IllegalStateException("Root scope counts do not match.");
            }
            Iterables2.zip(rootScopes, previousResult.rootScopes(), Tuple2::of).forEach(match -> {
                matchedBySharing.put(match._1(), match._2());
                externalMatches.put(match._1(), match._2());
            });
        }

        final IFuture<Result<S, L, D, R, T>> result;
        try {
            state = UnitState.INIT_TC;
            result = this.typeChecker.run(this, rootScopes).whenComplete((r, ex) -> {
                if(state == UnitState.INIT_TC) {
                    stateTransitionTrace = TransitionTrace.INITIALLY_STARTED;
                }
                if(inLocalPhase()) {
                    doCapture();
                }
                state = UnitState.DONE;
                resume();
                tryFinish();
            }).thenApply(r -> {
                return Result.<S, L, D, R, T>of(r, localCapture.get(), localScopeGraph.get(),
                        matchedBySharing.patchRange());
            });
        } catch(Exception e) {
            logger.error("Exception starting type-checker {}.", e);
            return doFinish(CompletableFuture.completedExceptionally(e));
        }

        // Start phantom units for all units that have not yet been restarted
        if(isIncrementalEnabled() && previousResult != null) {
            for(String removedId : Sets.difference(previousResult.subUnitResults().keySet(), addedUnitIds)) {
                @SuppressWarnings("unchecked") final IUnitResult<S, L, D, ? extends IOutput<S, L, D>> subResult =
                        (IUnitResult<S, L, D, ? extends IOutput<S, L, D>>) previousResult.subUnitResults()
                                .get(removedId);
                this.<Unit>doAddSubUnit(removedId,
                        (subself, subcontext) -> new PhantomUnit<>(subself, self, subcontext, edgeLabels, subResult),
                        new ArrayList<>(), true);
            }
        }

        if(state == UnitState.INIT_TC) {
            // runIncremental not called, so start eagerly
            doRestart(false);
        } else if(state == UnitState.DONE) {
            // Completed synchronously
            whenActive.complete(Unit.unit);
            confirmationResult.complete(Optional.of(PatchCollection.Immutable.of()));
        }

        return doFinish(result);
    }

    ///////////////////////////////////////////////////////////////////////////
    // IUnit2UnitProtocol interface, called by IUnit implementations
    ///////////////////////////////////////////////////////////////////////////

    @Override public IFuture<ConfirmResult<S, L, D>> _confirm(ScopePath<S, L> path, LabelWf<L> labelWF,
            DataWf<S, L, D> dataWF, boolean prevEnvEmpty) {
        assertConfirmationEnabled();
        stats.incomingConfirmations++;
        final IActorRef<? extends IUnit<S, L, D, ?>> sender = self.sender(TYPE);
        final ICompletableFuture<ConfirmResult<S, L, D>> result = new CompletableFuture<>();
        whenActive.whenComplete((__, ex) -> {
            if(ex != null && ex != Release.instance) {
                result.completeExceptionally(ex);
            } else {
                // Cannot immediately confirm on `Release.instance` exception
                // because subunits might change edges in shared scopes.
                confirmation.getConfirmation(new ConfirmationContext(sender))
                        .confirm(path, labelWF, dataWF, prevEnvEmpty).whenComplete(result::complete);
            }
        });
        return result;
    }

    @Override public IFuture<IQueryAnswer<S, L, D>> _queryPrevious(ScopePath<S, L> path, IQuery<S, L, D> query,
            DataWf<S, L, D> dataWF, DataLeq<S, L, D> dataEquiv) {
        assertPreviousResultProvided();
        return doQueryPrevious(self.sender(TYPE), previousResult.scopeGraph(), path, query, dataWF, dataEquiv);
    }

    @Override public IFuture<Optional<S>> _match(S previousScope) {
        assertOwnScope(previousScope);
        if(matchedBySharing.identityPatches().contains(previousScope)) {
            return CompletableFuture.completedFuture(Optional.of(previousScope));
        }
        if(matchedBySharing.patchDomain().contains(previousScope)) {
            return CompletableFuture.completedFuture(Optional.of(matchedBySharing.patch(previousScope)));
        }
        if(!changed && previousResult.result().localState() != null
                && previousResult.result().localState().scopes().contains(previousScope)) {
            return CompletableFuture.completedFuture(Optional.of(previousScope));
        }
        return super._match(previousScope);
    }

    @Override public IFuture<StateSummary<S, L, D>> _state() {
        // resume();
        if(state.equals(UnitState.ACTIVE)
                || (state == UnitState.DONE && stateTransitionTrace != TransitionTrace.RELEASED)) {
            return CompletableFuture.completedFuture(StateSummary.restart(process, dependentSet()));
        }

        if(state.equals(UnitState.RELEASED)
                || (state == UnitState.DONE && stateTransitionTrace == TransitionTrace.RELEASED)) {
            return CompletableFuture.completedFuture(StateSummary.released(process, dependentSet()));
        }

        return CompletableFuture.completedFuture(StateSummary.release(process, dependentSet()));
    }

    @Override public void _restart() {
        assertIncrementalEnabled();
        if(doRestart(true)) {
            stateTransitionTrace = TransitionTrace.RESTARTED;
        }
    }

    @Override public void _release() {
        if(state == UnitState.UNKNOWN) {
            doRelease(CapsuleUtil.toSet(this.addedQueries), CapsuleUtil.toSet(this.removedQueries),
                    PatchCollection.Immutable.<S>of().putAll(this.resultPatches),
                    PatchCollection.Immutable.<S>of().putAll(globalPatches));
        }
    }

    ///////////////////////////////////////////////////////////////////////////
    // ITypeCheckerContext interface, called by ITypeChecker implementations
    ///////////////////////////////////////////////////////////////////////////

    // NB. Invoke methods via `local` so that we have the same scheduling & ordering
    // guarantees as for remote calls.

    @Override public String id() {
        return self.id();
    }

    @SuppressWarnings("unchecked") @Override public <Q extends IOutput<S, L, D>, U extends IState<S, L, D>>
            IFuture<IUnitResult<S, L, D, Result<S, L, D, Q, U>>>
            add(String id, ITypeChecker<S, L, D, Q, U> unitChecker, List<S> rootScopes, boolean changed) {
        assertActive();

        // No previous result for subunit
        if(this.previousResult == null || !this.previousResult.subUnitResults().containsKey(id)) {
            this.addedUnitIds.__insert(id);

            return ifActive(whenContextActive(this.<Result<S, L, D, Q, U>>doAddSubUnit(id, (subself, subcontext) -> {
                return new TypeCheckerUnit<>(subself, self, subcontext, unitChecker, edgeLabels);
            }, rootScopes, false)._2()));
        }

        final IUnitResult<S, L, D, Result<S, L, D, Q, U>> subUnitPreviousResult;
        if(isIncrementalEnabled()) {

            subUnitPreviousResult =
                    (IUnitResult<S, L, D, Result<S, L, D, Q, U>>) this.previousResult.subUnitResults().get(id);


            // When a scope is shared, the shares must be consistent.
            // Also, it is not necessary that shared scopes are reachable from the root scopes
            // (A unit started by the Broker does not even have root scopes)
            // Therefore we enforce here that the current root scopes and the previous ones match.
            List<S> previousRootScopes = subUnitPreviousResult.rootScopes();

            int pSize = previousRootScopes.size();
            int cSize = rootScopes.size();
            if(cSize != pSize) {
                logger.error("Unit {} adds subunit {} with initial state but with different root scope count.");
                throw new IllegalStateException("Different root scope count.");
            }

            final BiMap.Transient<S> req = BiMap.Transient.of();
            for(int i = 0; i < rootScopes.size(); i++) {
                final S cRoot = rootScopes.get(i);
                final S pRoot = previousRootScopes.get(i);
                req.put(cRoot, pRoot);
                if(isOwner(cRoot)) {
                    matchedBySharing.put(cRoot, pRoot);
                }
            }

            if(isDifferEnabled()) {
                whenDifferActivated.thenAccept(__ -> {
                    if(!differ.matchScopes(req.freeze())) {
                        logger.error("Unit {} adds subunit {} with initial state but with different root scope count.");
                        throw new IllegalStateException("Could not match.");
                    }
                });
            }

            this.addedUnitIds.__insert(id);
        } else {
            subUnitPreviousResult = null;
        }
        final IFuture<IUnitResult<S, L, D, Result<S, L, D, Q, U>>> result =
                this.<Result<S, L, D, Q, U>>doAddSubUnit(id, (subself, subcontext) -> {
                    return new TypeCheckerUnit<S, L, D, Q, U>(subself, self, subcontext, unitChecker, edgeLabels,
                            changed, subUnitPreviousResult);
                }, rootScopes, false)._2();

        return ifActive(whenContextActive(result));
    }

    @Override public IFuture<IUnitResult<S, L, D, Unit>> add(String id, IScopeGraphLibrary<S, L, D> library,
            List<S> rootScopes) {
        assertActive();

        final IFuture<IUnitResult<S, L, D, Unit>> result = this.<Unit>doAddSubUnit(id, (subself, subcontext) -> {
            IUnitResult<S, L, D, ?> previousResult = this.previousResult != null ?
                    this.previousResult.subUnitResults().get(id) : null;

            return new ScopeGraphLibraryUnit<>(subself, self, subcontext, edgeLabels, library, previousResult);
        }, rootScopes, false)._2();

        this.addedUnitIds.__insert(id);
        return ifActive(whenContextActive(result));
    }

    @Override public void initScope(S root, Iterable<L> labels, boolean sharing) {
        assertActive();

        final List<EdgeOrData<L>> edges = stream(labels).map(EdgeOrData::edge).collect(Collectors.toList());

        doInitShare(self, root, edges, sharing);
    }

    @Override public S freshScope(String baseName, Iterable<L> edgeLabels, boolean data, boolean sharing) {
        assertActive();
        doImplicitActivate();

        final S scope = doFreshScope(baseName, edgeLabels, data, sharing);

        return scope;
    }

    @Override public S stableFreshScope(String name, Iterable<L> edgeLabels, boolean data) {
        assertActive();

        final S scope = doStableFreshScope(name, edgeLabels, data);

        return scope;
    }

    @Override public void shareLocal(S scope) {
        assertActive();

        doAddShare(self, scope);
    }

    @Override public void setDatum(S scope, D datum) {
        assertActive();

        doSetDatum(scope, datum);
        localScopeGraph.set(localScopeGraph.get().setDatum(scope, datum));
    }

    @Override public void addEdge(S source, L label, S target) {
        assertActive();
        doImplicitActivate();

        doAddEdge(self, source, label, target);
        localScopeGraph.set(localScopeGraph.get().addEdge(source, label, target));
    }

    @Override public void closeEdge(S source, L label) {
        assertActive();

        doCloseLabel(self, source, EdgeOrData.edge(label));
    }

    @Override public void closeScope(S scope) {
        assertActive();

        doCloseScope(self, scope);
    }

    @Override public IFuture<? extends java.util.Set<IResolutionPath<S, L, D>>> query(S scope, IQuery<S, L, D> query,
            DataWf<S, L, D> dataWF, DataLeq<S, L, D> dataEquiv, DataWf<S, L, D> dataWfInternal,
            DataLeq<S, L, D> dataEquivInternal) {
        assertActive();
        doImplicitActivate();

        final ScopePath<S, L> path = new ScopePath<>(scope);
        final IFuture<IQueryAnswer<S, L, D>> result =
                doQuery(self, self, true, path, query, dataWF, dataEquiv, dataWfInternal, dataEquivInternal);
        final IFuture<IQueryAnswer<S, L, D>> ret;
        if(result.isDone()) {
            ret = result;
        } else {
            final Query<S, L, D> wf = Query.of(self, path, query, dataWF, dataEquiv, result);
            waitFor(wf, self);
            ret = result.whenComplete((env, ex) -> {
                granted(wf, self);
                resume();
            });
        }
        stats.localQueries += 1;
        return ifActive(ret.thenCompose(ans -> {
            final IFuture<IQueryAnswer<S, L, D>> res = CompletableFuture.completedFuture(ans);
            // TODO: Does the relation q.transitiveQueries().isEmpty() <=> q non-local always hold?
            if(ans.transitiveQueries().isEmpty() && ans.predicateQueries().isEmpty()) {
                return res;
            } else {
                return whenContextActive(res);
            }
        })).thenApply(ans -> {
            return CapsuleUtil.toSet(ans.env());
        });
    }

    @Override public IFuture<? extends java.util.Set<IResolutionPath<S, L, D>>> query(S scope, LabelWf<L> labelWF,
            LabelOrder<L> labelOrder, DataWf<S, L, D> dataWF, DataLeq<S, L, D> dataEquiv,
            @Nullable DataWf<S, L, D> dataWfInternal, @Nullable DataLeq<S, L, D> dataEquivInternal) {
        return query(scope, new NameResolutionQuery<>(labelWF, labelOrder, edgeLabels), dataWF, dataEquiv,
                dataWfInternal, dataEquivInternal);
    }

    @Override public IFuture<? extends java.util.Set<IResolutionPath<S, L, D>>> query(S scope,
            StateMachine<L> stateMachine, LabelWf<L> labelWf, DataWf<S, L, D> dataWF, DataLeq<S, L, D> dataEquiv,
            DataWf<S, L, D> dataWfInternal, DataLeq<S, L, D> dataEquivInternal) {
        return query(scope, new StateMachineQuery<>(stateMachine, labelWf), dataWF, dataEquiv, dataWfInternal,
                dataEquivInternal);
    }

    @Override public <Q> IFuture<R> runIncremental(Function1<Optional<T>, IFuture<Q>> runLocalTypeChecker,
            Function1<R, Q> extractLocal, Function2<Q, IPatchCollection.Immutable<S>, Q> patch,
            Function2<Q, Throwable, IFuture<R>> combine) {
        assertInState(UnitState.INIT_TC);
        state = UnitState.UNKNOWN;
        if(!isIncrementalEnabled() || changed) {
            logger.debug("Unit changed or no previous result was available.");
            stateTransitionTrace = TransitionTrace.INITIALLY_STARTED;
            doRestart(false);
            return runLocalTypeChecker.apply(Optional.empty()).compose(combine::apply);
        }

        // Invariant: added units are marked as changed.
        // Therefore, if unit is not changed, previousResult cannot be null.
        assertPreviousResultProvided();

        if(previousResult.result().localState() != null) {
            doRestore(previousResult.result().localState());
        }

        doConfirmQueries();

        return confirmationResult.thenCompose(patches -> {
            if(patches.isPresent()) {
                final Q previousLocalResult = extractLocal.apply(previousResult.result().analysis());
                return combine.apply(patch.apply(previousLocalResult, patches.get()), null);
            } else {
                final Optional<T> initialState =
                        Optional.ofNullable(previousResult.result().localState()).map(StateCapture::typeCheckerState);
                return runLocalTypeChecker.apply(initialState).compose(combine::apply);
            }
        });
    }

    ///////////////////////////////////////////////////////////////////////////
    // Implementation -- confirmation, restart and reuse
    ///////////////////////////////////////////////////////////////////////////

    private void doConfirmQueries() {
        logger.debug("Confirming queries from previous result.");
        assertInState(UnitState.UNKNOWN);
        assertIncrementalEnabled();
        assertPreviousResultProvided();
        resume();

        if(previousResult.queries().isEmpty()) {
            logger.debug("Releasing - no queries in previous result.");
            doRelease(CapsuleUtil.immutableSet(), CapsuleUtil.immutableSet(), PatchCollection.Immutable.of(),
                    PatchCollection.Immutable.of());
            return;
        }

        final IConfirmation<S, L, D> confirmaton = confirmation.getConfirmation(new ConfirmationContext(self));

        // @formatter:off
        final List<IFuture<SC<Boolean, Boolean>>> futures = previousResult.queries().stream()
                .map(confirmaton::confirm)
                .<IFuture<SC<Boolean, Boolean>>>map(intermediateFuture -> {
                    return intermediateFuture.thenApply(intermediate -> intermediate.match(
                        () -> SC.shortCircuit(false),
                        (addedQueries, removedQueries, resultPatches, globalPatches) -> {
                            // Store intermediate set of patches, to be used on deadlock.
                            this.resultPatches.putAll(resultPatches);
                            this.globalPatches.putAll(globalPatches);
                            this.addedQueries.__insertAll(addedQueries);
                            this.removedQueries.__insertAll(removedQueries);
                            return SC.<Boolean, Boolean>of(true);
                        }
                    ));
                }).collect(Collectors.toList());
        // @formatter:on

        // @formatter:off
        AggregateFuture.ofShortCircuitable(patchCollections -> true, futures).whenComplete((r, ex) -> {
                if(ex == Release.instance) {
                    logger.debug("Confirmation received release.");
                    // Do nothing, unit is already released by deadlock resolution.
                } else if(ex != null) {
                    logger.error("Failure in confirmation.", ex);
                    failures.add(ex);
                } else {
                    if(r) {
                        logger.debug("Queries confirmed - releasing.");
                        this.doRelease(addedQueries.freeze(), removedQueries.freeze(), resultPatches.freeze(), globalPatches.freeze());
                    } else {
                        // No confirmation, hence restart
                        logger.debug("Query confirmation denied - restarting.");
                        if(doRestart(true)) {
                            stateTransitionTrace = TransitionTrace.RESTARTED;
                        }
                    }
                }
            });
        // @formatter:on
    }

    private void doRelease(Set.Immutable<IRecordedQuery<S, L, D>> addedQueries,
            Set.Immutable<IRecordedQuery<S, L, D>> removedQueries, IPatchCollection.Immutable<S> resultPatches,
            IPatchCollection.Immutable<S> globalPatches) {
        assertIncrementalEnabled();
        if(state == UnitState.UNKNOWN) {
            logger.debug("{} releasing.", this);
            logger.trace("Patches: result: {}; global: {}.", resultPatches, globalPatches);
            assertPreviousResultProvided();

            // TODO When a unit has active subunits, the matching differ cannot be used, because these units may add/remove edges.
            // However, a full-blown scope graph differ may impose quite some overhead. We should invent a smarter solution for this.
            // - Perhaps only maintain local scope graphs (i.e. no edge sharing), and relay queries to subunits
            //   - Con: very similar to epsilon edges, which degraded performance a lot, and complicated shadowing
            // - Somehow inform differ of local edges and subunit edges, and perform diff based on those.

            // Copy scopes
            // No need to patch, because they can only be local, and local state is reused.
            this.scopes.__insertAll(previousResult.scopes());

            final IPatchCollection.Immutable<S> localPatches = matchedBySharing.freeze();
            // @formatter:off
            final IScopeGraphDiffer<S, L, D> differ = localPatches.isIdentity() && this.addedUnitIds.isEmpty() ?
                new MatchingDiffer<S, L, D>(differOps(), differContext(typeChecker::internalData), resultPatches.allPatches()) :
                new ScopeGraphDiffer<>(differContext(typeChecker::internalData), new StaticDifferContext<>(previousResult.scopeGraph(),
                        previousResult.scopes(), new DifferDataOps()), differOps(), edgeLabels);
            // @formatter:on

            final Collection<S> openScopes =
                    this.addedUnitIds.isEmpty() ? Collections.emptySet() : previousResult.rootScopes();
            if(isConfirmationEnabled()) {
                this.envDiffer = new IndexedEnvDiffer<>(new EnvDiffer<>(envDifferContext, differOps()));
            }
            initDiffer(differ, previousResult.scopeGraph(), previousResult.scopes(),
                    previousResult.result().sharedScopes(), localPatches, openScopes, ImmutableMultimap.of());

            logger.debug("Rebuilding scope graph.");
            // @formatter:off
            final Patcher<S, L, D> patcher = new Patcher.Builder<S, L, D>()
                .patchSources(localPatches)
                .patchEdgeTargets(resultPatches)
                .patchDatumSources(localPatches)
                .patchDatums(resultPatches, context::substituteScopes)
                .build();

            final IScopeGraph.Immutable<S, L, D> patchedLocalScopeGraph = patcher.<Boolean>apply(
                previousResult.result().scopeGraph(),
                (oldSource, newSource) -> isOwner(newSource),
                (oldSource, newSource, lbl, oldTarget, newTarget, sourceLocal) -> {
                    if(!sourceLocal) {
                        doAddEdge(self, newSource, lbl, newTarget);
                    }
                },
                Patcher.DataPatchCallback.noop()
            );

            scopeGraph.set(scopeGraph.get().addAll(patchedLocalScopeGraph));
            localScopeGraph.set(localScopeGraph.get().addAll(patchedLocalScopeGraph));

            // initialize all scopes that are pending, and close all open labels.
            // these should be set by the now reused scopegraph.
            logger.debug("Close pending tokens.");
            // @formatter:off
            final IWaitFor.Cases<S, L, D> cases = IWaitFor.cases(
                initScope -> doInitShare(self, initScope.scope(), CapsuleUtil.immutableSet(), false),
                closeScope -> {},
                closeLabel -> doCloseLabel(self, closeLabel.scope(), closeLabel.label()),
                query -> {},
                pQuery -> {},
                confirm -> {},
                match -> {},
                result -> {},
                typeCheckerState -> {},
                differResult -> {},
                differState -> {},
                envDifferState -> {},
                unitAdd -> {}
            );
            // @formatter:on
            for(IWaitFor<S, L, D> wf: ownTokens()) {
                wf.visit(cases);
            }

            pendingExternalDatums.asMap().forEach((d, futures) -> futures.elementSet().forEach(future -> {
                final D datum = previousResult.result().analysis().getExternalRepresentation(d);
                final D result = context.substituteScopes(datum, resultPatches.patches().invert());
                self.complete(future, result, null);
            }));

            final IPatchCollection.Immutable<S> allPatches = resultPatches.putAll(globalPatches);
            java.util.Set<IRecordedQuery<S, L, D>> newRecordedQueries =
                    previousResult.queries().stream().map(q -> q.patch(allPatches)).collect(Collectors.toSet());
            java.util.Set<IRecordedQuery<S, L, D>> patchedRemovedQueries =
                    removedQueries.stream().map(q -> q.patch(allPatches)).collect(Collectors.toSet());
            newRecordedQueries.removeAll(patchedRemovedQueries);
            newRecordedQueries.addAll(addedQueries);
            recordedQueries.addAll(newRecordedQueries);

            stateTransitionTrace = TransitionTrace.RELEASED;

            self.complete(confirmationResult, Optional.of(resultPatches), null);
            // Cancel all futures waiting for activation
            self.complete(whenActive, null, Release.instance);
            state = UnitState.RELEASED;

            resume();
            tryFinish(); // FIXME needed?
            logger.debug("{} released.", this);
        }
    }

    private boolean doRestart(boolean external) {
        if(state == UnitState.INIT_TC || state == UnitState.UNKNOWN) {
            logger.debug("{} restarting.", this);
            state = UnitState.ACTIVE;
            self.complete(confirmationResult, Optional.empty(), null);
            self.complete(whenActive, Unit.unit, null);
            // Not activating context here, as `doRestart` is also used before non-incremental runs.
            if(isDifferEnabled()) {
                final IDifferContext<S, L, D> context = differContext(typeChecker::internalData);
                final IDifferOps<S, L, D> differOps = differOps();
                if(previousResult != null) {
                    final IDifferContext<S, L, D> differContext = new StaticDifferContext<>(previousResult.scopeGraph(),
                            previousResult.scopes(), new DifferDataOps());

                    if(isConfirmationEnabled()) {
                        this.envDiffer = new IndexedEnvDiffer<>(new EnvDiffer<>(envDifferContext, differOps()));
                    }

                    if(external) {
                        final StateCapture<S, L, D, T> capture = previousResult.result().localState();
                        final java.util.Set<S> openScopes = Sets.union(capture.openScopes().elementSet(),
                                capture.unInitializedScopes().elementSet());

                        initDiffer(new ScopeGraphDiffer<>(context, differContext, differOps, edgeLabels),
                                capture.scopeGraph(), capture.scopes(), previousResult.result().sharedScopes(),
                                matchedBySharing.freeze(), openScopes, capture.openEdges());
                    } else {
                        initDiffer(new ScopeGraphDiffer<>(context, differContext, differOps, edgeLabels),
                                this.rootScopes, previousResult.rootScopes());
                    }
                } else {
                    initDiffer(new AddingDiffer<>(context, differOps, edgeLabels), Collections.emptyList(),
                            Collections.emptyList());
                }
            }
            pendingExternalDatums.asMap().forEach((d, futures) -> {
                typeChecker.getExternalDatum(d).whenComplete((d2, ex) -> futures.forEach(future -> {
                    future.complete(d2, ex);
                }));
            });

            resume();
            tryFinish(); // FIXME needed?
            logger.debug("{} restarted.", this);
            return true;
        }
        return false;
    }

    private void doImplicitActivate() {
        // If invoked synchronously, do a implicit transition to ACTIVE
{'fraction_non_alphanumeric': 0.12762885474555954, 'fraction_numerical': 0.0010896807235480004, 'mean_word_length': 2.3254971237034017, 'pattern_counts': {'":': 0, '<': 455, '<?xml version=': 0, '>': 539, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6987032', 'n_tokens_mistral': 20396, 'n_tokens_neox': 19839, 'n_words': 4218}