Merge pull request #37 from nagilooh/design-space-exploration

Design-space exploration improvements

13 files changed, 296 insertions, 211 deletions

diff --git a/subprojects/store-dse/src/main/java/tools/refinery/store/dse/DesignSpaceExplorationAdapter.java b/subprojects/store-dse/src/main/java/tools/refinery/store/dse/DesignSpaceExplorationAdapter.java
index 8963a496..5aed5298 100644
--- a/subprojects/store-dse/src/main/java/tools/refinery/store/dse/DesignSpaceExplorationAdapter.java
+++ b/subprojects/store-dse/src/main/java/tools/refinery/store/dse/DesignSpaceExplorationAdapter.java
@@ -38,7 +38,9 @@ public interface DesignSpaceExplorationAdapter extends ModelAdapter {
 
         public boolean backtrack();
 
-        public Fitness calculateFitness();
+        public boolean backtrack(String reason);
+
+        public Fitness getFitness();
 
         public void newSolution();
 
@@ -48,9 +50,7 @@ public interface DesignSpaceExplorationAdapter extends ModelAdapter {
 
         public boolean fireActivation(Activation activation);
 
-        public void fireRandomActivation();
-
-        public boolean isCurrentInTrajectory();
+        public boolean fireRandomActivation();
 
         public List<Version> getTrajectory();
 
@@ -63,4 +63,6 @@ public interface DesignSpaceExplorationAdapter extends ModelAdapter {
         public void setRandom(Random random);
 
         public void setRandom(long seed);
+
+        public List<Version> getSolutions();
 }
diff --git a/subprojects/store-dse/src/main/java/tools/refinery/store/dse/DesignSpaceExplorationStoreAdapter.java b/subprojects/store-dse/src/main/java/tools/refinery/store/dse/DesignSpaceExplorationStoreAdapter.java
index 186bfebb..0252748d 100644
--- a/subprojects/store-dse/src/main/java/tools/refinery/store/dse/DesignSpaceExplorationStoreAdapter.java
+++ b/subprojects/store-dse/src/main/java/tools/refinery/store/dse/DesignSpaceExplorationStoreAdapter.java
@@ -6,6 +6,24 @@
 package tools.refinery.store.dse;
 
 import tools.refinery.store.adapter.ModelStoreAdapter;
+import tools.refinery.store.dse.internal.TransformationRule;
+import tools.refinery.store.dse.objectives.Objective;
+import tools.refinery.store.model.Model;
+import tools.refinery.store.query.dnf.RelationalQuery;
+
+import java.util.List;
+import java.util.Set;
 
 public interface DesignSpaceExplorationStoreAdapter extends ModelStoreAdapter {
+
+        @Override
+        DesignSpaceExplorationAdapter createModelAdapter(Model model);
+
+        Set<TransformationRule> getTransformationSpecifications();
+
+        Set<RelationalQuery> getGlobalConstraints();
+
+        List<Objective> getObjectives();
+
+        Strategy getStrategy();
 }
diff --git a/subprojects/store-dse/src/main/java/tools/refinery/store/dse/Strategy.java b/subprojects/store-dse/src/main/java/tools/refinery/store/dse/Strategy.java
index e240f478..c60a4410 100644
--- a/subprojects/store-dse/src/main/java/tools/refinery/store/dse/Strategy.java
+++ b/subprojects/store-dse/src/main/java/tools/refinery/store/dse/Strategy.java
@@ -7,7 +7,7 @@ package tools.refinery.store.dse;
 
 public interface Strategy {
 
-        public void initStrategy(DesignSpaceExplorationAdapter designSpaceExplorationAdapter);
+        void initialize(DesignSpaceExplorationAdapter designSpaceExplorationAdapter);
 
-        public void explore();
+        void explore();
 }
diff --git a/subprojects/store-dse/src/main/java/tools/refinery/store/dse/internal/DesignSpaceExplorationAdapterImpl.java b/subprojects/store-dse/src/main/java/tools/refinery/store/dse/internal/DesignSpaceExplorationAdapterImpl.java
index b32e9696..220f0b2d 100644
--- a/subprojects/store-dse/src/main/java/tools/refinery/store/dse/internal/DesignSpaceExplorationAdapterImpl.java
+++ b/subprojects/store-dse/src/main/java/tools/refinery/store/dse/internal/DesignSpaceExplorationAdapterImpl.java
@@ -33,27 +33,24 @@ public class DesignSpaceExplorationAdapterImpl implements DesignSpaceExploration
         private final Model model;
         private final ModelQueryAdapter queryEngine;
         private final DesignSpaceExplorationStoreAdapterImpl storeAdapter;
-        private final LinkedHashSet<TransformationRule> transformationRules;
-        private final LinkedHashSet<RelationalQuery> globalConstraints;
+        private final Set<TransformationRule> transformationRules;
+        private final Set<RelationalQuery> globalConstraints;
         private final List<Objective> objectives;
         private final LinkedHashSet<ResultSet<Boolean>> globalConstraintResultSets = new LinkedHashSet<>();
         private final Interpretation<Integer> sizeInterpretation;
         private final Strategy strategy;
 
         private ObjectiveComparatorHelper objectiveComparatorHelper;
-        private List<Version> trajectory = new LinkedList<>();
-        private Fitness lastFitness;
-        private final LinkedHashSet<Version> solutions = new LinkedHashSet<>();
-        private Map<Version, LinkedHashSet<Activation>> statesAndUntraversedActivations;
-        private Map<Version, LinkedHashSet<Activation>> statesAndTraversedActivations;
+        private List<Version> trajectory = new ArrayList<>();
+        private Map<Version, Version> parents = new HashMap<>();
+        private final List<Version> solutions = new ArrayList<>();
+        private Map<Version, List<Activation>> statesAndTraversedActivations;
         private Random random = new Random();
         private boolean isNewState = false;
         private final boolean isVisualizationEnabled;
         private final ModelVisualizerAdapter modelVisualizerAdapter;
 
-        public List<Version> getTrajectory() {
-                return new LinkedList<>(trajectory);
-        }
+        private final Map<Version, Fitness> fitnessCache = new HashMap<>();
 
         public DesignSpaceExplorationAdapterImpl(Model model, DesignSpaceExplorationStoreAdapterImpl storeAdapter) {
                 this.model = model;
@@ -72,14 +69,18 @@ public class DesignSpaceExplorationAdapterImpl implements DesignSpaceExploration
                 }
 
                 objectives = storeAdapter.getObjectives();
-                statesAndUntraversedActivations = new HashMap<>();
                 statesAndTraversedActivations = new HashMap<>();
                 strategy = storeAdapter.getStrategy();
+                strategy.initialize(this);
                 modelVisualizerAdapter = model.tryGetAdapter(ModelVisualizerAdapter.class).orElse(null);
                 isVisualizationEnabled = modelVisualizerAdapter != null;
 
         }
 
+        public List<Version> getTrajectory() {
+                return new ArrayList<>(trajectory);
+        }
+
         @Override
         public Model getModel() {
                 return model;
@@ -91,13 +92,13 @@ public class DesignSpaceExplorationAdapterImpl implements DesignSpaceExploration
         }
 
         @Override
-        public LinkedHashSet<Version> explore() {
+        public List<Version> explore() {
                 var state = model.commit();
                 trajectory.add(state);
-                statesAndUntraversedActivations.put(state, getAllActivations());
-                statesAndTraversedActivations.put(state, new LinkedHashSet<>());
-                strategy.initStrategy(this);
                 strategy.explore();
+                if (isVisualizationEnabled) {
+                        modelVisualizerAdapter.visualize();
+                }
                 return solutions;
         }
 
@@ -137,14 +138,22 @@ public class DesignSpaceExplorationAdapterImpl implements DesignSpaceExploration
 
         @Override
         public boolean backtrack() {
+                return backtrack("");
+        }
+        @Override
+        public boolean backtrack(String reason) {
                 if (trajectory.size() < 2) {
                         return false;
                 }
+                var currentState = model.getState();
+                if (!parents.containsKey(currentState)) {
+                        return false;
+                }
                 if (isVisualizationEnabled) {
                         modelVisualizerAdapter.addTransition(trajectory.get(trajectory.size() - 1),
-                                        trajectory.get(trajectory.size() - 2), "backtrack");
+                                        trajectory.get(trajectory.size() - 2), "backtrack(" + reason + ")");
                 }
-                model.restore(trajectory.get(trajectory.size() - 2));
+                model.restore(parents.get(model.getState()));
                 trajectory.remove(trajectory.size() - 1);
                 return true;
         }
@@ -156,7 +165,7 @@ public class DesignSpaceExplorationAdapterImpl implements DesignSpaceExploration
 //                      modelVisualizerAdapter.addTransition(this.trajectory.get(trajectory.size() - 1),
 //                                      trajectory.get(trajectory.size() - 1), "restore");
 //              }
-                this.trajectory = trajectory;
+                this.trajectory = new ArrayList<>(trajectory);
 
         }
 
