From 76a34bc9570f3be59dbb7d877e83448fcc5461b3 Mon Sep 17 00:00:00 2001
From: lukpo <lukas.posch@student.tugraz.at>
Date: Thu, 15 Jul 2021 11:51:12 +0200
Subject: [PATCH] expanded compute method in
 SparseNondeterministicStepBoundedHorizonHelper.* with resultMaybestates and
 choiceValues

---
 ...ndeterministicStepBoundedHorizonHelper.cpp | 35 ++++++++++++++-----
 ...NondeterministicStepBoundedHorizonHelper.h |  3 +-
 2 files changed, 29 insertions(+), 9 deletions(-)

diff --git a/src/storm/modelchecker/helper/finitehorizon/SparseNondeterministicStepBoundedHorizonHelper.cpp b/src/storm/modelchecker/helper/finitehorizon/SparseNondeterministicStepBoundedHorizonHelper.cpp
index 68d7631d0..a37d11782 100644
--- a/src/storm/modelchecker/helper/finitehorizon/SparseNondeterministicStepBoundedHorizonHelper.cpp
+++ b/src/storm/modelchecker/helper/finitehorizon/SparseNondeterministicStepBoundedHorizonHelper.cpp
@@ -25,7 +25,7 @@ namespace storm {
             }
 
             template<typename ValueType>
