1 files changed, 99 insertions, 0 deletions

diff --git a/Solvers/VIATRA-Solver/hu.bme.mit.inf.dslreasoner.viatrasolver.reasoner/src/hu/bme/mit/inf/dslreasoner/viatrasolver/reasoner/dse/ViatraReasonerSolutionSaver.xtend b/Solvers/VIATRA-Solver/hu.bme.mit.inf.dslreasoner.viatrasolver.reasoner/src/hu/bme/mit/inf/dslreasoner/viatrasolver/reasoner/dse/ViatraReasonerSolutionSaver.xtend
new file mode 100644
index 00000000..5877778e
--- /dev/null
+++ b/Solvers/VIATRA-Solver/hu.bme.mit.inf.dslreasoner.viatrasolver.reasoner/src/hu/bme/mit/inf/dslreasoner/viatrasolver/reasoner/dse/ViatraReasonerSolutionSaver.xtend
@@ -0,0 +1,99 @@
+package hu.bme.mit.inf.dslreasoner.viatrasolver.reasoner.dse
+
+import java.util.HashMap
+import java.util.Map
+import org.eclipse.viatra.dse.api.DSEException
+import org.eclipse.viatra.dse.api.Solution
+import org.eclipse.viatra.dse.api.SolutionTrajectory
+import org.eclipse.viatra.dse.base.ThreadContext
+import org.eclipse.viatra.dse.objectives.Fitness
+import org.eclipse.viatra.dse.objectives.IObjective
+import org.eclipse.viatra.dse.objectives.ObjectiveComparatorHelper
+import org.eclipse.viatra.dse.solutionstore.SolutionStore.ISolutionSaver
+import org.eclipse.xtend.lib.annotations.Accessors
+
+/**
+ * Based on {@link org.eclipse.viatra.dse.solutionstore.SolutionStore.BestSolutionSaver}.
+ */
+class ViatraReasonerSolutionSaver implements ISolutionSaver {
+        @Accessors val solutionCopier = new SolutionCopier
+        val boolean hasExtremalObjectives
+        val ObjectiveComparatorHelper comparatorHelper
+        val Map<SolutionTrajectory, Fitness> trajectories = new HashMap
+
+        @Accessors(PUBLIC_SETTER) var Map<Object, Solution> solutionsCollection
+
+        new(IObjective[][] leveledExtremalObjectives) {
+                comparatorHelper = new ObjectiveComparatorHelper(leveledExtremalObjectives)
+                hasExtremalObjectives = leveledExtremalObjectives.exists[!empty]
+        }
+
+        override saveSolution(ThreadContext context, Object id, SolutionTrajectory solutionTrajectory) {
+                if (hasExtremalObjectives) {
+                        saveBestSolutionOnly(context, id, solutionTrajectory)
+                } else {
+                        basicSaveSolution(context, id, solutionTrajectory)
+                }
+        }
+
+        private def saveBestSolutionOnly(ThreadContext context, Object id, SolutionTrajectory solutionTrajectory) {
+                val fitness = context.lastFitness
+                val dominatedTrajectories = newArrayList
+                for (entry : trajectories.entrySet) {
+                        val isLastFitnessBetter = comparatorHelper.compare(fitness, entry.value)
+                        if (isLastFitnessBetter < 0) {
+                                // Found a trajectory that dominates the current one, no need to save
+                                return false
+                        }
+                        if (isLastFitnessBetter > 0) {
+                                dominatedTrajectories += entry.key
+                        }
+                }
+                // We must save the new trajectory before removing dominated trajectories
+                // to avoid removing the current solution when it is reachable only via dominated trajectories.
+                val solutionSaved = basicSaveSolution(context, id, solutionTrajectory)
+                for (dominatedTrajectory : dominatedTrajectories) {
+                        trajectories -= dominatedTrajectory
+                        val dominatedSolution = dominatedTrajectory.solution
+                        if (!dominatedSolution.trajectories.remove(dominatedTrajectory)) {
+                                throw new DSEException(
+                                        "Dominated solution is not reachable from dominated trajectory. This should never happen!")
+                        }
+                        if (dominatedSolution.trajectories.empty) {
+                                val dominatedSolutionId = dominatedSolution.stateCode
+                                solutionCopier.markAsObsolete(dominatedSolutionId)
+                                solutionsCollection -= dominatedSolutionId
+                        }
+                }
+                solutionSaved
+        }
+
+        private def basicSaveSolution(ThreadContext context, Object id, SolutionTrajectory solutionTrajectory) {
+                val fitness = context.lastFitness
+                var boolean solutionSaved = false
+                var dseSolution = solutionsCollection.get(id)
+                if (dseSolution === null) {
+                        solutionCopier.copySolution(context, id)
+                        dseSolution = new Solution(id, solutionTrajectory)
+                        solutionsCollection.put(id, dseSolution)
+                        solutionSaved = true
+                } else {
+                        solutionSaved = dseSolution.trajectories.add(solutionTrajectory)
+                }
+                if (solutionSaved) {
+                        solutionTrajectory.solution = dseSolution
+                        trajectories.put(solutionTrajectory, fitness)
+                }
+                solutionSaved
+        }
+        
+        def isFitnessDominated(Fitness fitness) {
+                for (existingFitness : trajectories.values) {
+                        val isNewFitnessBetter = comparatorHelper.compare(fitness, existingFitness)
+                        if (isNewFitnessBetter < 0) {
+                                return true
+                        }
+                }
+                false
+        }
+}