@@ -171,7 +180,16 @@ public class DesignSpaceExplorationAdapterImpl implements DesignSpaceExploration
         }
 
         @Override
-        public Fitness calculateFitness() {
+        public List<Version> getSolutions() {
+                return solutions;
+        }
+
+        @Override
+        public Fitness getFitness() {
+        return fitnessCache.computeIfAbsent(model.getState(), s -> calculateFitness());
+        }
+
+        private Fitness calculateFitness() {
                 Fitness result = new Fitness();
                 boolean satisfiesHardObjectives = true;
                 for (Objective objective : objectives) {
@@ -183,8 +201,6 @@ public class DesignSpaceExplorationAdapterImpl implements DesignSpaceExploration
                 }
                 result.setSatisfiesHardObjectives(satisfiesHardObjectives);
 
-                lastFitness = result;
-
                 return result;
         }
 
@@ -203,15 +219,19 @@ public class DesignSpaceExplorationAdapterImpl implements DesignSpaceExploration
         }
 
         public LinkedHashSet<Activation> getUntraversedActivations() {
-//              return statesAndUntraversedActivations.get(model.getState());
-                LinkedHashSet<Activation> untraversedActivations = new LinkedHashSet<>();
-                for (Activation activation : getAllActivations()) {
-                        if (!statesAndTraversedActivations.get(model.getState()).contains(activation)) {
-                                untraversedActivations.add(activation);
+                var traversedActivations = statesAndTraversedActivations.get(model.getState());
+                if (traversedActivations == null) {
+                        return new LinkedHashSet<>(getAllActivations());
+                }
+                else {
+                        LinkedHashSet<Activation> untraversedActivations = new LinkedHashSet<>();
+                        for (Activation activation : getAllActivations()) {
+                                if (!traversedActivations.contains(activation)) {
+                                        untraversedActivations.add(activation);
+                                }
                         }
+                        return untraversedActivations;
                 }
-
-                return untraversedActivations;
         }
 
         @Override
@@ -220,37 +240,29 @@ public class DesignSpaceExplorationAdapterImpl implements DesignSpaceExploration
                         return false;
                 }
                 var previousState = model.getState();
-                if (!statesAndUntraversedActivations.get(previousState).contains(activation)) {
-//                      TODO: throw exception?
-                        return false;
-                }
                 if (!activation.fire()) {
                         return false;
                 }
-                statesAndUntraversedActivations.get(previousState).remove(activation);
-                statesAndTraversedActivations.get(previousState).add(activation);
+                statesAndTraversedActivations.computeIfAbsent(previousState, s -> new ArrayList<>()).add(activation);
                 var newState = model.commit();
                 trajectory.add(newState);
-                isNewState = !statesAndUntraversedActivations.containsKey(newState);
-                statesAndUntraversedActivations.put(newState, getAllActivations());
-                statesAndTraversedActivations.put(newState, new LinkedHashSet<>());
+                parents.put(newState, previousState);
+                isNewState = !statesAndTraversedActivations.containsKey(newState);
                 if (isVisualizationEnabled) {
                         if (isNewState) {
-                                modelVisualizerAdapter.addState(newState);
+                                modelVisualizerAdapter.addState(newState, getFitness().values());
                         }
-                        modelVisualizerAdapter.addTransition(trajectory.get(trajectory.size() - 2),
-                                        trajectory.get(trajectory.size() - 1), activation.transformationRule().getName(),
+                        modelVisualizerAdapter.addTransition(previousState, newState, activation.transformationRule().getName(),
                                         activation.activation());
                 }
                 return true;
         }
 
         @Override
-        public void fireRandomActivation() {
+        public boolean fireRandomActivation() {
                 var activations = getUntraversedActivations();
                 if (activations.isEmpty()) {
-//                      TODO: throw exception
-                        return;
+                        return false;
                 }
                 int index = random.nextInt(activations.size());
                 var iterator = activations.iterator();
@@ -258,31 +270,21 @@ public class DesignSpaceExplorationAdapterImpl implements DesignSpaceExploration
                         iterator.next();
                 }
                 var activationId = iterator.next();
-                fireActivation(activationId);
+                return fireActivation(activationId);
         }
 
-        @Override
-        public boolean isCurrentInTrajectory() {
-                return trajectory.contains(model.getState());
-        }
-
-        public LinkedHashSet<Activation> getAllActivations() {
-                LinkedHashSet<Activation> result = new LinkedHashSet<>();
+        public List<Activation> getAllActivations() {
+                List<Activation> result = new LinkedList<>();
                 for (var rule : transformationRules) {
-                        result.addAll(rule.getAllActivations());
+                        result.addAll(rule.getAllActivationsAsList());
                 }
                 return result;
         }
 
         public boolean isCurrentStateAlreadyTraversed() {
-//              TODO: check isomorphism?
                 return !isNewState;
         }
 
-        public Fitness getLastFitness() {
-                return lastFitness;
-        }
-
         public ObjectiveComparatorHelper getObjectiveComparatorHelper() {
                 if (objectiveComparatorHelper == null) {
                         objectiveComparatorHelper = new ObjectiveComparatorHelper(objectives);
diff --git a/subprojects/store-dse/src/main/java/tools/refinery/store/dse/internal/DesignSpaceExplorationStoreAdapterImpl.java b/subprojects/store-dse/src/main/java/tools/refinery/store/dse/internal/DesignSpaceExplorationStoreAdapterImpl.java
index 09925ae7..fea39886 100644
--- a/subprojects/store-dse/src/main/java/tools/refinery/store/dse/internal/DesignSpaceExplorationStoreAdapterImpl.java
+++ b/subprojects/store-dse/src/main/java/tools/refinery/store/dse/internal/DesignSpaceExplorationStoreAdapterImpl.java
@@ -5,27 +5,26 @@
  */
 package tools.refinery.store.dse.internal;
 
-import tools.refinery.store.adapter.ModelAdapter;
-import tools.refinery.store.model.Model;
-import tools.refinery.store.model.ModelStore;
-import tools.refinery.store.query.dnf.RelationalQuery;
 import tools.refinery.store.dse.DesignSpaceExplorationStoreAdapter;
 import tools.refinery.store.dse.Strategy;
 import tools.refinery.store.dse.objectives.Objective;
+import tools.refinery.store.model.Model;
+import tools.refinery.store.model.ModelStore;
+import tools.refinery.store.query.dnf.RelationalQuery;
 
-import java.util.LinkedHashSet;
 import java.util.List;
+import java.util.Set;
 
 public class DesignSpaceExplorationStoreAdapterImpl implements DesignSpaceExplorationStoreAdapter {
         private final ModelStore store;
-        private final LinkedHashSet<TransformationRule> transformationSpecifications;
-        private final LinkedHashSet<RelationalQuery> globalConstraints;
+        private final Set<TransformationRule> transformationSpecifications;
+        private final Set<RelationalQuery> globalConstraints;
         private final List<Objective> objectives;
         private final Strategy strategy;
 
         public DesignSpaceExplorationStoreAdapterImpl(ModelStore store,
-                                                                                                  LinkedHashSet<TransformationRule> transformationSpecifications,
-                                                                                                  LinkedHashSet<RelationalQuery> globalConstraints,
+                                                                                                  Set<TransformationRule> transformationSpecifications,
+                                                                                                  Set<RelationalQuery> globalConstraints,
                                                                                                   List<Objective> objectives, Strategy strategy) {
                         this.store = store;
                         this.transformationSpecifications = transformationSpecifications;
@@ -40,22 +39,26 @@ public class DesignSpaceExplorationStoreAdapterImpl implements DesignSpaceExplor
         }
 
         @Override
-        public ModelAdapter createModelAdapter(Model model) {
+        public DesignSpaceExplorationAdapterImpl createModelAdapter(Model model) {
                 return new DesignSpaceExplorationAdapterImpl(model, this);
         }
 
-        public LinkedHashSet<TransformationRule> getTransformationSpecifications() {
+        @Override
+        public Set<TransformationRule> getTransformationSpecifications() {
                 return transformationSpecifications;
         }
 
-        public LinkedHashSet<RelationalQuery> getGlobalConstraints() {
+        @Override
+        public Set<RelationalQuery> getGlobalConstraints() {
                 return globalConstraints;
         }
 
+        @Override
         public List<Objective> getObjectives() {
                 return objectives;
         }
 
+        @Override
         public Strategy getStrategy() {
                 return strategy;
         }
diff --git a/subprojects/store-dse/src/main/java/tools/refinery/store/dse/internal/TransformationRule.java b/subprojects/store-dse/src/main/java/tools/refinery/store/dse/internal/TransformationRule.java
index 015d4815..8123c0d6 100644
--- a/subprojects/store-dse/src/main/java/tools/refinery/store/dse/internal/TransformationRule.java
+++ b/subprojects/store-dse/src/main/java/tools/refinery/store/dse/internal/TransformationRule.java
@@ -14,8 +14,7 @@ import tools.refinery.store.query.resultset.OrderedResultSet;
 import tools.refinery.store.query.resultset.ResultSet;
 import tools.refinery.store.tuple.Tuple;
 
-import java.util.LinkedHashSet;
-import java.util.Random;
+import java.util.*;
 
 public class TransformationRule {
 
@@ -66,11 +65,11 @@ public class TransformationRule {
                 return precondition;
         }
 
-        public ResultSet<Boolean> getAllActivationsAsSets() {
+        public ResultSet<Boolean> getAllActivationsAsResultSet() {
                 return activations;
         }
 
-        public LinkedHashSet<Activation> getAllActivations() {
+        public Set<Activation> getAllActivations() {
                 var result = new LinkedHashSet<Activation>();
                 var cursor = activations.getAll();
                 while (cursor.move()) {
@@ -79,6 +78,15 @@ public class TransformationRule {
                 return result;
         }
 
+        public List<Activation> getAllActivationsAsList() {
+                var result = new ArrayList<Activation>();
+                var cursor = activations.getAll();
+                while (cursor.move()) {
+                        result.add(new Activation(this, cursor.getKey()));
+                }
+                return result;
+        }
+
         public Activation getRandomActivation() {
                 return new Activation(this, activations.getKey(random.nextInt(activations.size())));
         }
diff --git a/subprojects/store-dse/src/main/java/tools/refinery/store/dse/objectives/AlwaysSatisfiedRandomHardObjective.java b/subprojects/store-dse/src/main/java/tools/refinery/store/dse/objectives/AlwaysSatisfiedRandomHardObjective.java
new file mode 100644
index 00000000..327d5e2f
--- /dev/null
+++ b/subprojects/store-dse/src/main/java/tools/refinery/store/dse/objectives/AlwaysSatisfiedRandomHardObjective.java
@@ -0,0 +1,56 @@
+/*******************************************************************************
+ * Copyright (c) 2010-2016, Andras Szabolcs Nagy, Zoltan Ujhelyi and Daniel Varro
+ * Copyright (c) 2023 The Refinery Authors <https://refinery.tools/>
+ * This program and the accompanying materials are made available under the
+ * terms of the Eclipse Public License v. 2.0 which is available at
+ * http://www.eclipse.org/legal/epl-v20.html.
+ *
+ * SPDX-License-Identifier: EPL-2.0
+ *******************************************************************************/
+package tools.refinery.store.dse.objectives;
+
+import tools.refinery.store.dse.DesignSpaceExplorationAdapter;
+
+import java.util.Random;
+
+/**
+ * This hard objective is fulfilled in any circumstances. Use it if all states should be regarded as a valid solution.
+ *
+ * @author Andras Szabolcs Nagy
+ *
+ */
+public class AlwaysSatisfiedRandomHardObjective extends BaseObjective {
+
+        private static final String DEFAULT_NAME = "AlwaysSatisfiedDummyHardObjective";
+        private static final Random random = new Random(0);
+
+        public AlwaysSatisfiedRandomHardObjective() {
+                super(DEFAULT_NAME);
+        }
+
+        public AlwaysSatisfiedRandomHardObjective(String name) {
+                super(name);
+        }
+
+        @Override
+        public Double getFitness(DesignSpaceExplorationAdapter context) {
+//              return 0d;
+                return random.nextDouble();
+        }
+
+        @Override
+        public boolean isHardObjective() {
+                return true;
+        }
+
+        @Override
+        public boolean satisfiesHardObjective(Double fitness) {
+                return true;
+        }
+
+        @Override
+        public Objective createNew() {
+                return this;
+        }
+
+}
diff --git a/subprojects/store-dse/src/main/java/tools/refinery/store/dse/objectives/BaseObjective.java b/subprojects/store-dse/src/main/java/tools/refinery/store/dse/objectives/BaseObjective.java
index 7df33efe..b76598fb 100644
--- a/subprojects/store-dse/src/main/java/tools/refinery/store/dse/objectives/BaseObjective.java
+++ b/subprojects/store-dse/src/main/java/tools/refinery/store/dse/objectives/BaseObjective.java
@@ -30,7 +30,7 @@ public abstract class BaseObjective implements Objective {
         protected boolean isThereFitnessConstraint = false;
         protected Comparator<Double> fitnessConstraintComparator;
 
-        public BaseObjective(String name) {
+        protected BaseObjective(String name) {
                 Objects.requireNonNull(name, "Name of the objective cannot be null.");
                 this.name = name;
         }
diff --git a/subprojects/store-dse/src/main/java/tools/refinery/store/dse/objectives/Comparators.java b/subprojects/store-dse/src/main/java/tools/refinery/store/dse/objectives/Comparators.java
index 476504b0..181397b3 100644
--- a/subprojects/store-dse/src/main/java/tools/refinery/store/dse/objectives/Comparators.java
+++ b/subprojects/store-dse/src/main/java/tools/refinery/store/dse/objectives/Comparators.java
@@ -15,7 +15,7 @@ public class Comparators {
 
         private Comparators() { /*Utility class constructor*/ }
 
-        public static final Comparator<Double> HIGHER_IS_BETTER = (o1, o2) -> o1.compareTo(o2);
+        public static final Comparator<Double> HIGHER_IS_BETTER = Double::compareTo;
 
         public static final Comparator<Double> LOWER_IS_BETTER = (o1, o2) -> o2.compareTo(o1);
 
diff --git a/subprojects/store-dse/src/main/java/tools/refinery/store/dse/objectives/Fitness.java b/subprojects/store-dse/src/main/java/tools/refinery/store/dse/objectives/Fitness.java
index 92709d3e..b1dc4442 100644
--- a/subprojects/store-dse/src/main/java/tools/refinery/store/dse/objectives/Fitness.java
+++ b/subprojects/store-dse/src/main/java/tools/refinery/store/dse/objectives/Fitness.java
@@ -27,4 +27,20 @@ public class Fitness extends HashMap<String, Double> {
         public String toString() {
                 return super.toString() + " hardObjectives=" + satisfiesHardObjectives;
         }
+
+        @Override
+        public boolean equals(Object other) {
+                if (other == null) return false;
+                if (getClass() != other.getClass()) return false;
+                if (!super.equals(other)) return false;
+                return satisfiesHardObjectives == ((Fitness) other).satisfiesHardObjectives;
+        }
+
+        @Override
+        public int hashCode() {
+                int h = super.hashCode();
+                h = h * 31 + (satisfiesHardObjectives ? 1 : 0);
+                return h;
+        }
+
 }
diff --git a/subprojects/store-dse/src/main/java/tools/refinery/store/dse/objectives/ObjectiveComparatorHelper.java b/subprojects/store-dse/src/main/java/tools/refinery/store/dse/objectives/ObjectiveComparatorHelper.java
index 1d676562..eb03eeaf 100644
--- a/subprojects/store-dse/src/main/java/tools/refinery/store/dse/objectives/ObjectiveComparatorHelper.java
+++ b/subprojects/store-dse/src/main/java/tools/refinery/store/dse/objectives/ObjectiveComparatorHelper.java
@@ -49,6 +49,7 @@ public class ObjectiveComparatorHelper {
                         }
                 }
                 if (o2HasBetterFitness) {
+                        return -1;
                 } else if (o1HasBetterFitness) {
                         return 1;
                 }
diff --git a/subprojects/store-dse/src/main/java/tools/refinery/store/dse/strategy/BestFirstStrategy.java b/subprojects/store-dse/src/main/java/tools/refinery/store/dse/strategy/BestFirstStrategy.java
index 8648864c..92d878ce 100644
--- a/subprojects/store-dse/src/main/java/tools/refinery/store/dse/strategy/BestFirstStrategy.java
+++ b/subprojects/store-dse/src/main/java/tools/refinery/store/dse/strategy/BestFirstStrategy.java
@@ -16,49 +16,52 @@ import tools.refinery.store.dse.internal.Activation;
 import tools.refinery.store.dse.objectives.Fitness;
 import tools.refinery.store.dse.objectives.ObjectiveComparatorHelper;
 
-import java.util.Collection;
-import java.util.Iterator;
-import java.util.List;
-import java.util.PriorityQueue;
+import java.util.*;
 
 public class BestFirstStrategy implements Strategy {
 
         private DesignSpaceExplorationAdapter dseAdapter;
 
-        private int maxDepth;
+        private int maxDepth = Integer.MAX_VALUE;
+        private int maxSolutions = Integer.MAX_VALUE;
         private boolean backTrackIfSolution = true;
         private boolean onlyBetterFirst = false;
 
         private PriorityQueue<TrajectoryWithFitness> trajectoriesToExplore;
 
-        private static class TrajectoryWithFitness {
-
-                public List<Version> trajectory;
-                public Fitness fitness;
+        private record TrajectoryWithFitness(List<Version> trajectory, Fitness fitness) {
+                @Override
+                public String toString() {
+                                return trajectory.toString() + fitness.toString();
+                        }
 
-                public TrajectoryWithFitness(List<Version> trajectory, Fitness fitness) {
-                        super();
-                        this.trajectory = trajectory;
-                        this.fitness = fitness;
+                @Override
+                public int hashCode() {
+                        return trajectory.get(trajectory.size() - 1).hashCode();
                 }
 
                 @Override
-                public String toString() {
-                        return trajectory.toString() + fitness.toString();
+                public boolean equals(Object obj) {
+                        if (obj instanceof TrajectoryWithFitness other) {
+                                return Objects.equals(trajectory.get(trajectory.size() - 1), other.trajectory.get(other.trajectory.size() - 1));
+//                              return trajectory.equals(((TrajectoryWithFitness) obj).trajectory);
+                        }
+                        return false;
                 }
-
         }
 
-        public BestFirstStrategy() {
-                this(-1);
+        public BestFirstStrategy withDepthLimit(int maxDepth) {
+                if (maxDepth >= 0) {
+                        this.maxDepth = maxDepth;
+                }
+                return this;
         }
 
-        public BestFirstStrategy(int maxDepth) {
-                if (maxDepth < 0) {
-                        this.maxDepth = Integer.MAX_VALUE;
-                } else {
-                        this.maxDepth = maxDepth;
+        public BestFirstStrategy withSolutionLimit(int maxSolutions) {
+                if (maxSolutions >= 0) {
+                        this.maxSolutions = maxSolutions;
                 }
+                return this;
         }
 
         public BestFirstStrategy continueIfHardObjectivesFulfilled() {
@@ -72,28 +75,31 @@ public class BestFirstStrategy implements Strategy {
         }
 
         @Override
-        public void initStrategy(DesignSpaceExplorationAdapter designSpaceExplorationAdapter) {
+        public void initialize(DesignSpaceExplorationAdapter designSpaceExplorationAdapter) {
                 this.dseAdapter = designSpaceExplorationAdapter;
                 final ObjectiveComparatorHelper objectiveComparatorHelper = dseAdapter.getObjectiveComparatorHelper();
 
-                trajectoriesToExplore = new PriorityQueue<TrajectoryWithFitness>(11,
+                trajectoriesToExplore = new PriorityQueue<>(11,
                                 (o1, o2) -> objectiveComparatorHelper.compare(o2.fitness, o1.fitness));
         }
 
         @Override
         public void explore() {
+                if (maxSolutions == 0) {
+                        return;
+                }
                 final ObjectiveComparatorHelper objectiveComparatorHelper = dseAdapter.getObjectiveComparatorHelper();
 
                 boolean globalConstraintsAreSatisfied = dseAdapter.checkGlobalConstraints();
                 if (!globalConstraintsAreSatisfied) {
-                        // "Global constraint is not satisfied in the first state. Terminate.");
+                        // Global constraint is not satisfied in the first state. Terminate.
                         return;
                 }
 
-                final Fitness firstFitness = dseAdapter.calculateFitness();
+                final Fitness firstFitness = dseAdapter.getFitness();
                 if (firstFitness.isSatisfiesHardObjectives()) {
                         dseAdapter.newSolution();
-                        // "First state is a solution. Terminate.");
+                        // First state is a solution. Terminate.
                         if (backTrackIfSolution) {
                                 return;
                         }
@@ -103,21 +109,20 @@ public class BestFirstStrategy implements Strategy {
                         return;
                 }
 
-                final List<Version> firstTrajectory = dseAdapter.getTrajectory();
-                TrajectoryWithFitness currentTrajectoryWithFitness = new TrajectoryWithFitness(firstTrajectory, firstFitness);
-                trajectoriesToExplore.add(currentTrajectoryWithFitness);
+
+                var firstTrajectoryWithFitness = new TrajectoryWithFitness(dseAdapter.getTrajectory(), firstFitness);
+                trajectoriesToExplore.add(firstTrajectoryWithFitness);
+                TrajectoryWithFitness currentTrajectoryWithFitness = null;
 
                 mainLoop: while (true) {
 
                         if (currentTrajectoryWithFitness == null) {
                                 if (trajectoriesToExplore.isEmpty()) {
-                                        // "State space is fully traversed.");
+                                        // State space is fully traversed.
                                         return;
                                 } else {
                                         currentTrajectoryWithFitness = trajectoriesToExplore.element();
-//                                      if (logger.isDebugEnabled()) {
-//                                               "New trajectory is chosen: " + currentTrajectoryWithFitness);
-//                                      }
+                                        // New trajectory is chosen: " + currentTrajectoryWithFitness
                                         dseAdapter.restoreTrajectory(currentTrajectoryWithFitness.trajectory);
                                 }
                         }
@@ -130,32 +135,34 @@ public class BestFirstStrategy implements Strategy {
                         while (iterator.hasNext()) {
                                 final Activation nextActivation = iterator.next();
                                 if (!iterator.hasNext()) {
-                                        // "Last untraversed activation of the state.");
+                                        // Last untraversed activation of the state.
                                         trajectoriesToExplore.remove(currentTrajectoryWithFitness);
                                 }
 
-//                              if (logger.isDebugEnabled()) {
-//                                       "Executing new activation: " + nextActivation);
-//                              }
+                                // Executing new activation
                                 dseAdapter.fireActivation(nextActivation);
                                 if (dseAdapter.isCurrentStateAlreadyTraversed()) {
-                                        // "The new state is already visited.");
+                                        // The new state is already visited.
                                         dseAdapter.backtrack();
                                 } else if (!dseAdapter.checkGlobalConstraints()) {
-                                        // "Global constraint is not satisfied.");
+                                        // Global constraint is not satisfied.
                                         dseAdapter.backtrack();
                                 } else {
-                                        final Fitness nextFitness = dseAdapter.calculateFitness();
+                                        final Fitness nextFitness = dseAdapter.getFitness();
                                         if (nextFitness.isSatisfiesHardObjectives()) {
                                                 dseAdapter.newSolution();
-                                                // "Found a solution.");
+                                                var solutions = dseAdapter.getSolutions().size();
+                                                if (solutions >= maxSolutions) {
+                                                        return;
+                                                }
+                                                // Found a solution.
                                                 if (backTrackIfSolution) {
                                                         dseAdapter.backtrack();
                                                         continue;
                                                 }
                                         }
                                         if (dseAdapter.getDepth() >= maxDepth) {
-                                                // "Reached max depth.");
+                                                // Reached max depth.
                                                 dseAdapter.backtrack();
                                                 continue;
                                         }
@@ -167,33 +174,30 @@ public class BestFirstStrategy implements Strategy {
                                         int compare = objectiveComparatorHelper.compare(currentTrajectoryWithFitness.fitness,
                                                         nextTrajectoryWithFitness.fitness);
                                         if (compare < 0) {
-                                                // "Better fitness, moving on: " + nextFitness);
+                                                // Better fitness, moving on
                                                 currentTrajectoryWithFitness = nextTrajectoryWithFitness;
                                                 continue mainLoop;
                                         } else if (compare == 0) {
                                                 if (onlyBetterFirst) {
-                                                        // "Equally good fitness, backtrack: " + nextFitness);
+                                                        // Equally good fitness, backtrack
                                                         dseAdapter.backtrack();
-                                                        continue;
                                                 } else {
-                                                        // "Equally good fitness, moving on: " + nextFitness);
+                                                        // Equally good fitness, moving on
                                                         currentTrajectoryWithFitness = nextTrajectoryWithFitness;
                                                         continue mainLoop;
                                                 }
                                         } else {
-                                                // "Worse fitness.");
+                                                //"Worse fitness
                                                 currentTrajectoryWithFitness = null;
                                                 continue mainLoop;
                                         }
                                 }
                         }
 
-                        // "State is fully traversed.");
-                        trajectoriesToExplore.remove(currentTrajectoryWithFitness);
+                        // State is fully traversed.
                         currentTrajectoryWithFitness = null;
 
                 }
-                // "Interrupted.");
-
+                // Interrupted.
         }
 }
diff --git a/subprojects/store-dse/src/main/java/tools/refinery/store/dse/strategy/DepthFirstStrategy.java b/subprojects/store-dse/src/main/java/tools/refinery/store/dse/strategy/DepthFirstStrategy.java
index 1405789b..0a0caa7e 100644
--- a/subprojects/store-dse/src/main/java/tools/refinery/store/dse/strategy/DepthFirstStrategy.java
+++ b/subprojects/store-dse/src/main/java/tools/refinery/store/dse/strategy/DepthFirstStrategy.java
@@ -1,117 +1,92 @@
-/*******************************************************************************
- * Copyright (c) 2010-2016, Andras Szabolcs Nagy, Zoltan Ujhelyi and Daniel Varro
- * Copyright (c) 2023 The Refinery Authors <https://refinery.tools/>
- * This program and the accompanying materials are made available under the
- * terms of the Eclipse Public License v. 2.0 which is available at
- * http://www.eclipse.org/legal/epl-v20.html.
+/*
+ * SPDX-FileCopyrightText: 2021-2023 The Refinery Authors <https://refinery.tools/>
  *
  * SPDX-License-Identifier: EPL-2.0
- *******************************************************************************/
+ */
 package tools.refinery.store.dse.strategy;
 
 import tools.refinery.store.dse.DesignSpaceExplorationAdapter;
 import tools.refinery.store.dse.Strategy;
-import tools.refinery.store.dse.internal.Activation;
 import tools.refinery.store.dse.objectives.Fitness;
 
-import java.util.Collection;
-
 public class DepthFirstStrategy implements Strategy {
 
         private DesignSpaceExplorationAdapter dseAdapter;
 
-        private int maxDepth;
-        private boolean backTrackIfSolution = true;
+        private int maxDepth = Integer.MAX_VALUE;
+        private int maxSolutions = Integer.MAX_VALUE;
+        private boolean backtrackFromSolution = true;
 
-        public DepthFirstStrategy() {
-                this(-1);
+        public DepthFirstStrategy withDepthLimit(int maxDepth) {
+                if (maxDepth >= 0) {
+                        this.maxDepth = maxDepth;
+                }
+                return this;
         }
 
-        public DepthFirstStrategy(int maxDepth) {
-                if (maxDepth < 0) {
-                        this.maxDepth = Integer.MAX_VALUE;
-                } else {
-                        this.maxDepth = maxDepth;
+        public DepthFirstStrategy withSolutionLimit(int maxSolutions) {
+                if (maxSolutions >= 0) {
+                        this.maxSolutions = maxSolutions;
                 }
+                return this;
         }
 
         public DepthFirstStrategy continueIfHardObjectivesFulfilled() {
-                backTrackIfSolution = false;
+                backtrackFromSolution = false;
                 return this;
         }
 
         @Override
-        public void initStrategy(DesignSpaceExplorationAdapter designSpaceExplorationAdapter) {
+        public void initialize(DesignSpaceExplorationAdapter designSpaceExplorationAdapter) {
                 this.dseAdapter = designSpaceExplorationAdapter;
         }
 
         @Override
         public void explore() {
-                mainloop: while (true) {
-                        var globalConstraintsAreSatisfied = dseAdapter.checkGlobalConstraints();
-                        if (!globalConstraintsAreSatisfied) {
-                                var isSuccessfulUndo = dseAdapter.backtrack();
-                                if (!isSuccessfulUndo) {
-//                                      "Global constraint is not satisfied and cannot backtrack."
-                                        break;
-                                }
-                                else {
-//                                      "Global constraint is not satisfied, backtrack."
-                                        continue;
-                                }
+                if (maxSolutions == 0) {
+                        return;
+                }
+                while (dseAdapter.getSolutions().size() < maxSolutions) {
+                        if (!checkAndHandleGlobalConstraints()) {
+                                return;
                         }
 
-                        Fitness fitness = dseAdapter.calculateFitness();
+                        Fitness fitness = dseAdapter.getFitness();
                         if (fitness.isSatisfiesHardObjectives()) {
                                 dseAdapter.newSolution();
-                                if (backTrackIfSolution) {
-                                        var isSuccessfulUndo = dseAdapter.backtrack();
-                                        if (!isSuccessfulUndo) {
-//                                      "Found a solution but cannot backtrack."
-                                                break;
-                                        } else {
-//                                      "Found a solution, backtrack."
-                                                continue;
-                                        }
+                                if (backtrackFromSolution && !dseAdapter.backtrack()) {
+                                        return;
                                 }
                         }
 
-                        var depth = dseAdapter.getDepth();
-                        if (dseAdapter.getDepth() >= maxDepth) {
-                                var isSuccessfulUndo = dseAdapter.backtrack();
-                                if (!isSuccessfulUndo) {
-//                                      "Reached max depth but cannot backtrack."
-                                        break;
-                                }
+                        if (!checkAndHandleDepth()) {
+                                return;
                         }
 
-                        Collection<Activation> activations;
-                        do {
-                                activations = dseAdapter.getUntraversedActivations();
-                                if (activations.isEmpty()) {
-                                        if (!dseAdapter.backtrack()) {
-        //                                      "No more transitions from current state and cannot backtrack."
-                                                break mainloop;
-                                        }
-                                        else {
-        //                                      "No more transitions from current state, backtrack."
-                                                continue;
-                                        }
-                                }
-                        } while (activations.isEmpty());
-
-                        dseAdapter.fireRandomActivation();
-//                      if (dseAdapter.isCurrentInTrajectory()) {
-//                              if (!dseAdapter.backtrack()) {
-////                                    TODO: throw exception
-////                                    "The new state is present in the trajectory but cannot backtrack. Should never happen!"
-//                                      break;
-//                              }
-//                              else {
-////                                    "The new state is already visited in the trajectory, backtrack."
-//                                      continue;
-//                              }
-//                      }
+                        if (!backtrackToLastUntraversed()) {
+                                return;
+                        }
+
+                        if (!dseAdapter.fireRandomActivation()) {
+                                return;
+                        }
+                }
+        }
+
+        private boolean checkAndHandleGlobalConstraints() {
+                return dseAdapter.checkGlobalConstraints() || dseAdapter.backtrack();
+        }
+
+        private boolean checkAndHandleDepth() {
+                return dseAdapter.getDepth() < maxDepth || dseAdapter.backtrack();
+        }
+
+        private boolean backtrackToLastUntraversed() {
+                while (dseAdapter.getUntraversedActivations().isEmpty()) {
+                        if (!dseAdapter.backtrack()) {
+                                return false;
+                        }
                 }
+                return true;
         }
 }