feat: cancellation token for ModelStore

2 files changed, 15 insertions, 8 deletions

diff --git a/subprojects/store-reasoning/src/main/java/tools/refinery/store/reasoning/translator/multiobject/MultiObjectInitializer.java b/subprojects/store-reasoning/src/main/java/tools/refinery/store/reasoning/translator/multiobject/MultiObjectInitializer.java
index f11ab46b..89918155 100644
--- a/subprojects/store-reasoning/src/main/java/tools/refinery/store/reasoning/translator/multiobject/MultiObjectInitializer.java
+++ b/subprojects/store-reasoning/src/main/java/tools/refinery/store/reasoning/translator/multiobject/MultiObjectInitializer.java
@@ -30,9 +30,9 @@ class MultiObjectInitializer implements PartialModelInitializer {
 
         @Override
         public void initialize(Model model, ModelSeed modelSeed) {
-                var intervals = initializeIntervals(modelSeed);
-                initializeExists(intervals, modelSeed);
-                initializeEquals(intervals, modelSeed);
+                var intervals = initializeIntervals(model, modelSeed);
+                initializeExists(intervals, model, modelSeed);
+                initializeEquals(intervals, model, modelSeed);
                 var countInterpretation = model.getInterpretation(countSymbol);
                 var uniqueTable = new HashMap<CardinalityInterval, CardinalityInterval>();
                 for (int i = 0; i < intervals.length; i++) {
@@ -47,12 +47,13 @@ class MultiObjectInitializer implements PartialModelInitializer {
         }
 
         @NotNull
-        private CardinalityInterval[] initializeIntervals(ModelSeed modelSeed) {
+        private CardinalityInterval[] initializeIntervals(Model model, ModelSeed modelSeed) {
                 var intervals = new CardinalityInterval[modelSeed.getNodeCount()];
                 if (modelSeed.containsSeed(MultiObjectTranslator.COUNT_SYMBOL)) {
                         Arrays.fill(intervals, CardinalityIntervals.ONE);
                         var cursor = modelSeed.getCursor(MultiObjectTranslator.COUNT_SYMBOL, CardinalityIntervals.ONE);
                         while (cursor.move()) {
+                                model.checkCancelled();
                                 int i = cursor.getKey().get(0);
                                 checkNodeId(intervals, i);
                                 intervals[i] = cursor.getValue();
@@ -70,12 +71,13 @@ class MultiObjectInitializer implements PartialModelInitializer {
                 return intervals;
         }
 
-        private void initializeExists(CardinalityInterval[] intervals, ModelSeed modelSeed) {
+        private void initializeExists(CardinalityInterval[] intervals, Model model, ModelSeed modelSeed) {
                 if (!modelSeed.containsSeed(ReasoningAdapter.EXISTS_SYMBOL)) {
                         return;
                 }
                 var cursor = modelSeed.getCursor(ReasoningAdapter.EXISTS_SYMBOL, TruthValue.UNKNOWN);
                 while (cursor.move()) {
+                        model.checkCancelled();
                         int i = cursor.getKey().get(0);
                         checkNodeId(intervals, i);
                         switch (cursor.getValue()) {
@@ -89,13 +91,14 @@ class MultiObjectInitializer implements PartialModelInitializer {
                 }
         }
 
-        private void initializeEquals(CardinalityInterval[] intervals, ModelSeed modelSeed) {
+        private void initializeEquals(CardinalityInterval[] intervals, Model model, ModelSeed modelSeed) {
                 if (!modelSeed.containsSeed(ReasoningAdapter.EQUALS_SYMBOL)) {
                         return;
                 }
                 var seed = modelSeed.getSeed(ReasoningAdapter.EQUALS_SYMBOL);
                 var cursor = seed.getCursor(TruthValue.FALSE, modelSeed.getNodeCount());
                 while (cursor.move()) {
+                        model.checkCancelled();
                         var key = cursor.getKey();
                         int i = key.get(0);
                         int otherIndex = key.get(1);
@@ -116,6 +119,7 @@ class MultiObjectInitializer implements PartialModelInitializer {
                         }
                 }
                 for (int i = 0; i < intervals.length; i++) {
+                        model.checkCancelled();
                         if (seed.get(Tuple.of(i, i)) == TruthValue.FALSE) {
                                 throw new TranslationException(ReasoningAdapter.EQUALS_SYMBOL, "Inconsistent equality for node " + i);
                         }
diff --git a/subprojects/store-reasoning/src/main/java/tools/refinery/store/reasoning/translator/typehierarchy/TypeHierarchyInitializer.java b/subprojects/store-reasoning/src/main/java/tools/refinery/store/reasoning/translator/typehierarchy/TypeHierarchyInitializer.java
index c74f1e78..233e43f0 100644
--- a/subprojects/store-reasoning/src/main/java/tools/refinery/store/reasoning/translator/typehierarchy/TypeHierarchyInitializer.java
+++ b/subprojects/store-reasoning/src/main/java/tools/refinery/store/reasoning/translator/typehierarchy/TypeHierarchyInitializer.java
@@ -31,20 +31,23 @@ public class TypeHierarchyInitializer implements PartialModelInitializer {
                 var inferredTypes = new InferredType[modelSeed.getNodeCount()];
                 Arrays.fill(inferredTypes, typeHierarchy.getUnknownType());
                 for (var type : typeHierarchy.getAllTypes()) {
-                        initializeType(type, inferredTypes, modelSeed);
+                        model.checkCancelled();
+                        initializeType(type, inferredTypes, model, modelSeed);
                 }
                 var typeInterpretation = model.getInterpretation(typeSymbol);
                 var uniqueTable = new HashMap<InferredType, InferredType>();
                 for (int i = 0; i < inferredTypes.length; i++) {
+                        model.checkCancelled();
                         var uniqueType = uniqueTable.computeIfAbsent(inferredTypes[i], Function.identity());
                         typeInterpretation.put(Tuple.of(i), uniqueType);
                 }
         }
 
-        private void initializeType(PartialRelation type, InferredType[] inferredTypes, ModelSeed modelSeed) {
+        private void initializeType(PartialRelation type, InferredType[] inferredTypes, Model model, ModelSeed modelSeed) {
                 var cursor = modelSeed.getCursor(type, TruthValue.UNKNOWN);
                 var analysisResult = typeHierarchy.getAnalysisResult(type);
                 while (cursor.move()) {
+                        model.checkCancelled();
                         var i = cursor.getKey().get(0);
                         checkNodeId(inferredTypes, i);
                         var value = cursor.getValue();