feat: cancellation token for ModelStore

15 files changed, 90 insertions, 21 deletions

diff --git a/subprojects/store/src/main/java/tools/refinery/store/model/Model.java b/subprojects/store/src/main/java/tools/refinery/store/model/Model.java
index e2ab72e7..c4ce5207 100644
--- a/subprojects/store/src/main/java/tools/refinery/store/model/Model.java
+++ b/subprojects/store/src/main/java/tools/refinery/store/model/Model.java
@@ -8,11 +8,9 @@ package tools.refinery.store.model;
 import tools.refinery.store.adapter.ModelAdapter;
 import tools.refinery.store.map.Version;
 import tools.refinery.store.map.Versioned;
-import tools.refinery.store.model.internal.VersionedInterpretation;
 import tools.refinery.store.representation.AnySymbol;
 import tools.refinery.store.representation.Symbol;
 
-import java.util.Map;
 import java.util.Optional;
 
 public interface Model extends Versioned {
@@ -38,4 +36,6 @@ public interface Model extends Versioned {
         void addListener(ModelListener listener);
 
         void removeListener(ModelListener listener);
+
+        void checkCancelled();
 }
diff --git a/subprojects/store/src/main/java/tools/refinery/store/model/ModelStore.java b/subprojects/store/src/main/java/tools/refinery/store/model/ModelStore.java
index 89382b3a..61abf126 100644
--- a/subprojects/store/src/main/java/tools/refinery/store/model/ModelStore.java
+++ b/subprojects/store/src/main/java/tools/refinery/store/model/ModelStore.java
@@ -26,6 +26,8 @@ public interface ModelStore {
 
         <T extends ModelStoreAdapter> T getAdapter(Class<T> adapterType);
 
+        void checkCancelled();
+
         static ModelStoreBuilder builder() {
                 return new ModelStoreBuilderImpl();
         }
diff --git a/subprojects/store/src/main/java/tools/refinery/store/model/ModelStoreBuilder.java b/subprojects/store/src/main/java/tools/refinery/store/model/ModelStoreBuilder.java
index 8f652762..9b2b38c3 100644
--- a/subprojects/store/src/main/java/tools/refinery/store/model/ModelStoreBuilder.java
+++ b/subprojects/store/src/main/java/tools/refinery/store/model/ModelStoreBuilder.java
@@ -8,12 +8,15 @@ package tools.refinery.store.model;
 import tools.refinery.store.adapter.ModelAdapterBuilder;
 import tools.refinery.store.representation.AnySymbol;
 import tools.refinery.store.representation.Symbol;
+import tools.refinery.store.util.CancellationToken;
 
 import java.util.Collection;
 import java.util.List;
 import java.util.Optional;
 
 public interface ModelStoreBuilder {
+        ModelStoreBuilder cancellationToken(CancellationToken cancellationToken);
+
         default ModelStoreBuilder symbols(AnySymbol... symbols) {
                 return symbols(List.of(symbols));
         }
diff --git a/subprojects/store/src/main/java/tools/refinery/store/model/internal/ModelImpl.java b/subprojects/store/src/main/java/tools/refinery/store/model/internal/ModelImpl.java
index 2b12d5a6..d11e431b 100644
--- a/subprojects/store/src/main/java/tools/refinery/store/model/internal/ModelImpl.java
+++ b/subprojects/store/src/main/java/tools/refinery/store/model/internal/ModelImpl.java
@@ -13,23 +13,26 @@ import tools.refinery.store.model.*;
 import tools.refinery.store.representation.AnySymbol;
 import tools.refinery.store.representation.Symbol;
 import tools.refinery.store.tuple.Tuple;
+import tools.refinery.store.util.CancellationToken;
 
 import java.util.*;
 
 public class ModelImpl implements Model {
-        private final ModelStore store;
+        private final ModelStoreImpl store;
         private Version state;
         private LinkedHashMap<? extends AnySymbol, ? extends VersionedInterpretation<?>> interpretations;
         private final List<ModelAdapter> adapters;
         private final List<ModelListener> listeners = new ArrayList<>();
+        private final CancellationToken cancellationToken;
         private boolean uncommittedChanges;
         private ModelAction pendingAction = ModelAction.NONE;
         private Version restoringToState = null;
 
-        ModelImpl(ModelStore store, Version state, int adapterCount) {
+        ModelImpl(ModelStoreImpl store, Version state, int adapterCount) {
                 this.store = store;
                 this.state = state;
                 adapters = new ArrayList<>(adapterCount);
+                cancellationToken = store.getCancellationToken();
         }
 
         void setInterpretations(LinkedHashMap<? extends AnySymbol, ? extends VersionedInterpretation<?>> interpretations) {
@@ -88,6 +91,7 @@ public class ModelImpl implements Model {
 
         @Override
         public Version commit() {
+                checkCancelled();
                 if (hasPendingAction()) {
                         throw pendingActionError("commit");
                 }
@@ -106,6 +110,7 @@ public class ModelImpl implements Model {
                         Version[] interpretationVersions = new Version[interpretations.size()];
                         int j = 0;
                         for (var interpretationEntry : interpretations.entrySet()) {
+                                checkCancelled();
                                 interpretationVersions[j++] = interpretationEntry.getValue().commit();
                         }
                         ModelVersion modelVersion = new ModelVersion(interpretationVersions);
@@ -125,6 +130,7 @@ public class ModelImpl implements Model {
 
         @Override
         public void restore(Version version) {
+                checkCancelled();
                 if (hasPendingAction()) {
                         throw pendingActionError("restore to %s".formatted(version));
                 }
@@ -140,6 +146,7 @@ public class ModelImpl implements Model {
                         }
                         int j = 0;
                         for (var interpretation : interpretations.values()) {
+                                checkCancelled();
                                 interpretation.restore(ModelVersion.getInternalVersion(version, j++));
                         }
 
@@ -187,4 +194,9 @@ public class ModelImpl implements Model {
         public void removeListener(ModelListener listener) {
                 listeners.remove(listener);
         }
+
+        @Override
+        public void checkCancelled() {
+                cancellationToken.checkCancelled();
+        }
 }
diff --git a/subprojects/store/src/main/java/tools/refinery/store/model/internal/ModelStoreBuilderImpl.java b/subprojects/store/src/main/java/tools/refinery/store/model/internal/ModelStoreBuilderImpl.java
index 2dde7a4c..53ecde5e 100644
--- a/subprojects/store/src/main/java/tools/refinery/store/model/internal/ModelStoreBuilderImpl.java
+++ b/subprojects/store/src/main/java/tools/refinery/store/model/internal/ModelStoreBuilderImpl.java
@@ -16,15 +16,29 @@ import tools.refinery.store.model.ModelStoreConfiguration;
 import tools.refinery.store.representation.AnySymbol;
 import tools.refinery.store.representation.Symbol;
 import tools.refinery.store.tuple.Tuple;
+import tools.refinery.store.util.CancellationToken;
 
 import java.util.*;
 
 public class ModelStoreBuilderImpl implements ModelStoreBuilder {
+        private CancellationToken cancellationToken;
         private final LinkedHashSet<AnySymbol> allSymbols = new LinkedHashSet<>();
         private final LinkedHashMap<SymbolEquivalenceClass<?>, List<AnySymbol>> equivalenceClasses = new LinkedHashMap<>();
         private final List<ModelAdapterBuilder> adapters = new ArrayList<>();
 
         @Override
+        public ModelStoreBuilder cancellationToken(CancellationToken cancellationToken) {
+                if (this.cancellationToken != null) {
+                        throw new IllegalStateException("Cancellation token was already set");
+                }
+                if (cancellationToken == null) {
+                        throw new IllegalStateException("Cancellation token must not be null");
+                }
+                this.cancellationToken = cancellationToken;
+                return this;
+        }
+
+        @Override
         public <T> ModelStoreBuilder symbol(Symbol<T> symbol) {
                 if (!allSymbols.add(symbol)) {
                         // No need to add symbol twice.
@@ -75,7 +89,8 @@ public class ModelStoreBuilderImpl implements ModelStoreBuilder {
                 for (var entry : equivalenceClasses.entrySet()) {
                         createStores(stores, entry.getKey(), entry.getValue());
                 }
-                var modelStore = new ModelStoreImpl(stores, adapters.size());
+                var modelStore = new ModelStoreImpl(stores, adapters.size(), cancellationToken == null ?
+                                CancellationToken.NONE : cancellationToken);
                 for (var adapterBuilder : adapters) {
                         var storeAdapter = adapterBuilder.build(modelStore);
                         modelStore.addAdapter(storeAdapter);
@@ -86,8 +101,8 @@ public class ModelStoreBuilderImpl implements ModelStoreBuilder {
         private <T> void createStores(Map<AnySymbol, VersionedMapStore<Tuple, ?>> stores,
                                                                   SymbolEquivalenceClass<T> equivalenceClass, List<AnySymbol> symbols) {
                 int size = symbols.size();
-                VersionedMapStoreFactory<Tuple,T> mapFactory = VersionedMapStore
-                                .<Tuple,T>builder()
+                VersionedMapStoreFactory<Tuple, T> mapFactory = VersionedMapStore
+                                .<Tuple, T>builder()
                                 .strategy(VersionedMapStoreFactoryBuilder.StoreStrategy.DELTA)
                                 .defaultValue(equivalenceClass.defaultValue())
                                 .build();
diff --git a/subprojects/store/src/main/java/tools/refinery/store/model/internal/ModelStoreImpl.java b/subprojects/store/src/main/java/tools/refinery/store/model/internal/ModelStoreImpl.java
index a320a618..fd4cc160 100644
--- a/subprojects/store/src/main/java/tools/refinery/store/model/internal/ModelStoreImpl.java
+++ b/subprojects/store/src/main/java/tools/refinery/store/model/internal/ModelStoreImpl.java
@@ -14,16 +14,20 @@ import tools.refinery.store.model.ModelDiffCursor;
 import tools.refinery.store.model.ModelStore;
 import tools.refinery.store.representation.AnySymbol;
 import tools.refinery.store.tuple.Tuple;
+import tools.refinery.store.util.CancellationToken;
 
 import java.util.*;
 
 public class ModelStoreImpl implements ModelStore {
         private final LinkedHashMap<? extends AnySymbol, ? extends VersionedMapStore<Tuple, ?>> stores;
         private final List<ModelStoreAdapter> adapters;
+        private final CancellationToken cancellationToken;
 
-        ModelStoreImpl(LinkedHashMap<? extends AnySymbol, ? extends VersionedMapStore<Tuple, ?>> stores, int adapterCount) {
+        ModelStoreImpl(LinkedHashMap<? extends AnySymbol, ? extends VersionedMapStore<Tuple, ?>> stores, int adapterCount,
+                                   CancellationToken cancellationToken) {
                 this.stores = stores;
                 adapters = new ArrayList<>(adapterCount);
+                this.cancellationToken = cancellationToken;
         }
 
         @Override
@@ -100,4 +104,13 @@ public class ModelStoreImpl implements ModelStore {
         void addAdapter(ModelStoreAdapter adapter) {
                 adapters.add(adapter);
         }
+
+        @Override
+        public void checkCancelled() {
+                cancellationToken.checkCancelled();
+        }
+
+        CancellationToken getCancellationToken() {
+                return cancellationToken;
+        }
 }
diff --git a/subprojects/store/src/main/java/tools/refinery/store/model/internal/VersionedInterpretation.java b/subprojects/store/src/main/java/tools/refinery/store/model/internal/VersionedInterpretation.java
index 71df3962..dcf0ad08 100644
--- a/subprojects/store/src/main/java/tools/refinery/store/model/internal/VersionedInterpretation.java
+++ b/subprojects/store/src/main/java/tools/refinery/store/model/internal/VersionedInterpretation.java
@@ -75,6 +75,7 @@ public abstract class VersionedInterpretation<T> implements Interpretation<T> {
         @Override
         public T put(Tuple key, T value) {
                 checkKey(key);
+                model.checkCancelled();
                 model.markAsChanged();
                 var oldValue = map.put(key, value);
                 valueChanged(key, oldValue, value, false);
@@ -83,15 +84,12 @@ public abstract class VersionedInterpretation<T> implements Interpretation<T> {
 
         @Override
         public void putAll(Cursor<Tuple, T> cursor) {
-                if (listeners.isEmpty()) {
-                        map.putAll(cursor);
-                        return;
-                }
                 model.markAsChanged();
                 if (cursor.getDependingMaps().contains(map)) {
                         List<Tuple> keys = new ArrayList<>();
                         List<T> values = new ArrayList<>();
                         while (cursor.move()) {
+                                model.checkCancelled();
                                 keys.add(cursor.getKey());
                                 values.add(cursor.getValue());
                         }
diff --git a/subprojects/store/src/main/java/tools/refinery/store/statecoding/StateCodeCalculatorFactory.java b/subprojects/store/src/main/java/tools/refinery/store/statecoding/StateCodeCalculatorFactory.java
index 04e17a13..809205e4 100644
--- a/subprojects/store/src/main/java/tools/refinery/store/statecoding/StateCodeCalculatorFactory.java
+++ b/subprojects/store/src/main/java/tools/refinery/store/statecoding/StateCodeCalculatorFactory.java
@@ -7,9 +7,10 @@ package tools.refinery.store.statecoding;
 
 import org.eclipse.collections.api.set.primitive.IntSet;
 import tools.refinery.store.model.Interpretation;
+import tools.refinery.store.model.Model;
 
 import java.util.List;
 
 public interface StateCodeCalculatorFactory {
-        StateCodeCalculator create(List<? extends Interpretation<?>> interpretations, IntSet individuals);
+        StateCodeCalculator create(Model model, List<? extends Interpretation<?>> interpretations, IntSet individuals);
 }
diff --git a/subprojects/store/src/main/java/tools/refinery/store/statecoding/internal/StateCoderStoreAdapterImpl.java b/subprojects/store/src/main/java/tools/refinery/store/statecoding/internal/StateCoderStoreAdapterImpl.java
index 89586bfb..2cb6f3c1 100644
--- a/subprojects/store/src/main/java/tools/refinery/store/statecoding/internal/StateCoderStoreAdapterImpl.java
+++ b/subprojects/store/src/main/java/tools/refinery/store/statecoding/internal/StateCoderStoreAdapterImpl.java
@@ -68,7 +68,7 @@ public class StateCoderStoreAdapterImpl implements StateCoderStoreAdapter {
         @Override
         public StateCoderAdapter createModelAdapter(Model model) {
                 var interpretations = symbols.stream().map(model::getInterpretation).toList();
-                var coder = codeCalculatorFactory.create(interpretations, individuals);
+                var coder = codeCalculatorFactory.create(model, interpretations, individuals);
                 return new StateCoderAdapterImpl(this, coder, model);
         }
 }
diff --git a/subprojects/store/src/main/java/tools/refinery/store/statecoding/neighbourhood/AbstractNeighbourhoodCalculator.java b/subprojects/store/src/main/java/tools/refinery/store/statecoding/neighbourhood/AbstractNeighbourhoodCalculator.java
index 5d390da2..4cef6786 100644
--- a/subprojects/store/src/main/java/tools/refinery/store/statecoding/neighbourhood/AbstractNeighbourhoodCalculator.java
+++ b/subprojects/store/src/main/java/tools/refinery/store/statecoding/neighbourhood/AbstractNeighbourhoodCalculator.java
@@ -10,6 +10,7 @@ import org.eclipse.collections.api.map.primitive.MutableIntLongMap;
 import org.eclipse.collections.api.set.primitive.IntSet;
 import tools.refinery.store.model.AnyInterpretation;
 import tools.refinery.store.model.Interpretation;
+import tools.refinery.store.model.Model;
 import tools.refinery.store.statecoding.ObjectCode;
 import tools.refinery.store.tuple.Tuple;
 import tools.refinery.store.tuple.Tuple0;
@@ -17,13 +18,16 @@ import tools.refinery.store.tuple.Tuple0;
 import java.util.*;
 
 public abstract class AbstractNeighbourhoodCalculator {
+        protected final Model model;
         protected final List<AnyInterpretation> nullImpactValues;
         protected final LinkedHashMap<AnyInterpretation, long[]> impactValues;
         protected final MutableIntLongMap individualHashValues = IntLongMaps.mutable.empty();
 
         protected static final long PRIME = 31;
 
-        protected AbstractNeighbourhoodCalculator(List<? extends AnyInterpretation> interpretations, IntSet individuals) {
+        protected AbstractNeighbourhoodCalculator(Model model, List<? extends AnyInterpretation> interpretations,
+                                                                                          IntSet individuals) {
+                this.model = model;
                 this.nullImpactValues = new ArrayList<>();
                 this.impactValues = new LinkedHashMap<>();
                 // Random isn't used for cryptographical purposes but just to assign distinguishable identifiers to symbols.
diff --git a/subprojects/store/src/main/java/tools/refinery/store/statecoding/neighbourhood/LazyNeighbourhoodCalculator.java b/subprojects/store/src/main/java/tools/refinery/store/statecoding/neighbourhood/LazyNeighbourhoodCalculator.java
index c188a839..04335141 100644
--- a/subprojects/store/src/main/java/tools/refinery/store/statecoding/neighbourhood/LazyNeighbourhoodCalculator.java
+++ b/subprojects/store/src/main/java/tools/refinery/store/statecoding/neighbourhood/LazyNeighbourhoodCalculator.java
@@ -12,6 +12,7 @@ import org.eclipse.collections.api.set.primitive.IntSet;
 import tools.refinery.store.map.Cursor;
 import tools.refinery.store.model.AnyInterpretation;
 import tools.refinery.store.model.Interpretation;
+import tools.refinery.store.model.Model;
 import tools.refinery.store.statecoding.StateCodeCalculator;
 import tools.refinery.store.statecoding.StateCoderResult;
 import tools.refinery.store.tuple.Tuple;
@@ -19,8 +20,9 @@ import tools.refinery.store.tuple.Tuple;
 import java.util.List;
 
 public class LazyNeighbourhoodCalculator extends AbstractNeighbourhoodCalculator implements StateCodeCalculator {
-        public LazyNeighbourhoodCalculator(List<? extends AnyInterpretation> interpretations, IntSet individuals) {
-                super(interpretations, individuals);
+        public LazyNeighbourhoodCalculator(Model model, List<? extends AnyInterpretation> interpretations,
+                                                                           IntSet individuals) {
+                super(model, interpretations, individuals);
         }
 
         public StateCoderResult calculateCodes() {
diff --git a/subprojects/store/src/main/java/tools/refinery/store/statecoding/neighbourhood/NeighbourhoodCalculator.java b/subprojects/store/src/main/java/tools/refinery/store/statecoding/neighbourhood/NeighbourhoodCalculator.java
index 1442c915..5e6de53b 100644
--- a/subprojects/store/src/main/java/tools/refinery/store/statecoding/neighbourhood/NeighbourhoodCalculator.java
+++ b/subprojects/store/src/main/java/tools/refinery/store/statecoding/neighbourhood/NeighbourhoodCalculator.java
@@ -8,6 +8,7 @@ package tools.refinery.store.statecoding.neighbourhood;
 import org.eclipse.collections.api.set.primitive.IntSet;
 import tools.refinery.store.map.Cursor;
 import tools.refinery.store.model.Interpretation;
+import tools.refinery.store.model.Model;
 import tools.refinery.store.statecoding.ObjectCode;
 import tools.refinery.store.statecoding.StateCodeCalculator;
 import tools.refinery.store.statecoding.StateCoderResult;
@@ -21,17 +22,19 @@ public class NeighbourhoodCalculator extends AbstractNeighbourhoodCalculator imp
         private ObjectCodeImpl previousObjectCode = new ObjectCodeImpl();
         private ObjectCodeImpl nextObjectCode = new ObjectCodeImpl();
 
-        public NeighbourhoodCalculator(List<? extends Interpretation<?>> interpretations, IntSet individuals) {
-                super(interpretations, individuals);
+        public NeighbourhoodCalculator(Model model, List<? extends Interpretation<?>> interpretations, IntSet individuals) {
+                super(model, interpretations, individuals);
         }
 
         public StateCoderResult calculateCodes() {
+                model.checkCancelled();
                 previousObjectCode.clear();
                 nextObjectCode.clear();
                 initializeWithIndividuals(previousObjectCode);
 
                 int rounds = 0;
                 do {
+                        model.checkCancelled();
                         constructNextObjectCodes(previousObjectCode, nextObjectCode);
                         var tempObjectCode = previousObjectCode;
                         previousObjectCode = nextObjectCode;
@@ -60,6 +63,7 @@ public class NeighbourhoodCalculator extends AbstractNeighbourhoodCalculator imp
 
         private void constructNextObjectCodes(ObjectCodeImpl previous, ObjectCodeImpl next) {
                 for (var impactValueEntry : this.impactValues.entrySet()) {
+                        model.checkCancelled();
                         Interpretation<?> interpretation = (Interpretation<?>) impactValueEntry.getKey();
                         var cursor = interpretation.getAll();
                         int arity = interpretation.getSymbol().arity();
diff --git a/subprojects/store/src/main/java/tools/refinery/store/util/CancellationToken.java b/subprojects/store/src/main/java/tools/refinery/store/util/CancellationToken.java
new file mode 100644
index 00000000..be294013
--- /dev/null
+++ b/subprojects/store/src/main/java/tools/refinery/store/util/CancellationToken.java
@@ -0,0 +1,13 @@
+/*
+ * SPDX-FileCopyrightText: 2023 The Refinery Authors <https://refinery.tools/>
+ *
+ * SPDX-License-Identifier: EPL-2.0
+ */
+package tools.refinery.store.util;
+
+@FunctionalInterface
+public interface CancellationToken {
+        CancellationToken NONE = () -> {};
+
+        void checkCancelled();
+}
diff --git a/subprojects/store/src/test/java/tools/refinery/store/statecoding/EquivalenceTest.java b/subprojects/store/src/test/java/tools/refinery/store/statecoding/EquivalenceTest.java
index 3c35849e..f2d2f7b5 100644
--- a/subprojects/store/src/test/java/tools/refinery/store/statecoding/EquivalenceTest.java
+++ b/subprojects/store/src/test/java/tools/refinery/store/statecoding/EquivalenceTest.java
@@ -192,7 +192,8 @@ class EquivalenceTest {
                 ModelStore store = ModelStore.builder()
                                 .symbols(person, age, friend, parents, population)
                                 .with(StateCoderAdapter.builder()
-                                                .stateCodeCalculatorFactory((p1, p2) -> calculator))
+                                                .stateCodeCalculatorFactory((ignoredModel, ignoredInterpretations, ignoredIndividuals) ->
+                                                                calculator))
                                 .build();
 
                 var stateCoder = store.getAdapter(StateCoderStoreAdapter.class);
diff --git a/subprojects/store/src/test/java/tools/refinery/store/statecoding/StateCoderBuildTest.java b/subprojects/store/src/test/java/tools/refinery/store/statecoding/StateCoderBuildTest.java
index 0b738005..0928aa8e 100644
--- a/subprojects/store/src/test/java/tools/refinery/store/statecoding/StateCoderBuildTest.java
+++ b/subprojects/store/src/test/java/tools/refinery/store/statecoding/StateCoderBuildTest.java
@@ -124,7 +124,8 @@ class StateCoderBuildTest {
                 var store = ModelStore.builder()
                                 .symbols(friend)
                                 .with(StateCoderAdapter.builder()
-                                                .stateCodeCalculatorFactory((interpretations, individuals) -> mock))
+                                                .stateCodeCalculatorFactory((ignoredModel, ignoredInterpretations, ignoredIndividuals) ->
+                                                                mock))
                                 .build();
 
                 var model = store.createEmptyModel();