-            std::vector<ValueType> SparseNondeterministicStepBoundedHorizonHelper<ValueType>::compute(Environment const& env, storm::solver::SolveGoal<ValueType>&& goal, storm::storage::SparseMatrix<ValueType> const& transitionMatrix, storm::storage::SparseMatrix<ValueType> const& backwardTransitions, storm::storage::BitVector const& phiStates, storm::storage::BitVector const& psiStates,   uint64_t lowerBound, uint64_t upperBound,  ModelCheckerHint const& hint)
+            std::vector<ValueType> SparseNondeterministicStepBoundedHorizonHelper<ValueType>::compute(Environment const& env, storm::solver::SolveGoal<ValueType>&& goal, storm::storage::SparseMatrix<ValueType> const& transitionMatrix, storm::storage::SparseMatrix<ValueType> const& backwardTransitions, storm::storage::BitVector const& phiStates, storm::storage::BitVector const& psiStates, uint64_t lowerBound, uint64_t upperBound, ModelCheckerHint const& hint, storm::storage::BitVector& resultMaybeStates, std::vector<ValueType>& choiceValues)
             {
                 std::vector<ValueType> result(transitionMatrix.getRowGroupCount(), storm::utility::zero<ValueType>());
                 storm::storage::BitVector makeZeroColumns;
@@ -61,13 +61,27 @@ namespace storm {
 
                     auto multiplier = storm::solver::MultiplierFactory<ValueType>().create(env, submatrix);
                     if (lowerBound == 0) {
-                        multiplier->repeatedMultiplyAndReduce(env, goal.direction(), subresult, &b, upperBound);
+                        if(goal.isShieldingTask())
+                        {
+                            multiplier->repeatedMultiplyAndReduceWithChoices(env, goal.direction(), subresult, &b, upperBound, nullptr, choiceValues, transitionMatrix.getRowGroupIndices());
+                        } else {
+                            multiplier->repeatedMultiplyAndReduce(env, goal.direction(), subresult, &b, upperBound);
+                        }
                     } else {
-                        multiplier->repeatedMultiplyAndReduce(env, goal.direction(), subresult, &b, upperBound - lowerBound + 1);
-                        storm::storage::SparseMatrix<ValueType> submatrix = transitionMatrix.getSubmatrix(true, maybeStates, maybeStates, false);
-                        auto multiplier = storm::solver::MultiplierFactory<ValueType>().create(env, submatrix);
-                        b = std::vector<ValueType>(b.size(), storm::utility::zero<ValueType>());
-                        multiplier->repeatedMultiplyAndReduce(env, goal.direction(), subresult, &b, lowerBound - 1);
+                        if(goal.isShieldingTask())
+                        {
+                            multiplier->repeatedMultiplyAndReduceWithChoices(env, goal.direction(), subresult, &b, upperBound - lowerBound + 1, nullptr, choiceValues, transitionMatrix.getRowGroupIndices());
+                            storm::storage::SparseMatrix<ValueType> submatrix = transitionMatrix.getSubmatrix(true, maybeStates, maybeStates, false);
+                            auto multiplier = storm::solver::MultiplierFactory<ValueType>().create(env, submatrix);
+                            b = std::vector<ValueType>(b.size(), storm::utility::zero<ValueType>());
+                            multiplier->repeatedMultiplyAndReduceWithChoices(env, goal.direction(), subresult, &b, lowerBound - 1, nullptr, choiceValues, transitionMatrix.getRowGroupIndices());
+                        } else {
+                            multiplier->repeatedMultiplyAndReduce(env, goal.direction(), subresult, &b, upperBound - lowerBound + 1);
+                            storm::storage::SparseMatrix<ValueType> submatrix = transitionMatrix.getSubmatrix(true, maybeStates, maybeStates, false);
+                            auto multiplier = storm::solver::MultiplierFactory<ValueType>().create(env, submatrix);
+                            b = std::vector<ValueType>(b.size(), storm::utility::zero<ValueType>());
+                            multiplier->repeatedMultiplyAndReduce(env, goal.direction(), subresult, &b, lowerBound - 1);
+                        }
                     }
                     // Set the values of the resulting vector accordingly.
                     storm::utility::vector::setVectorValues(result, maybeStates, subresult);
@@ -75,10 +89,15 @@ namespace storm {
                 if (lowerBound == 0) {
                     storm::utility::vector::setVectorValues(result, psiStates, storm::utility::one<ValueType>());
                 }
+
+                //TODO: check if this works with nullptr as default for resultMaybeStates
+                if(!resultMaybeStates.empty())
+                {
+                    resultMaybeStates = maybeStates;
+                }
                 return result;
             }
 
-
             template class SparseNondeterministicStepBoundedHorizonHelper<double>;
             template class SparseNondeterministicStepBoundedHorizonHelper<storm::RationalNumber>;
         }
diff --git a/src/storm/modelchecker/helper/finitehorizon/SparseNondeterministicStepBoundedHorizonHelper.h b/src/storm/modelchecker/helper/finitehorizon/SparseNondeterministicStepBoundedHorizonHelper.h
index 98fc49c2c..0a0a0fc51 100644
--- a/src/storm/modelchecker/helper/finitehorizon/SparseNondeterministicStepBoundedHorizonHelper.h
+++ b/src/storm/modelchecker/helper/finitehorizon/SparseNondeterministicStepBoundedHorizonHelper.h
@@ -4,6 +4,7 @@
 #include "storm/modelchecker/hints/ModelCheckerHint.h"
 #include "storm/modelchecker/prctl/helper/SolutionType.h"
 #include "storm/storage/SparseMatrix.h"
+#include "storm/storage/BitVector.h"
 #include "storm/utility/solver.h"
 #include "storm/solver/SolveGoal.h"
 
@@ -15,7 +16,7 @@ namespace storm {
             class SparseNondeterministicStepBoundedHorizonHelper {
             public:
                 SparseNondeterministicStepBoundedHorizonHelper();
-                std::vector<ValueType> compute(Environment const& env, storm::solver::SolveGoal<ValueType>&& goal, storm::storage::SparseMatrix<ValueType> const& transitionMatrix, storm::storage::SparseMatrix<ValueType> const& backwardTransitions,  storm::storage::BitVector const& phiStates, storm::storage::BitVector const& psiStates, uint64_t lowerBound, uint64_t upperBound, ModelCheckerHint const& hint = ModelCheckerHint());
+                std::vector<ValueType> compute(Environment const& env, storm::solver::SolveGoal<ValueType>&& goal, storm::storage::SparseMatrix<ValueType> const& transitionMatrix, storm::storage::SparseMatrix<ValueType> const& backwardTransitions, storm::storage::BitVector const& phiStates, storm::storage::BitVector const& psiStates, uint64_t lowerBound, uint64_t upperBound, ModelCheckerHint const& hint = ModelCheckerHint(), storm::storage::BitVector& resultMaybeStates = nullptr, std::vector<ValueType>& choiceValues = nullptr);
             private:
             };