From 37d36c52b3d1e2c573415db89b065d21f4d6805c Mon Sep 17 00:00:00 2001
From: Lukas Posch <lukas.posch@student.tugraz.at>
Date: Mon, 15 Feb 2021 15:46:07 +0100
Subject: [PATCH] fill up the result vector for ~relevantStates

---
 .../rpatl/helper/SparseSmgRpatlHelper.cpp       |  7 +++----
 .../rpatl/helper/internal/GameViHelper.cpp      | 17 +++++++++++++++++
 .../rpatl/helper/internal/GameViHelper.h        |  5 +++++
 3 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/src/storm/modelchecker/rpatl/helper/SparseSmgRpatlHelper.cpp b/src/storm/modelchecker/rpatl/helper/SparseSmgRpatlHelper.cpp
index 57ea8ba29..b30cea558 100644
--- a/src/storm/modelchecker/rpatl/helper/SparseSmgRpatlHelper.cpp
+++ b/src/storm/modelchecker/rpatl/helper/SparseSmgRpatlHelper.cpp
@@ -94,11 +94,10 @@ namespace storm {
                 }
 
                 viHelper.performValueIteration(env, x, b, goal.direction());
-
-                // TODO: here is the debug output for all states, should I fill up the vector again with 0 for the states i left out
-                //  e.g. phi states are [ 2 3 ], the probability output is for the first state (state 2) but the initial sate is 0
-
+                //STORM_LOG_DEBUG(storm::utility::vector::toString(x));
+                viHelper.fillResultVector(x, relevantStates, psiStates);
                 STORM_LOG_DEBUG(storm::utility::vector::toString(x));
+
                 if (produceScheduler) {
                     scheduler = std::make_unique<storm::storage::Scheduler<ValueType>>(expandScheduler(viHelper.extractScheduler(), psiStates));
                     STORM_LOG_DEBUG("IsPartial?" << scheduler->isPartialScheduler());
diff --git a/src/storm/modelchecker/rpatl/helper/internal/GameViHelper.cpp b/src/storm/modelchecker/rpatl/helper/internal/GameViHelper.cpp
index fb8882a5d..bff9bc052 100644
--- a/src/storm/modelchecker/rpatl/helper/internal/GameViHelper.cpp
+++ b/src/storm/modelchecker/rpatl/helper/internal/GameViHelper.cpp
@@ -133,6 +133,23 @@ namespace storm {
                     return true;
                 }
 
+                template <typename ValueType>
+                void GameViHelper<ValueType>::fillResultVector(std::vector<ValueType>& result, storm::storage::BitVector relevantStates, storm::storage::BitVector psiStates)
+                {
+                    std::vector<ValueType> filled_vector = std::vector<ValueType>(relevantStates.size(), storm::utility::zero<ValueType>());
+                    uint bit_index = 0;
+                    for(uint i = 0; i < filled_vector.size(); i++) {
+                        if (relevantStates.get(i)) {
+                            filled_vector.at(i) = result.at(bit_index);
+                            bit_index++;
+                        }
+                        else if (psiStates.get(i)) {
+                            filled_vector.at(i) = 1;
+                        }
+                    }
+                    result = filled_vector;
+                }
+
                 template <typename ValueType>
                 void GameViHelper<ValueType>::setProduceScheduler(bool value) {
                     _produceScheduler = value;
diff --git a/src/storm/modelchecker/rpatl/helper/internal/GameViHelper.h b/src/storm/modelchecker/rpatl/helper/internal/GameViHelper.h
index 92f963851..f3e353665 100644
--- a/src/storm/modelchecker/rpatl/helper/internal/GameViHelper.h
+++ b/src/storm/modelchecker/rpatl/helper/internal/GameViHelper.h
@@ -25,6 +25,11 @@ namespace storm {
 
                     void performValueIteration(Environment const& env, std::vector<ValueType>& x, std::vector<ValueType> b, storm::solver::OptimizationDirection const dir);
 
+                    /*!
+                     * Fills the result vector to the original size with ones for being psiStates, zeros for being not phiStates
+                     */
+                    void fillResultVector(std::vector<ValueType>& result, storm::storage::BitVector relevantStates, storm::storage::BitVector psiStates);
+
                     /*h
                      * Sets whether an optimal scheduler shall be constructed during the computation
                      */